Update app.py

app.py

@@ -1,45 +1,38 @@
 import torch
-import torch.nn as nn
-from torchvision import models, transforms
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
 import gradio as gr
 
-# 🚨 Use class order from training (check ImageFolder class_to_idx)
+# Class labels in the same order as training
 class_names = ["Normal", "Cancer", "Malignant"]
 
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406],
-                         [0.229, 0.224, 0.225])
-])
+# Load the ViT feature extractor
+extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
 
-# Load model
-def get_model():
-    model = models.vgg16(pretrained=False)
-    model.classifier[6] = nn.Linear(4096, 3)
-    return model
-
-model = get_model()
-model.load_state_dict(torch.load("distilled_vgg16.pth", map_location="cpu"))
+# Load the ViT model and custom weights
+model = AutoModelForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224-in21k", num_labels=3
+)
+model.classifier = torch.nn.Linear(model.classifier.in_features, 3)  # Replace classifier
+model.load_state_dict(torch.load("distill_vgg16.pth", map_location="cpu"))  # Your uploaded file
 model.eval()
 
-# Predict
+# Prediction function
 def predict(img: Image.Image):
-    image = transform(img).unsqueeze(0)
+    inputs = extractor(images=img, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(image)
-        probs = torch.softmax(outputs, dim=1)[0]
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=1)[0]
         pred = torch.argmax(probs).item()
     return f"Prediction: {class_names[pred]}\nPlease consult a doctor for further diagnosis."
 
-# Gradio UI
+# Gradio interface
 interface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
+    inputs=gr.Image(type="pil", label="Upload CT Image"),
     outputs=gr.Textbox(label="Diagnosis"),
-    title="Lung Cancer Detector (VGG16 + KD)",
-    description="Upload a CT scan. Model will predict: Normal, Cancer, or Malignant."
+    title="Lung Cancer Classifier",
+    description="This Vision Transformer model predicts lung conditions from CT scans. Classes: Normal, Cancer, Malignant."
 )
 
 if __name__ == "__main__":