Update

app.py

@@ -22,19 +22,23 @@ class Net2D(nn.Module):
         return x[:, 0] if x.size(1) == 1 else x
 
 
+def rescale(x):
+    x = x / 255.0
+    x = x - 0.5
+    x = x * 2.0
+    return 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)
+weights = {k.replace("model.", ""): v for k, v in weights.items()}
+model = Net2D(weights).eval()
 
 
 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
+    img = rescale(img)
     with torch.no_grad():
         grade = torch.softmax(model(img.float()), dim=1)[0]
     cats = ["None", "Mild", "Moderate", "Severe", "Proliferative"]
@@ -44,6 +48,7 @@ def predict(Image):
     return output_dict
 
 
+
 image = gr.Image(shape=(512, 512), image_mode="RGB")
 label = gr.Label(label="Grade")
 
@@ -51,12 +56,11 @@ 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"]
+    examples=["examples/0.png", "examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png"]
     )
 
 
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    demo.launch(debug=True, share=True)