sukhmani1303/lung-cancer-vit-model

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.