gradio app

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:system"
+}

app.py

@@ -0,0 +1,55 @@
+import torch
+from torchvision.transforms import v2
+import gradio
+from PIL import Image
+from mae_model.models.linear_predictor import Predictor
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+LABELS = [
+    "adipose",
+    "background",
+    "debris",
+    "lymphocytes",
+    "mucus",
+    "smooth muscle",
+    "normal colon mucosa",
+    "cancer-associated stroma",
+    "colorectal adenocarcinoma epithelium",
+]
+
+
+model = Predictor(n_labels=len(LABELS))
+model.load_state_dict(torch.load("checkpoints/model_linear_v2_epoch_100.pt"))
+model.to(device)
+model.eval()
+
+
+tf = v2.Compose([
+    v2.ToImage(),
+    v2.ToDtype(torch.float32, scale=True),
+])
+
+def predict(image):
+    if image is None:
+        return None
+
+    pil_img = Image.fromarray(image.astype('uint8'), 'RGB')
+    img_tensor = tf(pil_img).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        outputs = model(img_tensor)
+        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+    
+    return {LABELS[i]: float(probabilities[i]) for i in range(len(LABELS))}
+
+demo = gradio.Interface(
+    fn=predict,
+    inputs=gradio.Image(), 
+    outputs=gradio.Label(num_top_classes=3),
+    title="OrganMNIST Image Classification",
+    description="Upload a tissue patch image for classification"
+)
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+gradio
+torch
+torchvision
+pillow
+git+https://github.com/HasanAli5/MAE-Model-MedMNIST-Predictor.git#egg=mae_model