update

README.md

@@ -1,10 +1,9 @@
 ---
 title: Masksup Segmentation Demo
-emoji: 👁
+emoji: 🛋️ 🛏️ 📦
 colorFrom: green
 colorTo: gray
 sdk: gradio
-sdk_version: 3.14.0
 app_file: app.py
 pinned: false
 license: afl-3.0

app.py

@@ -19,18 +19,22 @@ DEVICE = "cpu"
 print(DEVICE)
 
 # Load color map
-cmap = np.load('cmap.npy')
+cmap = np.load("cmap.npy")
 
 # Make directories
 os.system("mkdir ./models")
 
 # Get model weights
 if not os.path.exists("./models/masksupnyu39.31d.pth"):
-    os.system("wget -O ./models/masksupnyu39.31d.pth https://github.com/hasibzunair/masksup-segmentation/releases/download/v0.1/masksupnyu39.31iou.pth")
+    os.system(
+        "wget -O ./models/masksupnyu39.31d.pth https://github.com/hasibzunair/masksup-segmentation/releases/download/v0.1/masksupnyu39.31iou.pth"
+    )
 
 # Load model
 model = NestedUNet(num_classes=40)
-checkpoint = torch.load("./models/masksupnyu39.31d.pth", map_location=torch.device('cpu'))
+checkpoint = torch.load(
+    "./models/masksupnyu39.31d.pth", map_location=torch.device("cpu")
+)
 model.load_state_dict(checkpoint)
 model = model.to(DEVICE)
 model.eval()
@@ -49,14 +53,23 @@ def inference(img_path):
     )
 
     image = transforms_image(image)
-    image = image[None, :]
+    image = image[None, :]  # batch dimension
+
     # Predict
     with torch.no_grad():
         output = torch.sigmoid(model(image.to(DEVICE).float()))
-        output = torch.softmax(output, dim=1).argmax(dim=1)[0].float().cpu().numpy().astype(np.uint8)
+        output = (
+            torch.softmax(output, dim=1)
+            .argmax(dim=1)[0]
+            .float()
+            .cpu()
+            .numpy()
+            .astype(np.uint8)
+        )
         pred = cmap[output]
     return pred
 
+
 # App
 title = "Masked Supervised Learning for Semantic Segmentation"
 description = codecs.open("description.html", "r", "utf-8").read()
@@ -64,13 +77,17 @@ article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2210.009
 
 gr.Interface(
     inference,
-    gr.inputs.Image(type='filepath', label="Input Image"),
+    gr.inputs.Image(type="filepath", label="Input Image"),
     gr.outputs.Image(type="filepath", label="Predicted Output"),
-    examples=["./sample_images/a.png", "./sample_images/b.png", 
-    "./sample_images/c.png", "./sample_images/d.png"],
+    examples=[
+        "./sample_images/a.png",
+        "./sample_images/b.png",
+        "./sample_images/c.png",
+        "./sample_images/d.png",
+    ],
     title=title,
     description=description,
     article=article,
     allow_flagging=False,
     analytics_enabled=False,
-    ).launch(debug=True, enable_queue=True)
+).launch(debug=True, enable_queue=True)