| 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() | |