Upload with huggingface_hub

DESCRIPTION.md

@@ -1 +1 @@
-Image segmentation using DETR. Takes in both an inputu image and the desired confidence, and returns a segmented image.
+Simple image segmentation using gradio's AnnotatedImage component.

README.md

@@ -5,7 +5,7 @@ emoji: 🔥
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio
-sdk_version: 3.26.0
+sdk_version: 3.27.0
 app_file: run.py
 pinned: false
 ---

run.ipynb

@@ -1 +1 @@
-{"cells": [{"cell_type": "markdown", "id": 302934307671667531413257853548643485645, "metadata": {}, "source": ["# Gradio Demo: image_segmentation\n", "### Image segmentation using DETR. Takes in both an inputu image and the desired confidence, and returns a segmented image.\n", "        "]}, {"cell_type": "code", "execution_count": null, "id": 272996653310673477252411125948039410165, "metadata": {}, "outputs": [], "source": ["!pip install -q gradio transformers torch scipy numpy"]}, {"cell_type": "code", "execution_count": null, "id": 288918539441861185822528903084949547379, "metadata": {}, "outputs": [], "source": ["# Downloading files from the demo repo\n", "import os\n", "!wget -q https://github.com/gradio-app/gradio/raw/main/demo/image_segmentation/example_2.png"]}, {"cell_type": "code", "execution_count": null, "id": 44380577570523278879349135829904343037, "metadata": {}, "outputs": [], "source": ["import gradio as gr\n", "import torch\n", "import random\n", "import numpy as np\n", "from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation\n", "\n", "device = torch.device(\"cpu\")\n", "model = MaskFormerForInstanceSegmentation.from_pretrained(\"facebook/maskformer-swin-tiny-ade\").to(device)\n", "model.eval()\n", "preprocessor = MaskFormerFeatureExtractor.from_pretrained(\"facebook/maskformer-swin-tiny-ade\")\n", "\n", "def visualize_instance_seg_mask(mask):\n", "    image = np.zeros((mask.shape[0], mask.shape[1], 3))\n", "    labels = np.unique(mask)\n", "    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}\n", "    for i in range(image.shape[0]):\n", "      for j in range(image.shape[1]):\n", "        image[i, j, :] = label2color[mask[i, j]]\n", "    image = image / 255\n", "    return image\n", "\n", "def query_image(img):\n", "    target_size = (img.shape[0], img.shape[1])\n", "    inputs = preprocessor(images=img, return_tensors=\"pt\")\n", "    with torch.no_grad():\n", "        outputs = model(**inputs)\n", "    outputs.class_queries_logits = outputs.class_queries_logits.cpu()\n", "    outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()\n", "    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()\n", "    results = torch.argmax(results, dim=0).numpy()\n", "    results = visualize_instance_seg_mask(results)\n", "    return results\n", "\n", "demo = gr.Interface(\n", "    query_image, \n", "    inputs=[gr.Image()], \n", "    outputs=\"image\",\n", "    title=\"MaskFormer Demo\",\n", "    examples=[[\"example_2.png\"]]\n", ")\n", "\n", "demo.launch()"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}
+{"cells": [{"cell_type": "markdown", "id": 302934307671667531413257853548643485645, "metadata": {}, "source": ["# Gradio Demo: image_segmentation\n", "### Simple image segmentation using gradio's AnnotatedImage component.\n", "        "]}, {"cell_type": "code", "execution_count": null, "id": 272996653310673477252411125948039410165, "metadata": {}, "outputs": [], "source": ["!pip install -q gradio "]}, {"cell_type": "code", "execution_count": null, "id": 288918539441861185822528903084949547379, "metadata": {}, "outputs": [], "source": ["import gradio as gr\n", "import numpy as np\n", "import random\n", "\n", "with gr.Blocks() as demo:\n", "    section_labels = [\n", "        \"apple\",\n", "        \"banana\",\n", "        \"carrot\",\n", "        \"donut\",\n", "        \"eggplant\",\n", "        \"fish\",\n", "        \"grapes\",\n", "        \"hamburger\",\n", "        \"ice cream\",\n", "        \"juice\",\n", "    ]\n", "\n", "    with gr.Row():\n", "        num_boxes = gr.Slider(0, 5, 2, step=1, label=\"Number of boxes\")\n", "        num_segments = gr.Slider(0, 5, 1, step=1, label=\"Number of segments\")\n", "\n", "    with gr.Row():\n", "        img_input = gr.Image()\n", "        img_output = gr.AnnotatedImage().style(\n", "            color_map={\"banana\": \"#a89a00\", \"carrot\": \"#ffae00\"}\n", "        )\n", "\n", "    section_btn = gr.Button(\"Identify Sections\")\n", "    selected_section = gr.Textbox(label=\"Selected Section\")\n", "\n", "    def section(img, num_boxes, num_segments):\n", "        sections = []\n", "        for a in range(num_boxes):\n", "            x = random.randint(0, img.shape[1])\n", "            y = random.randint(0, img.shape[0])\n", "            w = random.randint(0, img.shape[1] - x)\n", "            h = random.randint(0, img.shape[0] - y)\n", "            sections.append(((x, y, x + w, y + h), section_labels[a]))\n", "        for b in range(num_segments):\n", "            x = random.randint(0, img.shape[1])\n", "            y = random.randint(0, img.shape[0])\n", "            r = random.randint(0, min(x, y, img.shape[1] - x, img.shape[0] - y))\n", "            mask = np.zeros(img.shape[:2])\n", "            for i in range(img.shape[0]):\n", "                for j in range(img.shape[1]):\n", "                    dist_square = (i - y) ** 2 + (j - x) ** 2\n", "                    if dist_square < r**2:\n", "                        mask[i, j] = round((r**2 - dist_square) / r**2 * 4) / 4\n", "            sections.append((mask, section_labels[b + num_boxes]))\n", "        return (img, sections)\n", "\n", "    section_btn.click(section, [img_input, num_boxes, num_segments], img_output)\n", "\n", "    def select_section(evt: gr.SelectData):\n", "        return section_labels[evt.index]\n", "\n", "    img_output.select(select_section, None, selected_section)\n", "\n", "\n", "demo.launch()\n"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}

run.py

@@ -1,42 +1,61 @@
 import gradio as gr
-import torch
-import random
 import numpy as np
-from transformers import MaskFormerFeatureExtractor, MaskFormerForInstanceSegmentation
-
-device = torch.device("cpu")
-model = MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-tiny-ade").to(device)
-model.eval()
-preprocessor = MaskFormerFeatureExtractor.from_pretrained("facebook/maskformer-swin-tiny-ade")
-
-def visualize_instance_seg_mask(mask):
-    image = np.zeros((mask.shape[0], mask.shape[1], 3))
-    labels = np.unique(mask)
-    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
-    for i in range(image.shape[0]):
-      for j in range(image.shape[1]):
-        image[i, j, :] = label2color[mask[i, j]]
-    image = image / 255
-    return image
-
-def query_image(img):
-    target_size = (img.shape[0], img.shape[1])
-    inputs = preprocessor(images=img, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
-    outputs.class_queries_logits = outputs.class_queries_logits.cpu()
-    outputs.masks_queries_logits = outputs.masks_queries_logits.cpu()
-    results = preprocessor.post_process_segmentation(outputs=outputs, target_size=target_size)[0].cpu().detach()
-    results = torch.argmax(results, dim=0).numpy()
-    results = visualize_instance_seg_mask(results)
-    return results
-
-demo = gr.Interface(
-    query_image, 
-    inputs=[gr.Image()], 
-    outputs="image",
-    title="MaskFormer Demo",
-    examples=[["example_2.png"]]
-)
-
-demo.launch()
+import random
+
+with gr.Blocks() as demo:
+    section_labels = [
+        "apple",
+        "banana",
+        "carrot",
+        "donut",
+        "eggplant",
+        "fish",
+        "grapes",
+        "hamburger",
+        "ice cream",
+        "juice",
+    ]
+
+    with gr.Row():
+        num_boxes = gr.Slider(0, 5, 2, step=1, label="Number of boxes")
+        num_segments = gr.Slider(0, 5, 1, step=1, label="Number of segments")
+
+    with gr.Row():
+        img_input = gr.Image()
+        img_output = gr.AnnotatedImage().style(
+            color_map={"banana": "#a89a00", "carrot": "#ffae00"}
+        )
+
+    section_btn = gr.Button("Identify Sections")
+    selected_section = gr.Textbox(label="Selected Section")
+
+    def section(img, num_boxes, num_segments):
+        sections = []
+        for a in range(num_boxes):
+            x = random.randint(0, img.shape[1])
+            y = random.randint(0, img.shape[0])
+            w = random.randint(0, img.shape[1] - x)
+            h = random.randint(0, img.shape[0] - y)
+            sections.append(((x, y, x + w, y + h), section_labels[a]))
+        for b in range(num_segments):
+            x = random.randint(0, img.shape[1])
+            y = random.randint(0, img.shape[0])
+            r = random.randint(0, min(x, y, img.shape[1] - x, img.shape[0] - y))
+            mask = np.zeros(img.shape[:2])
+            for i in range(img.shape[0]):
+                for j in range(img.shape[1]):
+                    dist_square = (i - y) ** 2 + (j - x) ** 2
+                    if dist_square < r**2:
+                        mask[i, j] = round((r**2 - dist_square) / r**2 * 4) / 4
+            sections.append((mask, section_labels[b + num_boxes]))
+        return (img, sections)
+
+    section_btn.click(section, [img_input, num_boxes, num_segments], img_output)
+
+    def select_section(evt: gr.SelectData):
+        return section_labels[evt.index]
+
+    img_output.select(select_section, None, selected_section)
+
+
+demo.launch()