MuhammadBilalCS50/dermamnist_efficientnetv2

1import torch2import torch.nn.functional as F3from PIL import Image4from torchvision import transforms5from torchvision.models import efficientnet_v2_s6import gradio as gr7 8DERMAMNIST_LABELS = {9    0: "actinic keratoses and intraepithelial carcinoma",10    1: "basal cell carcinoma",11    2: "benign keratosis-like lesions",12    3: "dermatofibroma",13    4: "melanoma",14    5: "melanocytic nevi",15    6: "vascular lesions",16}17 18DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")19CHECKPOINT_PATH = "best_model.pth"20 21 22def load_model():23    checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)24    image_size = checkpoint.get("image_size", 224)25    num_classes = checkpoint.get("num_classes", 7)26 27    model = efficientnet_v2_s(weights=None)28    in_features = model.classifier[1].in_features29    model.classifier[1] = torch.nn.Linear(in_features, num_classes)30 31    model.load_state_dict(checkpoint["model_state_dict"])32    model.to(DEVICE)33    model.eval()34    return model, image_size35 36 37MODEL, IMAGE_SIZE = load_model()38TRANSFORM = transforms.Compose([39    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),40    transforms.ToTensor(),41    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),42])43 44 45def predict(image: Image.Image):46    if image is None:47        return {"No image": 1.0}48 49    image = image.convert("RGB")50    x = TRANSFORM(image).unsqueeze(0).to(DEVICE)51 52    with torch.no_grad():53        logits = MODEL(x)54        probs = F.softmax(logits, dim=1).cpu().numpy()[0]55 56    return {DERMAMNIST_LABELS[i]: float(probs[i]) for i in range(len(probs))}57 58 59demo = gr.Interface(60    fn=predict,61    inputs=gr.Image(type="pil", label="Upload a dermatoscopic skin lesion image"),62    outputs=gr.Label(num_top_classes=5, label="Predicted lesion class"),63    title="DermaMNIST Classification with EfficientNetV2-S",64    description=(65        "Upload a skin lesion image. "66        "This demo uses an EfficientNetV2-S model trained from scratch on a randomized subset of DermaMNIST."67    ),68)69 70if __name__ == "__main__":71    demo.launch()72