app.py

import gradio as gr
import torch
import torchvision.transforms as transforms
from medmnist import INFO
from model import load_model
from PIL import Image


info = INFO["dermamnist"]
class_names = list(info["label"].values())


model = load_model()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])


def predict(image):
    if image is None:
        return {"Error": 1.0}
    image = image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0)

    with torch.no_grad():
        outputs = model(input_tensor)
        probs = torch.softmax(outputs, dim=1).squeeze().numpy()

    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="Skin Disease Classifier",
    description="Upload a skin image and our model will predict potential skin cancer(melanoma), tumor or moles using EfficientNet-B2 fine-tuned on DermMNIST."
)

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