Lung Cancer ViT Model

Classifies lung CT images as Normal or Cancer using a Vision Transformer (ViT).

Model Details

Dataset: IQ-OTH/NCCD Lung Cancer (~1097 images)
Model: timm.vit_base_patch16_224, fine-tuned
Classes: Normal, Cancer (Benign + Malignant)
Input: 224x224 RGB images
Performance: ~95% test accuracy (see report.txt and classification_report.txt)

Usage

import torch
import timm
import cv2
import numpy as np

# Load model
model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=2)
model.load_state_dict(torch.load('pytorch_model.bin'))
model.eval()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Preprocess image
img = cv2.imread('path_to_image.jpg')
img = cv2.resize(img, (224, 224))
img = img.astype(np.float32) / 255.0
img = torch.tensor(img).permute(2, 0, 1).unsqueeze(0).to(device)

# Predict
with torch.no_grad():
    outputs = model(img)
    probs = torch.softmax(outputs, dim=1)
    predicted_idx = outputs.max(1)[1].item()
    class_names = ['Cancer', 'Normal']
    confidence = probs[0][predicted_idx].item()
    print(f'Classified as: {class_names[predicted_idx]}, Confidence: {confidence:.4f}')

Files

pytorch_model.bin: Model weights
config.json: Model configuration
report.txt: Comprehensive training report
classification_report.txt: Test set classification metrics
confusion_matrix.png: Confusion matrix plot
roc_curve.png: ROC curve with AUC
training_plots.png: Training loss and validation accuracy plots

Training Report

See report.txt and classification_report.txt for details on dataset, hyperparameters, and performance.

Visualizations

Confusion Matrix: confusion_matrix.png
ROC Curve: roc_curve.png
Training Plots: training_plots.png

Medical Disclaimer

For educational purposes only. Consult healthcare professionals for diagnosis.

Inference App

Try the model interactively at Hugging Face Space.