Initial commit

.gitattributes

@@ -19,6 +19,7 @@
 *.pb filter=lfs diff=lfs merge=lfs -text
 *.pickle filter=lfs diff=lfs merge=lfs -text
 *.pkl filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,8 +1,8 @@
 ---
 title: Diabetic Retinopathy
-emoji: 🌍
-colorFrom: purple
-colorTo: pink
+emoji: 🐠
+colorFrom: pink
+colorTo: purple
 sdk: gradio
 sdk_version: 3.9.1
 app_file: app.py

app.py

@@ -0,0 +1,62 @@
+import gradio as gr
+import timm
+import torch 
+import torch.nn as nn
+
+
+class Net2D(nn.Module):
+
+    def __init__(self, weights):
+        super().__init__()
+        self.backbone = timm.create_model("tf_efficientnet_b6_ns", pretrained=False, global_pool="", num_classes=0)
+        self.pool_layer = nn.AdaptiveAvgPool2d(1)
+        self.dropout = nn.Dropout(0.5)
+        self.linear = nn.Linear(2304, 5)
+        self.load_state_dict(weights)
+
+    def forward(self, x):
+        x = self.backbone(x)
+        x = self.pool_layer(x).view(x.size(0), -1)
+        x = self.dropout(x)
+        x = self.linear(x)
+        return x[:, 0] if x.size(1) == 1 else x
+
+
+weights = torch.load("model0.ckpt", map_location=torch.device("cpu"))["state_dict"]
+weights = {k.replace("model.", "") : v for k, v in weights.items()}
+model = Net2D(weights)
+
+
+def predict(Image):
+    img = torch.from_numpy(Image)
+    img = img[:, :, [2, 1, 0]]
+    img = img.permute(2, 0, 1)
+    img = img.unsqueeze(0)
+    img = img / img.max()
+    img = img - 0.5
+    img = img * 2.0
+    with torch.no_grad():
+        grade = torch.softmax(model(img.float()), dim=1)[0]
+    cats = ["None", "Mild", "Moderate", "Severe", "Proliferative"]
+    output_dict = {}
+    for cat, value in zip(cats, grade):
+        output_dict[cat] = value.item()
+    return output_dict
+
+
+image = gr.Image(shape=(512, 512), image_mode="RGB")
+label = gr.Label(label="Grade")
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=image,
+    outputs=label,
+    examples=["examples/none.png", "examples/mild.png", "examples/moderate.png", "examples/severe.png",
+              "examples/proliferative.png"]
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(debug=True)
+
+

model0.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b44c74a6c75f747e48799bb9c20510bd27a6e900f33d5fd795955f59ce78eac4
+size 164258381

requirements.txt

@@ -0,0 +1,4 @@
+omegaconf
+timm
+torch
+