wellallyTech

Posted on Sep 16

Medical AI in Minutes: Building a Skin Lesion Classifier with DINOv2 and Few-Shot Learning 🩺💻

#computervision #machinelearning #programming #ai

In the world of medical AI, data is the ultimate currency—and it's incredibly expensive. When you are dealing with skin lesion detection or identifying rare conditions like melanoma, you rarely have the luxury of 100,000 labeled images. This is where Few-shot Learning and Meta’s DINOv2 change the game.

By leveraging DINOv2, a self-supervised Vision Transformer (ViT), we can extract incredibly rich features from images without needing a single label initially. Combined with a simple Logistic Regression classifier, we can build a surprisingly accurate screening tool with just a handful of examples per class. In this tutorial, we will explore how to implement a medical image classification pipeline using PyTorch, Scikit-learn, and DINOv2 to bridge the gap between pixels and diagnosis.

The Architecture: Why DINOv2?

Traditional supervised models trained on ImageNet often fail in medical contexts because they focus on "objectness" (e.g., "is this a dog?"). Medical images require a focus on texture, edges, and minute color variations.

DINOv2 is trained using self-distillation on a massive curated dataset, making its frozen features "plug-and-play" for downstream tasks. We use it as a powerful feature extractor, then pass those features to a linear head.

graph TD
    A[Input Image: Skin Lesion] --> B[Pre-processing: OpenCV/Resize]
    B --> C[DINOv2 Frozen Backbone]
    C --> D[CLS Token Extraction]
    D --> E[Feature Vector 1x768/1024]
    E --> F[Logistic Regression Classifier]
    F --> G[Prediction: Benign / Malignant]
    G --> H{Clinical Review}

Prerequisites

To follow along, ensure you have the following tech_stack installed:

PyTorch: For loading the DINOv2 model.
Scikit-learn: For the Logistic Regression classifier.
OpenCV: For image preprocessing.
Transformers/PIL: For data handling.

Step 1: Loading the DINOv2 Backbone

Meta provides DINOv2 via torch.hub, making it incredibly easy to load different scales (Small, Base, Large, Giant). For this example, we'll use dinov2_vits14 to keep it lightweight.

import torch

# Load the small DINOv2 model (ViT-S/14)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14').to(device)
model.eval()

print(f"✅ DINOv2 model loaded successfully on {device}!")

Step 2: Feature Extraction Pipeline

We need to transform our skin lesion images into high-dimensional vectors. DINOv2 processes images in 14x14 patches; we will use the [CLS] token as a global representation of the image.

import cv2
import numpy as np
from torchvision import transforms

def extract_features(img_path):
    # Load image with OpenCV
    img = cv2.imread(img_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # Preprocessing: DINOv2 expects 224x224 and specific normalization
    preprocess = transforms.Compose([
        transforms.ToPILImage(),
        transforms.Resize(224),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    input_tensor = preprocess(img).unsqueeze(0).to(device)

    with torch.no_grad():
        # Get the features (CLS token)
        features = model(input_tensor)

    return features.cpu().numpy().flatten()

Step 3: Few-Shot Training with Logistic Regression

Because DINOv2 features are so linearly separable, we don't need a deep neural network for the final step. A simple Logistic Regression from Scikit-learn is less prone to overfitting on small datasets.

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report

# Assume X_train contains extracted features and y_train contains labels
# Labels: 0 for Benign, 1 for Melanoma
X_train = [] # Populated using extract_features()
y_train = [] # Labels

# The "Few-shot" magic
classifier = LogisticRegression(max_iter=1000, C=1.0)
classifier.fit(X_train, y_train)

# Quick prediction logic
def predict_lesion(img_path):
    feat = extract_features(img_path).reshape(1, -1)
    prob = classifier.predict_proba(feat)[0]
    return "Malignant" if prob[1] > 0.5 else "Benign", prob[1]

🥑 Professional Insight: Scaling to Production

While this "Feature Extraction + Linear Head" approach is incredible for prototyping, production-grade medical systems require more robust handling of uncertainty and data drift.

For more advanced architectural patterns and production-ready AI implementations, I highly recommend exploring the deep-dive articles at WellAlly Tech Blog. They cover how to move from experimental notebooks to scalable, HIPAA-compliant AI microservices, which is the natural next step after mastering these Few-shot techniques.

Why This Works (The "Secret Sauce") 🧪

Semantic Consistency: DINOv2 is trained to recognize that different views of the same image should have similar representations. This makes it robust to changes in lighting or camera angles—common issues in smartphone-based dermatology.
No Overfitting: By freezing the DINOv2 backbone, we only train the few thousand parameters of the Logistic Regression model. This prevents the "memorization" of small medical datasets.
Efficiency: You can run this entire inference pipeline on a standard CPU once the features are extracted, making it perfect for edge devices or mobile apps.

Conclusion

We’ve just built a medical screening pipeline that bypasses the need for massive labeled datasets. By combining Meta’s DINOv2 with a Few-shot Learning mindset, we can tackle niche domains in healthcare where data is the biggest bottleneck.

Next Steps:

Try using the dinov2_vitl14 (Large) model for even better accuracy.
Implement UMAP to visualize your skin lesion features in 2D—you'll be amazed at how well they cluster!

What are you building with Vision Transformers? Let me know in the comments below! 🚀

DEV Community