Spaces:

hasibzunair
/

masksup-segmentation-demo

Running

App Files Files Community

hasibzunair commited on Dec 20, 2022

Commit

8a3583d

•

1 Parent(s): 3a4add5

add app files

Browse files

Files changed (9) hide show

cmap.npy +3 -0
description.html +10 -0
nyu.ipynb +165 -0
requirements.txt +6 -0
sample_images/a.png +0 -0
sample_images/b.png +0 -0
sample_images/c.png +0 -0
sample_images/d.png +0 -0
unetplusplus.py +142 -0

cmap.npy ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e5648506a4b5dbeb787e93f26b429cab659c3b66a4d579645edb2f24ba41a919
+size 848

description.html ADDED Viewed

	@@ -0,0 +1,10 @@

+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <title>Title</title>
+</head>
+<body>
+    This is a demo of our BMVC'2022 Oral paper <a href="https://arxiv.org/abs/2210.00923">Masked Supervised Learning for Semantic Segmentation</a>.</br>
+</body>
+</html>

nyu.ipynb ADDED Viewed

	@@ -0,0 +1,165 @@

+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import numpy as np\n",
+    "import cv2\n",
+    "import codecs\n",
+    "import torch\n",
+    "import torchvision.transforms as transforms\n",
+    "import gradio as gr\n",
+    "\n",
+    "from PIL import Image\n",
+    "\n",
+    "from unetplusplus import NestedUNet\n",
+    "\n",
+    "torch.manual_seed(0)\n",
+    "\n",
+    "if torch.cuda.is_available():\n",
+    "    torch.backends.cudnn.deterministic = True\n",
+    "\n",
+    "# Device\n",
+    "DEVICE = \"cpu\"\n",
+    "print(DEVICE)\n",
+    "\n",
+    "# Load color map\n",
+    "cmap = np.load('cmap.npy')\n",
+    "\n",
+    "# Make directories\n",
+    "os.system(\"mkdir ./models\")\n",
+    "\n",
+    "# Get model weights\n",
+    "if not os.path.exists(\"./models/masksupnyu39.31d.pth\"):\n",
+    "    os.system(\"wget -O ./models/masksupnyu39.31d.pth https://github.com/hasibzunair/masksup-segmentation/releases/download/v0.1/masksupnyu39.31iou.pth\")\n",
+    "\n",
+    "# Load model\n",
+    "model = NestedUNet(num_classes=40)\n",
+    "checkpoint = torch.load(\"./models/masksupnyu39.31d.pth\")\n",
+    "model.load_state_dict(checkpoint)\n",
+    "model = model.to(DEVICE)\n",
+    "model.eval()\n",
+    "\n",
+    "\n",
+    "# Main inference function\n",
+    "def inference(img_path):\n",
+    "    image = Image.open(img_path).convert(\"RGB\")\n",
+    "    transforms_image = transforms.Compose(\n",
+    "        [\n",
+    "            transforms.Resize((224, 224)),\n",
+    "            transforms.CenterCrop((224, 224)),\n",
+    "            transforms.ToTensor(),\n",
+    "            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
+    "        ]\n",
+    "    )\n",
+    "\n",
+    "    image = transforms_image(image)\n",
+    "    image = image[None, :]\n",
+    "    # Predict\n",
+    "    with torch.no_grad():\n",
+    "        output = torch.sigmoid(model(image.to(DEVICE).float()))\n",
+    "        output = torch.softmax(output, dim=1).argmax(dim=1)[0].float().cpu().numpy().astype(np.uint8)\n",
+    "        pred = cmap[output]\n",
+    "    return pred\n",
+    "\n",
+    "# App\n",
+    "title = \"Masked Supervised Learning for Semantic Segmentation\"\n",
+    "description = codecs.open(\"description.html\", \"r\", \"utf-8\").read()\n",
+    "article = \"<p style='text-align: center'><a href='https://arxiv.org/abs/2210.00923' target='_blank'>Masked Supervised Learning for Semantic Segmentation</a> | <a href='https://github.com/hasibzunair/masksup-segmentation' target='_blank'>Github</a></p>\"\n",
+    "\n",
+    "gr.Interface(\n",
+    "    inference,\n",
+    "    gr.inputs.Image(type='file', label=\"Input Image\"),\n",
+    "    gr.outputs.Image(type=\"file\", label=\"Predicted Output\"),\n",
+    "    examples=[\"./sample_images/a.png\", \"./sample_images/b.png\", \n",
+    "    \"./sample_images/c.png\", \"./sample_images/d.png\"],\n",
+    "    title=title,\n",
+    "    description=description,\n",
+    "    article=article,\n",
+    "    allow_flagging=False,\n",
+    "    analytics_enabled=False,\n",
+    "    ).launch(debug=True, enable_queue=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.12 ('fifa')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.12"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "5a4cff4f724f20f3784f32e905011239b516be3fadafd59414871df18d0dad63"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

requirements.txt ADDED Viewed

	@@ -0,0 +1,6 @@

+scipy==1.4.1
+torch
+h5py==2.10.0
+numpy==1.18.1
+opencv-python-headless==4.2.0.32
+Pillow

sample_images/a.png ADDED Viewed

sample_images/b.png ADDED Viewed

sample_images/c.png ADDED Viewed

sample_images/d.png ADDED Viewed

unetplusplus.py ADDED Viewed

	@@ -0,0 +1,142 @@

+import torch
+from torch import nn
+import torch.nn.functional as F
+__all__ = ['UNet', 'NestedUNet']
+"""Taken from https://github.com/4uiiurz1/pytorch-nested-unet"""
+class VGGBlock(nn.Module):
+    def __init__(self, in_channels, middle_channels, out_channels):
+        super().__init__()
+        self.relu = nn.ReLU(inplace=True)
+        self.conv1 = nn.Conv2d(in_channels, middle_channels, 3, padding=1)
+        self.bn1 = nn.BatchNorm2d(middle_channels)
+        self.conv2 = nn.Conv2d(middle_channels, out_channels, 3, padding=1)
+        self.bn2 = nn.BatchNorm2d(out_channels)
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+        return out
+class UNet(nn.Module):
+    def __init__(self, num_classes, input_channels=3, **kwargs):
+        super().__init__()
+        nb_filter = [32, 64, 128, 256, 512]
+        self.pool = nn.MaxPool2d(2, 2)
+        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
+        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
+        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
+        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
+        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])
+        self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
+        self.conv2_2 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
+        self.conv1_3 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
+        self.conv0_4 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
+        self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
+    def forward(self, input):
+        x0_0 = self.conv0_0(input)
+        x1_0 = self.conv1_0(self.pool(x0_0))
+        x2_0 = self.conv2_0(self.pool(x1_0))
+        x3_0 = self.conv3_0(self.pool(x2_0))
+        x4_0 = self.conv4_0(self.pool(x3_0))
+        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
+        x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1))
+        x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1))
+        x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1))
+        output = self.final(x0_4)
+        return output
+class NestedUNet(nn.Module):
+    """
+    U-Net Plus plus architecture
+    Reference: https://arxiv.org/abs/1807.10165
+    """
+    def __init__(self, num_classes=1, input_channels=3, deep_supervision=False, **kwargs):
+        super().__init__()
+        nb_filter = [32, 64, 128, 256, 512]
+        self.deep_supervision = deep_supervision
+        self.pool = nn.MaxPool2d(2, 2)
+        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.conv0_0 = VGGBlock(input_channels, nb_filter[0], nb_filter[0])
+        self.conv1_0 = VGGBlock(nb_filter[0], nb_filter[1], nb_filter[1])
+        self.conv2_0 = VGGBlock(nb_filter[1], nb_filter[2], nb_filter[2])
+        self.conv3_0 = VGGBlock(nb_filter[2], nb_filter[3], nb_filter[3])
+        self.conv4_0 = VGGBlock(nb_filter[3], nb_filter[4], nb_filter[4])
+        self.conv0_1 = VGGBlock(nb_filter[0]+nb_filter[1], nb_filter[0], nb_filter[0])
+        self.conv1_1 = VGGBlock(nb_filter[1]+nb_filter[2], nb_filter[1], nb_filter[1])
+        self.conv2_1 = VGGBlock(nb_filter[2]+nb_filter[3], nb_filter[2], nb_filter[2])
+        self.conv3_1 = VGGBlock(nb_filter[3]+nb_filter[4], nb_filter[3], nb_filter[3])
+        self.conv0_2 = VGGBlock(nb_filter[0]*2+nb_filter[1], nb_filter[0], nb_filter[0])
+        self.conv1_2 = VGGBlock(nb_filter[1]*2+nb_filter[2], nb_filter[1], nb_filter[1])
+        self.conv2_2 = VGGBlock(nb_filter[2]*2+nb_filter[3], nb_filter[2], nb_filter[2])
+        self.conv0_3 = VGGBlock(nb_filter[0]*3+nb_filter[1], nb_filter[0], nb_filter[0])
+        self.conv1_3 = VGGBlock(nb_filter[1]*3+nb_filter[2], nb_filter[1], nb_filter[1])
+        self.conv0_4 = VGGBlock(nb_filter[0]*4+nb_filter[1], nb_filter[0], nb_filter[0])
+        if self.deep_supervision:
+            self.final1 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
+            self.final2 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
+            self.final3 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
+            self.final4 = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
+        else:
+            self.final = nn.Conv2d(nb_filter[0], num_classes, kernel_size=1)
+    def forward(self, input):
+        x0_0 = self.conv0_0(input)
+        x1_0 = self.conv1_0(self.pool(x0_0))
+        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))
+        x2_0 = self.conv2_0(self.pool(x1_0))
+        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))
+        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))
+        x3_0 = self.conv3_0(self.pool(x2_0))
+        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))
+        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))
+        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))
+        x4_0 = self.conv4_0(self.pool(x3_0))
+        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))
+        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))
+        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))
+        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))
+        if self.deep_supervision:
+            output1 = self.final1(x0_1)
+            output2 = self.final2(x0_2)
+            output3 = self.final3(x0_3)
+            output4 = self.final4(x0_4)
+            return [output1, output2, output3, output4]
+        else:
+            output = self.final(x0_4)
+            return output