TurkishCodeMan/vit-lung-cancer
114
🫁 ViT Lung Cancer Classifier
Fine-tuned Vision Transformer (ViT-Base/16) for lung cancer CT image classification into 3 classes: normal, malignant, and benign.
📊 Model Details
🏷️ Label Mapping
📅 Dataset
The model was trained on a comprehensive lung cancer dataset containing global clinical and risk factor data.
🚀 Usage
Install
pip install transformers torch pillowInference
from transformers import ViTForImageClassification, ViTImageProcessor
from PIL import Image
import torch
model_id = "TurkishCodeMan/vit-lung-cancer"
processor = ViTImageProcessor.from_pretrained(model_id)
model = ViTForImageClassification.from_pretrained(model_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.eval().to(device)
def predict(image_path: str) -> dict:
img = Image.open(image_path).convert("RGB")
inputs = processor(images=img, return_tensors="pt").to(device)
with torch.no_grad():
logits = model(**inputs).logits
pred_id = logits.argmax(-1).item()
probs = torch.softmax(logits.float(), dim=-1)[0]
return {
"prediction": model.config.id2label[pred_id],
"probabilities": {
label: round(probs[i].item(), 4)
for i, label in model.config.id2label.items()
}
}
result = predict("lung_scan.jpg")
print(result)