Post

Image Classification

Posted Updated
By
4 min read
Image Classification

Project

Table of contents

Overview

This project demonstrates how to implement an image classifier using a deep learning framework. The classifier is trained to recognize medical images from the MedMNIST dataset, specifically PathMNIST. The workflow includes data loading, preprocessing, model definition, training, evaluation, and saving the model.

Details of Dataset

  • Dataset Used: MedMNIST (PathMNIST)
  • Source: MedMNIST GitHub
  • Description: PathMNIST contains histopathology images for multi-class classification. Each image is labeled according to tissue type or pathology.
  • Number of Classes: 9 (as set in the code)
  • Channels: 3 (RGB images)
1
2
3
4
5
6
7
8
9
10

# Example: Load MedMNIST PathMNIST dataset
import medmnist
from medmnist import INFO

title: "Medical Image Classification"
info = INFO[data_flag]
DataClass = getattr(medmnist, info['python_class'])

train_dataset = DataClass(split='train', transform=data_transform, download=True)
test_dataset = DataClass(split='test', transform=data_transform, download=True)

Data Preparation

First need to do normalization since it ensures that the input data has a consistent scale, which helps the neural network train more efficiently and converge faster. It prevents features with larger values from dominating the learning process and helps gradients flow better during backpropagation.

Which method of normalization is used?
Standardization is used, where each channel is normalized to have a mean of 0.5 and a standard deviation of 0.5:

1
2
3
4
5
6

import torchvision.transforms as transforms

date: 2025-10-10
  transforms.ToTensor(),
  transforms.Normalize(mean=[.5], std=[.5])
])

This scales pixel values from [0, 1] to [-1, 1].

Architecture of Network

Defining neural network:
The model is defined as a subclass of nn.Module with two main parts: convolutional layers for feature extraction and fully connected layers for classification.

Architecture of the used CNN

  • 3 convolutional layers:
    • Conv2d(3 → 8), kernel size 3, ReLU, MaxPool2d
    • Conv2d(8 → 16), kernel size 3, ReLU, MaxPool2d
    • Conv2d(16 → 32), kernel size 3, ReLU, MaxPool2d
  • Flattening layer
  • 2 fully connected layers:
    • Linear(32 → 16), ReLU
    • Linear(16 → 9) (output for 9 classes)

This structure is simple yet effective for small medical image datasets. The convolutional layers extract spatial features, pooling reduces dimensionality, and the fully connected layers perform classification. It balances learning capacity and computational efficiency.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

import torch.nn as nn
import torch.nn.functional as F

class MedMNISTModel(nn.Module):
  def __init__(self):
    super(MedMNISTModel, self).__init__()
    self.conv_model = nn.Sequential(
      nn.Conv2d(in_channels=3, out_channels=8, kernel_size=3),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2, stride=2),
      nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2, stride=2),
      nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3),
      nn.ReLU(),
      nn.MaxPool2d(kernel_size=2, stride=2)
    )
    self.classification_model = nn.Sequential(
      nn.Linear(in_features=32, out_features=16),
      nn.ReLU(),
      nn.Linear(in_features=16, out_features=9),
    )
    self.optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
    self.criterion = nn.CrossEntropyLoss()

  def forward(self, x):
    x = self.conv_model(x)
    x = torch.flatten(x, start_dim=1)
    x = self.classification_model(x)
    output = F.log_softmax(x, dim=1)
    return x

model = MedMNISTModel()

Training and Test

Loss function
CrossEntropyLoss is used, which is standard for multi-class classification. It measures the difference between predicted probabilities and true labels.

Optimizer
Adam optimizer (optim.Adam) is used for its adaptive learning rate and efficient convergence.

Settings
The main training loop uses 10 epochs (EPOCHS = 10). This can be adjusted for better performance.

Other settings regarding the CNN:

  • Batch size: 8 (controls how many samples are processed before updating the model)
  • Learning rate: 0.001 (determines the step size for each update)
  • Input channels: 3 (RGB images)
  • Number of classes: 9 (output layer size)
  • Device: Uses GPU (CUDA) if available, otherwise CPU
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# Training loop
import torch

EPOCHS = 10
device = "cuda:0" if torch.cuda.is_available() else "cpu"
model.to(device)

for epoch in range(EPOCHS):
  running_loss = 0.0
  for i, data in enumerate(train_loader):
    imgs, targets = data[0].to(device), data[1].to(device)
    outputs = model(imgs)
    loss = model.criterion(outputs, targets.squeeze())
    model.optimizer.zero_grad()
    loss.backward()
    model.optimizer.step()
    running_loss += loss.item()
    if i % 2000 == 1999:
      print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
      running_loss = 0.0

Evaluation

  • Accuracy is calculated using sklearn.metrics.accuracy_score.
  • Confusion matrix is visualized using seaborn.heatmap.
  • These metrics provide insight into model performance and class-wise prediction quality.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# Evaluation
predicted = []
targets = []
with torch.no_grad():
  for i, data in enumerate(test_loader):
    img, target = data[0].to(device), data[1].item()
    output = torch.argmax(model(img).cpu(), dim=1).detach().numpy()
    predicted.append(output)
    targets.append(target)

from sklearn.metrics import accuracy_score, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns

accuracy = accuracy_score(targets, predicted) * 100
print(f"Accuracy of the model is {accuracy:.2f}%")
cf_mt = confusion_matrix(targets, predicted)
plt.figure(figsize=(10, 10))
sns.heatmap(cf_mt, annot=True)

Additional Sections

Saving the Model:
The trained model is saved using:

1

torch.save(model.state_dict(), "model.pth")

Hardware Utilization:
The code checks for GPU availability and uses CUDA if possible for faster training.

Visualization:
Montage of training images and confusion matrix are visualized for better understanding.

Conclusion

This project covers the essential steps for building a medical image classifier:

  • Data loading and normalization
  • Model architecture design
  • Training and evaluation
  • Model saving and visualization

You can extend this project by experimenting with different architectures, datasets, and hyperparameters.

Resources

Project repository

GitHub Code: Medical Image Classification, Segmentation using U-Net, and Tabular Data Classification

Machine Learning and Deep Learning, Machine Learning, Computer Vision
image classification deep learning CNN computer vision medical imaging evaluation
This post is licensed under CC BY 4.0 by the author.