Update app.py

app.py

@@ -1,14 +1,13 @@
 import torch
 import torch.nn as nn
-import torchvision.transforms as transforms
-from torchvision import models
+from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
 
-# Class names
+# Define class labels
 class_names = ["Normal", "Cancer", "Malignant"]
 
-# Image preprocessing
+# Define preprocessing for CT scan images
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -16,35 +15,29 @@ transform = transforms.Compose([
                          std=[0.229, 0.224, 0.225])
 ])
 
-# Load model architecture
+# Define and load model
 def get_model():
     model = models.vgg16(pretrained=False)
     model.classifier[6] = nn.Linear(4096, 3)  # 3 output classes
     return model
 
-# Load trained weights
 model = get_model()
 model.load_state_dict(torch.load("distilled_vgg16.pth", map_location=torch.device("cpu")))
 model.eval()
 
-# Prediction function
+# Define prediction function
 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)
-        probabilities = torch.softmax(output, dim=1)[0]
-        predicted_class = torch.argmax(probabilities).item()
-        confidence = probabilities[predicted_class].item()
-    return f"Likely: {class_names[predicted_class]} (Confidence: {confidence * 100:.2f}%)"
+        outputs = model(image)
+        probabilities = torch.softmax(outputs, dim=1)[0]
+        top_class = torch.argmax(probabilities).item()
+        confidence = probabilities[top_class].item()
+    return f"Likely: {class_names[top_class]} (Confidence: {confidence*100:.2f}%)"
 
-# Gradio Interface
+# Gradio UI
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
     outputs=gr.Textbox(label="Prediction"),
-    title="Lung Cancer Detector (PyTorch VGG16)",
-    description="Upload a CT scan image to predict: Normal, Cancer, or Malignant."
-)
-
-if __name__ == "__main__":
-    interface.launch()
+    title="Lung Cancer Classifier (PyTorch VGG1