# Teaching a Computer to "See": A Practical Guide to Image Classification with CNNs & PyTorch

Some grocery apps can identify a vegetable just from a photo. Point the camera at ginger, galangal, or turmeric (three roots people mix up constantly at the market) and the app tells you which is which. Underneath, there's one core idea at work: **image classification**.

This post breaks that idea down from scratch, starting with pixels and ending with a CNN ready to train in PyTorch.

> 💡 **How to read this:** a few "Guess First" boxes are scattered through the post. Take a guess before checking the answer, it helps the concept stick.

* * *

## Table of Contents

1.  [What Image Classification Actually Is](#what-image-classification-actually-is)
    
2.  [Computer Vision vs Image Processing](#computer-vision-vs-image-processing)
    
3.  [CNN Anatomy: Backbone vs Head](#cnn-anatomy-backbone-vs-head)
    
4.  [Getting Comfortable with PyTorch](#getting-comfortable-with-pytorch)
    
5.  [Building a CNN from Scratch](#building-a-cnn-from-scratch)
    
6.  [Transfer Learning](#transfer-learning)
    
7.  [Preparing a Dataset](#preparing-a-dataset)
    
8.  [The Training Loop](#the-training-loop)
    
9.  [Cheat Sheet](#cheat-sheet)
    
10.  [Test Your Understanding](#test-your-understanding)
     

* * *

## What Image Classification Actually Is

The definition is simple: feed in an image, get back what it's a picture of. Input is an image, output is a class or label.

That simple idea powers a lot of everyday applications: face recognition, disease detection in medical scans, object recognition in self-driving systems, and yes, the grocery app that tells vegetables apart.

Before going further, one thing is worth internalizing: to a computer, an image is just numbers.

*   An image is made of **pixels**, the smallest unit of a picture.
    
*   On a computer, an image becomes a **matrix**: rows by columns, each cell holding a value for the light intensity at that point.
    
*   A color image is really three matrices stacked together: Red (R), Green (G), Blue (B), forming a single 3-dimensional block of data.
    

> 🎯 **Guess First:** A 100×100 color image is made up of how many matrices?
> 
> **Jawaban:** Three (R, G, B), each 100×100, stacked into a single 100×100×3 block.

## Computer Vision vs Image Processing

These two terms get mixed up often, but they serve different goals.

|  | Image Processing | Computer Vision |
| --- | --- | --- |
| Goal | Manipulate how an image looks | Interpret what an image contains |
| Input | Image | Image or video |
| Output | A modified image | An interpretation: description, coordinates, class |
| Examples | Sharpening, blurring, edge detection | Object detection, image classification |

One image processing technique underpins CNNs directly: **convolution**, sliding a kernel (a small matrix) across an image to produce a specific effect.

A natural question follows: do the numbers inside that kernel have to be set by hand? They don't. Kernel values start out randomly initialized, then get updated automatically through **training** until the model settles on the best values for its data. A hand-designed kernel is fixed and doesn't adapt; a trained one shapes itself around whatever data it sees.

That's the foundation for this post's main subject: **image classification with CNNs (Convolutional Neural Networks)**.

## CNN Anatomy: Backbone vs Head

A CNN works in two phases: look closely first, then decide. These are called the **Backbone** and the **Head**.

### The Backbone, feature extractor

The backbone learns visual features from an image: edges, corners, textures, patterns, eventually objects. It works progressively. Early layers pick up simple features like lines and basic shapes. Deeper layers capture increasingly complex, abstract features, from object parts to whole objects. As the network goes deeper, the spatial size of the image shrinks while the information packed into it gets denser.

Each convolutional layer contains kernels (filters) that sweep across the image and produce a feature map. Two parameters matter most:

*   `in_channels`, the number of input channels. For the first layer receiving a color image directly, this is 3 (R, G, B).
    
*   `out_channels`, the number of output channels, equal to the number of filters used. Use 32 filters, get 32 stacked feature maps out.
    

Multiple filters help because each one learns to look for something different: texture, color, edges. More filters mean the model captures a richer range of visual patterns.

> ⚠️ **A hard rule when stacking layers:** a layer's `in_channels` (or `in_features`) must match the previous layer's `out_channels` (or `out_features`). Skip this and the code throws an error the moment it runs. It's the single most common mistake when designing a CNN architecture for the first time.

A linear layer only accepts a 1-dimensional vector, while an image is a matrix. Force that combination and the image needs to be *flattened* first. Example: a 28×28 image flattens into a vector of 784 (28 times 28). For a large image, say 1000×1000 pixels, that flattened vector balloons and becomes impractical.

Convolutional layers are useful precisely because they accept matrix input directly. By the time the gradual extraction process is done, the representation is already compact and information-dense before it gets flattened.

There's no need to design a backbone from scratch either. Proven architectures already exist and can be reused directly: ResNet, VGG, EfficientNet.

### The Head, decision maker

Once the backbone finishes extracting features, the result (a vector, after flattening) moves into the **head** for classification. The head is typically one or more fully connected (FC) layers.

Every neuron in an FC layer connects to every neuron in the layer before it, hence "fully connected." The final output layer needs exactly as many neurons as there are classes. Classifying digits 0 through 9 means 10 output neurons; a 15-category vegetable dataset means 15.

The last layer usually applies **softmax**, converting raw outputs into a probability for each class. The class with the highest probability becomes the model's prediction.

> 🎯 **Guess First:** A model classifies animals into 3 classes: cat, dog, panda. Its softmax output is `[0.15, 0.15, 0.70]`. What does it predict?
> 
> **Jawaban:** Panda, since it has the highest probability of the three.

In short: the backbone turns a raw image into a feature representation through convolution, pooling, and non-linear activation. The head takes that representation and picks the most likely class based on softmax probabilities.

![Backbone to Head flow in a CNN](https://cdn.hashnode.com/uploads/covers/6a8ef2e3923670c989379174/2d2c47aa-18f8-42ec-9f42-17500662abd6.png align="center")

*A raw image gets distilled into shrinking feature maps, flattened into a vector, then classified through FC layers into a final prediction.*

## Getting Comfortable with PyTorch

PyTorch is a tool for building and training deep learning models. Its API feels a lot like NumPy, so prior NumPy experience speeds things up considerably. The difference is that PyTorch plugs directly into hardware accelerators: CUDA GPUs, Apple Silicon, or plain CPU.

A few basics worth knowing upfront:

*   Default data types: integers become `int64`, decimals become `float32`.
    
*   A tensor's default device is CPU. Moving it to GPU is a single call: `.to('cuda')`.
    
*   **Autograd** is PyTorch's mechanism for computing gradients automatically, the core of training. That computation happens automatically; there's no manual gradient math involved.
    

```python
import torch

x = torch.randn(2, 2)
x = x.to('cuda')

a = torch.ones(2, 3)
b = torch.zeros(2, 3)
c = a + b
```

## Building a CNN from Scratch

### Linear Layer

```python
import torch.nn as nn

layer = nn.Linear(in_features=5, out_features=10)
```

This layer takes 5 input values and produces 10 output values. Because it's fully connected, there are 5 times 10 connections behind the scenes.

### A Custom Model

Building a custom architecture means writing a class that extends `nn.Module`. This is mandatory: PyTorch is designed so that any model recognized this way can actually go through training.

```python
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.block1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=10, kernel_size=3),
            nn.ReLU(),
        )
        self.block2 = nn.Sequential(
            nn.Conv2d(in_channels=10, out_channels=10, kernel_size=3),
            nn.ReLU(),
        )
        self.flatten = nn.Flatten()
        self.classifier = nn.Linear(in_features=10 * 24 * 24, out_features=num_classes)

    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.flatten(x)
        x = self.classifier(x)
        return x
```

Three things worth double-checking:

1.  Block two's `in_channels` (10) has to match block one's `out_channels` (10).
    
2.  `in_features` on the classifier has to match the flattened output size. For a 28×28 input with no padding, each Conv2d with kernel size 3 shrinks each side by 2 pixels: 28 becomes 26, then 26 becomes 24. That leaves a 24×24 map, so `in_features` = 10 times 24 times 24. When the math gets confusing, printing the output shape after each block is the fastest way to check.
    
3.  The final layer's `out_features` has to match the number of classes in the dataset.
    

## Transfer Learning

Rather than designing an architecture from zero, a pretrained, already-proven backbone can be reused directly. `torchvision`'s model zoo has plenty: ResNet, VGG, EfficientNet.

The pattern: load the backbone, then override its final classifier layer to match the target number of classes. Layer naming differs across architectures.

```python
import torch.nn as nn
from torchvision import models

# EfficientNet-B1
model = models.efficientnet_b1(weights="IMAGENET1K_V1")
model.classifier[1] = nn.Linear(in_features=1280, out_features=15)

# ResNet-50
model = models.resnet50(weights="IMAGENET1K_V1")
model.fc = nn.Linear(in_features=2048, out_features=15)
```

These architectures already extract general visual features well. The remaining work is teaching the model to recognize the specific classes in the target dataset.

## Preparing a Dataset

Two common approaches exist for organizing data before training.

### Folder per Class

```plaintext
dataset/
├── train/
│   ├── class1/
│   │   ├── image1.jpg
│   │   └── image2.jpg
│   └── class2/
│       ├── image3.jpg
│       └── image4.jpg
├── val/
│   ├── class1/
│   └── class2/
```

Each image sits in the folder matching its class. This structure is recognized directly by built-in framework utilities: `ImageFolder` in PyTorch, `ImageDataGenerator` in TensorFlow.

```python
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor()
])

train_dataset = datasets.ImageFolder(root='dataset/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
```

The upside is simplicity and broad framework support. The downside shows up with random filenames or extra metadata needs, like multi-label data.

### CSV File

The alternative is a CSV listing image paths alongside their labels.

```plaintext
image_path,class_label
dataset/images/image1.jpg,0
dataset/images/image2.jpg,1
dataset/images/image3.jpg,0
```

```python
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms

class CustomDataset(Dataset):
    def __init__(self, csv_file, transform=None):
        self.data = pd.read_csv(csv_file)
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path = self.data.iloc[idx, 0]
        image = Image.open(img_path)
        label = int(self.data.iloc[idx, 1])
        if self.transform:
            image = self.transform(image)
        return image, label

transform = transforms.Compose([
    transforms.Resize((128, 128)),
    transforms.ToTensor()
])

train_dataset = CustomDataset(csv_file='dataset/train_labels.csv', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
```

This is more flexible for datasets carrying extra metadata or scattered across an unorganized directory structure. The tradeoff is custom code, since no framework ships a built-in loader for this shape.

> 📌 **Which one to pick?** Clean, simple, already-organized-by-category data: use folders. Complex data with extra metadata or messy file layouts: use CSV.

### Dataset vs DataLoader

A **Dataset** holds all the data, images and labels together. Pulling one sample from it returns one image plus one label. Feeding an entire dataset to a model at once is usually impossible given GPU memory limits.

A **DataLoader** splits the dataset into **batches**, small chunks processed incrementally. `DataLoader(batch_size=128)` produces batches shaped `[128, 3, 64, 64]`: 128 images, 3 RGB channels, 64×64 pixels, plus 128 matching labels.

## The Training Loop

Each epoch (one full pass through the training data) runs through the following steps.

```python
model.train()
train_loss, train_acc = 0, 0

for images, labels in train_loader:
    images, labels = images.to(device), labels.to(device)

    outputs = model(images)
    loss = criterion(outputs, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    train_loss += loss.item()
    train_acc += (outputs.argmax(1) == labels).float().mean().item()

train_loss /= len(train_loader)
train_acc /= len(train_loader)
```

The test/validation step follows the same shape, minus `loss.backward()` and `optimizer.step()`. That step is pure evaluation, not learning.

The loss function for multi-class problems: **Cross Entropy Loss**.

Finally, save the model whenever it improves, so it can be reused later without retraining from zero.

```python
if test_acc > best_acc:
    best_acc = test_acc
    torch.save(model.state_dict(), 'best_model.pth')
```

## Cheat Sheet

*   \[ \] A CNN is a Backbone (feature extraction) plus a Head (classification)
    
*   \[ \] A layer's `in_channels`/`in_features` must match the previous layer's `out_channels`/`out_features`
    
*   \[ \] `out_channels` equals the number of filters used
    
*   \[ \] The final output layer needs one neuron per class
    
*   \[ \] Linear layers need 1D input, so images must be flattened first
    
*   \[ \] Convolutional layers accept matrix input directly
    
*   \[ \] Kernel values come from training, not manual design
    
*   \[ \] A Dataset holds the data; a DataLoader splits it into batches
    
*   \[ \] Transfer learning reuses a pretrained backbone, overriding just the classifier
    
*   \[ \] The multi-class loss function of choice: Cross Entropy Loss
    

## Test Your Understanding

**1\. Why can't an image go straight into a linear layer?**

> **Jawaban:** A linear layer only accepts a 1-dimensional vector, while an image is a matrix. The fix is to flatten it first, or use a convolutional layer instead, which is built to accept matrix input directly.

**2\. If a convolutional layer uses 64 filters, what's its out\_channels?**

> **Jawaban:** 64 — `out_channels` always equals the number of filters used.

**3\. A dataset has 20 fruit classes. How many neurons should the final output layer have?**

> **Jawaban:** 20 — the number of output neurons must match the number of classes.

**4\. What's the difference between the training step and the test/validation step in a training loop?**

> **Jawaban:** They follow the same shape: forward pass, compute loss. The difference is the test/validation step skips optimization (backpropagation and parameter updates), since that step is purely for evaluation, not learning.

**5\. When does a CSV-based approach make more sense than folder-per-class?**

> **Jawaban:** When the dataset carries extra information (metadata, multi-label) or the images live scattered across a directory structure that isn't neatly organized by class. CSV trades a bit of custom code for more flexibility.

* * *

The next post picks up from [here](https://shaka-ai.hashnode.dev/flower-classification-training-fastapi-serving): training a real model to tell 102 flower species apart, all the way through to an API another app can call.

*Part of an ongoing Computer Vision series.*