PBC_VIT

Peripheral Blood Cell (PBC) Classification using Vision Transformers

This repository contains PyTorch model weights (teacher_model.pth and student_model.pth) for 8-type Peripheral Blood Cell classification. The models leverage Vision Transformer (ViT) architectures and Knowledge Distillation (KD) for efficient cellular analysis.

Overview

• Task: 8-Class Peripheral Blood Cell Image Classification (basophil, eosinophil, erythroblast, ig, lymphocyte, monocyte, neutrophil, platelet).
• Dataset: 17,092 RGB images.
• Teacher Model: Custom Vision Transformer trained from scratch.
• Student Model: Distilled vit_tiny_patch16_224 (5M parameters, ~22MB).

Model Performance

Class Labels Map

label2id = {
    "basophil": 0,
    "eosinophil": 1,
    "erythroblast": 2,
    "ig": 3,
    "lymphocyte": 4,
    "monocyte": 5,
    "neutrophil": 6,
    "platelet": 7,
}
id2label = {v: k for k, v in label2id.items()}

Usage

Loading the Models via timm / PyTorch

Load Distilled Student Model (student_model.pth)

import timm
import torch
from torchvision import transforms

# 1. Initialize Student Architecture
model = timm.create_model("vit_tiny_patch16_224", pretrained=False, num_classes=8)

# 2. Load Weight Checkpoint
state_dict = torch.load("student_model.pth", map_location="cpu")
model.load_state_dict(state_dict)
model.eval()

# 3. Preprocessing
transform = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        ),
    ]
)

Running Inference

from PIL import Image

image = Image.open("blood_cell_image.jpg").convert("RGB")
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    logits = model(input_tensor)
    probs = torch.nn.functional.softmax(logits, dim=-1)
    predicted_class_id = torch.argmax(probs, dim=-1).item()

print(
    f"Predicted Class: {id2label[predicted_class_id]} ({probs[0][predicted_class_id]*100:.2f})"
)

Model Training & Distillation Details

• Loss Function: Focal Loss (gamma=0.7, alpha=2) for class imbalance handling.
• Knowledge Distillation Loss: Combined Cross-Entropy loss and KL-Divergence on temperature-scaled soft probabilities.
• Optimization: Adam / Early stopping triggered on validation loss stabilization.