Spaces:
Sleeping
Sleeping
File size: 1,467 Bytes
2b85da0 37f5c64 2b85da0 794c694 2b85da0 37f5c64 794c694 2b85da0 37f5c64 2b85da0 794c694 2b85da0 37f5c64 794c694 2b85da0 794c694 37f5c64 794c694 37f5c64 2b85da0 37f5c64 72c8baa 37f5c64 72c8baa 3c1db92 2dec519 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
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()
|