app.py

import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import gradio as gr

# Define class labels
class_names = ["Normal", "Cancer", "Malignant"]

# Define preprocessing for CT scan images
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# Define and load model
def get_model():
    model = models.vgg16(pretrained=False)
    model.classifier[6] = nn.Linear(4096, 3)  # 3 output classes
    return model

model = get_model()
model.load_state_dict(torch.load("distilled_vgg16.pth", map_location=torch.device("cpu")))
model.eval()

# Define prediction function
def predict(img: Image.Image):
    image = transform(img).unsqueeze(0)
    with torch.no_grad():
        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 UI
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Prediction"),
    title="Lung Cancer Classifier (PyTorch VGG16)",
    description="Upload a lung CT scan image to classify it as Normal, Cancer, or Malignant."
)

if __name__ == "__main__":
    interface.launch()