Image Classification

Image classification assigns labels to images - is this a cat or a dog, a defect or not, a tumor or healthy tissue. Deep learning revolutionized this field, and transfer learning makes it accessible without massive datasets. Production systems need to handle real-world image variability, scale efficiently, and fail gracefully.

Where It Is Used

Image classification powers visual recognition across industries:

• Manufacturing: Defect detection, quality control
• Healthcare: Medical imaging, pathology, radiology
• Retail: Product recognition, visual search
• Agriculture: Crop disease detection, yield estimation
• Security: Content moderation, threat detection
• Automotive: Road sign recognition, scene understanding

Products: Google Photos, Pinterest Lens, medical diagnostic tools, industrial inspection systems.

Typical Data You Will See

Data sources: - Labeled image datasets - Web-scraped images - Production data with labels - Synthetic data and augmentation - Pre-labeled datasets (ImageNet, COCO)

What You Do in Practice

1. Define classes - What categories need to be distinguished?
2. Collect and label data - Quality labels are critical
3. Preprocess images - Resize, normalize, handle formats
4. Apply augmentation - Increase effective dataset size
5. Choose architecture - Transfer learning vs custom
6. Train model - Fine-tune or train from scratch
7. Evaluate thoroughly - Per-class metrics, error analysis
8. Optimize for deployment - Quantization, pruning, ONNX
9. Deploy with monitoring - Latency, confidence distribution
10. Handle edge cases - Out-of-distribution, adversarial

Baseline First

Start with a pre-trained model and simple fine-tuning:

import torch
import torchvision.models as models
from torchvision import transforms

# Load pre-trained ResNet
model = models.resnet50(weights='IMAGENET1K_V2')

# Replace final layer for your classes
num_classes = 10
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

# Freeze early layers for faster training
for param in model.parameters():
    param.requires_grad = False
for param in model.fc.parameters():
    param.requires_grad = True

# Standard preprocessing
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

Why transfer learning baseline: - Works well with limited data (100s of images) - Fast to train (minutes to hours) - State-of-the-art features from ImageNet - Easy to iterate and improve

Architecture Selection

Selection criteria: - Deployment target (server, mobile, edge) - Latency requirements - Dataset size - Accuracy requirements

Example with EfficientNet:

from torchvision.models import efficientnet_b0, EfficientNet_B0_Weights

# Load with latest weights
weights = EfficientNet_B0_Weights.IMAGENET1K_V1
model = efficientnet_b0(weights=weights)

# Modify classifier
model.classifier[1] = torch.nn.Linear(1280, num_classes)

Data Augmentation

Augmentation is crucial for generalization:

from torchvision import transforms
import albumentations as A

# Basic augmentation
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# Advanced augmentation with Albumentations
train_transform_albu = A.Compose([
    A.RandomResizedCrop(224, 224, scale=(0.8, 1.0)),
    A.HorizontalFlip(p=0.5),
    A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.1, rotate_limit=15),
    A.OneOf([
        A.GaussNoise(var_limit=(10, 50)),
        A.GaussianBlur(blur_limit=3),
        A.MotionBlur(blur_limit=3),
    ], p=0.3),
    A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    A.pytorch.ToTensorV2()
])

Augmentation strategies by domain:

Training Pipeline

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

def train_epoch(model, dataloader, criterion, optimizer, device):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

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

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = outputs.max(1)
        total += labels.size(0)
        correct += predicted.eq(labels).sum().item()

    return running_loss / len(dataloader), correct / total

# Training loop with scheduler
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.1)

for epoch in range(epochs):
    train_loss, train_acc = train_epoch(model, train_loader, 
                                         criterion, optimizer, device)
    val_loss, val_acc = evaluate(model, val_loader, criterion, device)
    scheduler.step()

    print(f"Epoch {epoch}: Train {train_acc:.3f}, Val {val_acc:.3f}")

Evaluation and Metrics

Comprehensive evaluation:

from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns

def evaluate_model(model, dataloader, class_names, device):
    model.eval()
    all_preds = []
    all_labels = []
    all_probs = []

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            outputs = model(images)
            probs = torch.softmax(outputs, dim=1)
            _, preds = outputs.max(1)

            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.numpy())
            all_probs.extend(probs.cpu().numpy())

    # Classification report
    print(classification_report(all_labels, all_preds, 
                                target_names=class_names))

    # Confusion matrix
    cm = confusion_matrix(all_labels, all_preds)
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', 
                xticklabels=class_names, yticklabels=class_names)
    plt.xlabel('Predicted')
    plt.ylabel('True')

    return all_preds, all_labels, all_probs

Failure Modes and Debugging Checklist

• [ ] Data leakage: Same image in train and test (different crops)
• [ ] Label noise: Mislabeled training examples
• [ ] Class imbalance: Model ignores rare classes
• [ ] Domain shift: Training data differs from production
• [ ] Overfitting: Perfect train, poor test performance
• [ ] Adversarial vulnerability: Small perturbations fool model
• [ ] Out-of-distribution: Model confident on unknown classes

Production Checklist

Before deployment: - [ ] Model tested on held-out data matching production - [ ] Per-class performance acceptable - [ ] Model size fits deployment target - [ ] Latency meets requirements - [ ] Confidence calibration checked - [ ] Edge cases documented

After deployment: - [ ] Prediction latency monitored - [ ] Confidence distribution tracked - [ ] Low-confidence predictions flagged - [ ] Periodic accuracy validation - [ ] Model versioning in place

Optimization for deployment:

import torch
import onnx

# Export to ONNX
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, "model.onnx",
    input_names=['image'],
    output_names=['predictions'],
    dynamic_axes={'image': {0: 'batch'}}
)

# Quantization for faster inference
model_quantized = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

Handling Edge Cases

Confidence thresholding:

def predict_with_threshold(model, image, threshold=0.8):
    """Predict with confidence threshold."""
    model.eval()
    with torch.no_grad():
        output = model(image.unsqueeze(0))
        probs = torch.softmax(output, dim=1)
        max_prob, pred_class = probs.max(1)

        if max_prob.item() < threshold:
            return None, max_prob.item()  # Uncertain
        return pred_class.item(), max_prob.item()

Out-of-distribution detection: - Monitor prediction entropy - Track distance from training distribution - Use ensemble disagreement - Implement human review pipeline

Tools and Stack

Example Mini-Case

Scenario: Manufacturing company needs to detect defects in product images (scratch, dent, discoloration, good).

Approach: 1. Data: 10K labeled images (7K good, 1K each defect type) 2. Baseline: ResNet-50 pretrained, fine-tuned - Accuracy: 89% 3. Augmentation: Rotation, brightness, noise - Accuracy: 93% 4. Class weights: Handle imbalance - Recall improved 15% 5. Deployment: ONNX + TensorRT on edge device

Results: - Overall accuracy: 94% - Defect recall: 97% (prioritized over precision) - Latency: 15ms on edge GPU - False positive rate: 3%

Key learnings: - Good lighting critical - trained with lighting augmentation - Edge cases needed human labeling review - Confidence threshold reduced false alarms in production

Related Pages

• Index: ML Field Guide
• Object Detection
• Binary Classification
• Multiclass Classification

Further Reading

• PyTorch Image Models (timm)
• Transfer Learning in Computer Vision
• Albumentations Documentation
• Model Deployment with ONNX

Try it live

Everything above runs in your browser — open Convolutional Neural Network Simulator and change the parameters while it is running. Nothing is installed, nothing is uploaded, the whole model lives in one tab.