Update app.py

app.py

@@ -5,10 +5,10 @@ from torchvision import models
 from PIL import Image
 import gradio as gr
 
-# Class labels (must match model output)
+# Define class names
 class_names = ["Normal", "Cancer", "Malignant"]
 
-# Define transform for CT scan input
+# Image transform
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -16,7 +16,7 @@ transform = transforms.Compose([
                          std=[0.229, 0.224, 0.225])
 ])
 
-# Load the VGG16 model structure and adjust classifier
+# Load VGG16 model and modify final layer
 def get_model():
     model = models.vgg16(pretrained=False)
     model.classifier[6] = nn.Linear(4096, 3)  # 3 output classes
@@ -27,22 +27,23 @@ model = get_model()
 model.load_state_dict(torch.load("distilled_vgg16.pth", map_location=torch.device("cpu")))
 model.eval()
 
-# Prediction function
+# Prediction function — return only top class
 def predict(img: Image.Image):
-    image = transform(img).unsqueeze(0)  # Add batch dimension
+    image = transform(img).unsqueeze(0)
     with torch.no_grad():
         output = model(image)
-        probs = torch.softmax(output, dim=1)[0]  # Convert logits to probabilities
-    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
+        predicted_class = torch.argmax(output, dim=1).item()
+    return f"Likely: {class_names[predicted_class]}"
 
-# Gradio UI
+# Gradio Interface
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Label(num_top_classes=3),
+    outputs=gr.Textbox(label="Prediction"),
     title="Lung Cancer Classifier (PyTorch)",
-    description="Upload a CT scan image to classify it as Normal, Cancer, or Malignant."
+    description="Upload a lung CT scan. The model predicts the most likely condition: Normal, Cancer, or Malignant."
 )
 
 if __name__ == "__main__":
     interface.launch()
+