| import gradio as gr
|
| import torch
|
| import torchvision.transforms as transforms
|
| from PIL import Image
|
| import torchvision.models as models
|
| import torch.nn as nn
|
|
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
|
| model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
|
| in_features = model.fc.in_features
|
| model.fc = nn.Sequential(
|
| nn.Linear(in_features, 512),
|
| nn.ReLU(),
|
| nn.Dropout(0.4),
|
| nn.Linear(512, 2)
|
| )
|
|
|
|
|
| model.load_state_dict(torch.load("best_model (2).pth", map_location=device))
|
| model.to(device)
|
| model.eval()
|
|
|
|
|
| transform = transforms.Compose([
|
| transforms.Lambda(lambda img: img.convert("RGB")),
|
| transforms.Resize((224, 224)),
|
| transforms.ToTensor(),
|
| transforms.Normalize(mean=[0.5]*3, std=[0.5]*3)
|
| ])
|
|
|
|
|
| class_names = ["NORMAL", "PNEUMONIA"]
|
|
|
|
|
| def classify_image(img):
|
| img = transform(img).unsqueeze(0).to(device)
|
| with torch.no_grad():
|
| outputs = model(img)
|
| probs = torch.nn.functional.softmax(outputs, dim=1)
|
| return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}
|
|
|
|
|
| interface = gr.Interface(
|
| fn=classify_image,
|
| inputs=gr.Image(type="pil"),
|
| outputs=gr.Label(num_top_classes=2),
|
| title="🩺 Pneumonia Classifier",
|
| description="Upload a chest X-ray image. The model predicts whether it's NORMAL or shows signs of PNEUMONIA."
|
| )
|
|
|
|
|
| interface.launch()
|
|
|
