formatting, more examples

app.py

@@ -13,76 +13,92 @@ import logging
 
 args = {
     "model_path": "model_last_epoch_34_torchvision0_3_state.ptw",
-    "device": torch.device('cuda' if torch.cuda.is_available() else 'cpu'),
-    "dxlabels": ["akiec", "bcc", "bkl", "df", "mel", "nv","vasc"]
+    "device": torch.device("cuda" if torch.cuda.is_available() else "cpu"),
+    "dxlabels": ["akiec", "bcc", "bkl", "df", "mel", "nv", "vasc"],
 }
 logging.warning(args)
 
+
 # No specific normalization was performed during training
 def normtransform(x):
     return x
 
+
 # Load model
 logging.warning("Loading model...")
 model = torchvision.models.resnet34()
-model.fc = torch.nn.Linear(model.fc.in_features, len(args.get("dxlabels")))
-model.load_state_dict(torch.load(args.get("model_path")))
+model.fc = torch.nn.Linear(model.fc.in_features, len(args["dxlabels"]))
+model.load_state_dict(torch.load(args["model_path"]))
 model.eval()
-model.to(device=args.get("device"))
+model.to(device=args["device"])
 torch.set_grad_enabled(False)
 logging.warning("Model loaded.")
 
-def predict(image: str)->dict:
+
+def predict(image: str) -> dict:
     global model, args
 
     logging.warning(f"Starting predict('{image}') ...")
-    prediction_tensor = torch.zeros([1, len(args.get("dxlabels"))]).to(device=args.get("device"))
+    prediction_tensor = torch.zeros([1, len(args["dxlabels"])]).to(
+        device=args["device"]
+    )
 
     # Test-time augmentations
     available_sizes = [224]
     target_sizes, hflips, rotations, crops = available_sizes, [0, 1], [0, 90], [0.8]
     aug_combos = [x for x in itertools.product(target_sizes, hflips, rotations, crops)]
-    
+
     # Load image
     img = Image.open(image)
-    img = img.convert('RGB')
+    img = img.convert("RGB")
 
     # Predict with Test-time augmentation
-    for (target_size, hflip, rotation, crop) in aug_combos:
-        tfm = transforms.Compose([
-            transforms.Resize(int(target_size // crop)),
-            transforms.CenterCrop(target_size),
-            transforms.RandomHorizontalFlip(hflip),
-            transforms.RandomRotation([rotation, rotation]),
-            transforms.ToTensor(),
-            # shades_of_grey_torch,
-            normtransform
-        ])
-        test_data = tfm(img).unsqueeze(0).to(device=args.get("device"))
-        running_preds = torch.FloatTensor().to(device=args.get("device"))
-        outputs = model(test_data)
-        prediction_tensor += running_preds
-    
+    for target_size, hflip, rotation, crop in aug_combos:
+        tfm = transforms.Compose(
+            [
+                transforms.Resize(int(target_size // crop)),
+                transforms.CenterCrop(target_size),
+                transforms.RandomHorizontalFlip(hflip),
+                transforms.RandomRotation([rotation, rotation]),
+                transforms.ToTensor(),
+                normtransform,
+            ]
+        )
+        test_data = tfm(img).unsqueeze(0).to(device=args["device"])
+        with torch.no_grad():
+            outputs = model(test_data)
+        prediction_tensor += outputs
+
     prediction_tensor /= len(aug_combos)
-    predictions = F.softmax(prediction_tensor, dim=1)[0].cpu().numpy()
+    predictions = F.softmax(prediction_tensor, dim=1)[0].detach().cpu().tolist()
     logging.warning(f"Returning {predictions=}")
-    return {args.get("dxlabels")[enu]: p for enu, p in enumerate(predictions)}
+    return {args["dxlabels"][enu]: p for enu, p in enumerate(predictions)}
+
+
+description = """
+Research image classification model for multi-class predictions of common dermatologic tumors, the model was trained 
+on the [HAM10000 dataset](https://www.nature.com/articles/sdata2018161).
+
+This is the model used in the publication 
+[Tschandl P. et al. Nature Medicine 2020](https://www.nature.com/articles/s41591-020-0942-0)
+where human-computer interaction of such a system was analyzed.
 
-description = '''Research artifact for multi-class predictions of common 
-dermatologic tumors. This is the model used in the publication 
-[Tschandl P. et al. Nature Medicine 2020](https://www.nature.com/articles/s41591-020-0942-0).
+Instructions for uploading: Ensure the image is not blurred, the lesion is centered and in focus, and no black/white
+vignette is in the surrounding. The image should depict the whole lesion, and not a zoomed-in part.
 
-For education and research use only. 
-**DO NOT use this to obtain medical advice!**
-If you have a skin change in question, seek contact to your physician.'''
+For education and research use only. **DO NOT use this to obtain medical advice!**
+If you have a skin change in question, seek contact to your physician."""
 
 logging.warning("Starting Gradio interface...")
 gr.Interface(
     predict,
     inputs=gr.Image(label="Upload a dermatoscopic image", type="filepath"),
-    outputs=gr.Label(num_top_classes=len(args.get("dxlabels"))),
+    outputs=gr.Label(num_top_classes=len(args["dxlabels"])),
     title="Dermatoscopic classification",
     description=description,
-    allow_flagging="manual",
-    examples=[os.path.join(os.path.dirname(__file__), "ISIC_0024306.jpg")]
+    allow_flagging="never",
+    examples=[
+        os.path.join(os.path.dirname(__file__), "images", x)
+        for x in ["ISIC_0024306.jpg", "ISIC_0024315.jpg", "ISIC_0024318.jpg"]
+    ],
 ).launch()