initial commit

README.md

@@ -1,12 +0,0 @@
----
-title: Mmdetection Space
-emoji: 🦀
-colorFrom: pink
-colorTo: pink
-sdk: gradio
-sdk_version: 3.1.4
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python
+
+from __future__ import annotations
+
+import argparse
+import os
+import pathlib
+import subprocess
+
+if os.getenv("SYSTEM") == "spaces":
+    import mim
+
+    mim.uninstall("mmcv-full", confirm_yes=True)
+    mim.install("mmcv-full==1.6.1", is_yes=True)
+
+    subprocess.call("pip uninstall -y opencv-python".split())
+    subprocess.call("pip uninstall -y opencv-python-headless".split())
+    subprocess.call("pip install opencv-python-headless==4.5.5.64".split())
+
+import cv2
+import gradio as gr
+import numpy as np
+
+from model import AppModel
+
+## Edit and 
+DESCRIPTION = """# MMDetection
+This is an unofficial demo for [https://github.com/open-mmlab/mmdetection](https://github.com/open-mmlab/mmdetection).
+"""
+FOOTER = '<img id="visitor-badge" src="https://visitor-badge.glitch.me/badge?page_id=hf-technical-mmdetection" alt="visitor badge" />'
+
+DEFAULT_MODEL_TYPE = "detection"
+DEFAULT_MODEL_NAMES = {
+    "detection": "faster_rcnn"
+}
+DEFAULT_MODEL_NAME = DEFAULT_MODEL_NAMES[DEFAULT_MODEL_TYPE]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", type=str, default="cpu")
+    parser.add_argument("--theme", type=str)
+    parser.add_argument("--share", action="store_true")
+    parser.add_argument("--port", type=int)
+    parser.add_argument("--disable-queue", dest="enable_queue", action="store_false")
+    return parser.parse_args()
+
+
+
+def update_input_image(image: np.ndarray) -> dict:
+    if image is None:
+        return gr.Image.update(value=None)
+    scale = 1500 / max(image.shape[:2])
+    if scale < 1:
+        image = cv2.resize(image, None, fx=scale, fy=scale)
+    return gr.Image.update(value=image)
+
+
+def update_model_name(model_type: str) -> dict:
+    model_dict = getattr(AppModel, f"{model_type.upper()}_MODEL_DICT")
+    model_names = list(model_dict.keys())
+    model_name = DEFAULT_MODEL_NAMES[model_type]
+    return gr.Dropdown.update(choices=model_names, value=model_name)
+
+
+def update_visualization_score_threshold(model_type: str) -> dict:
+    return gr.Slider.update(visible=model_type != "panoptic_segmentation")
+
+
+def update_redraw_button(model_type: str) -> dict:
+    return gr.Button.update(visible=model_type != "panoptic_segmentation")
+
+
+def set_example_image(example: list) -> dict:
+    return gr.Image.update(value=example[0])
+
+
+def main():
+    args = parse_args()
+    model = AppModel(DEFAULT_MODEL_NAME, args.device)
+
+    with gr.Blocks(theme=args.theme, css="style.css") as demo:
+        gr.Markdown(DESCRIPTION)
+
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    input_image = gr.Image(label="Input Image", type="numpy")
+                with gr.Group():
+                    with gr.Row():
+                        model_type = gr.Radio(
+                            list(DEFAULT_MODEL_NAMES.keys()),
+                            value=DEFAULT_MODEL_TYPE,
+                            label="Model Type",
+                        )
+                    with gr.Row():
+                        model_name = gr.Dropdown(
+                            model.model_list(),
+                            value=DEFAULT_MODEL_NAME,
+                            label="Model",
+                        )
+                with gr.Row():
+                    run_button = gr.Button(value="Run")
+                    prediction_results = gr.Variable()
+            with gr.Column():
+                with gr.Row():
+                    visualization = gr.Image(label="Result", type="numpy")
+                with gr.Row():
+                    visualization_score_threshold = gr.Slider(
+                        0,
+                        1,
+                        step=0.05,
+                        value=0.3,
+                        label="Visualization Score Threshold",
+                    )
+                with gr.Row():
+                    redraw_button = gr.Button(value="Redraw")
+
+        with gr.Row():
+            paths = sorted(pathlib.Path("images").rglob("*.jpeg"))
+            example_images = gr.Dataset(
+                components=[input_image], samples=[[path.as_posix()] for path in paths]
+            )
+
+        gr.Markdown(FOOTER)
+
+        input_image.change(
+            fn=update_input_image, inputs=input_image, outputs=input_image
+        )
+
+        model_type.change(fn=update_model_name, inputs=model_type, outputs=model_name)
+        model_type.change(
+            fn=update_visualization_score_threshold,
+            inputs=model_type,
+            outputs=visualization_score_threshold,
+        )
+        model_type.change(
+            fn=update_redraw_button, inputs=model_type, outputs=redraw_button
+        )
+
+        model_name.change(fn=model.set_model, inputs=model_name, outputs=None)
+        run_button.click(
+            fn=model.run,
+            inputs=[
+                model_name,
+                input_image,
+                visualization_score_threshold,
+            ],
+            outputs=[
+                prediction_results,
+                visualization,
+            ],
+        )
+        redraw_button.click(
+            fn=model.visualize_detection_results,
+            inputs=[
+                input_image,
+                prediction_results,
+                visualization_score_threshold,
+            ],
+            outputs=visualization,
+        )
+        example_images.click(
+            fn=set_example_image, inputs=example_images, outputs=input_image
+        )
+
+    demo.launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
+    )
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from mmdet.apis import inference_detector, init_detector_from_hf_hub
+
+MODEL_DICT = {"faster_rcnn": {"repo_id": "blesot/Faster-R-CNN-Object-detection"}, "mask_rcnn": {"repo_id": "blesot/Mask-RCNN"}}
+
+
+class Model:
+
+    def __init__(self, model_name: str, device: str | torch.device):
+        self.device = torch.device(device)
+        self._load_all_models_once()
+        self.model_name = model_name
+        self.model = self._load_model(model_name)
+
+    def _load_all_models_once(self) -> None:
+        for name in MODEL_DICT.keys():
+            self._load_model(name)
+
+    def _load_model(self, name: str) -> nn.Module:
+        dic = MODEL_DICT[name]
+        return init_detector_from_hf_hub(dic['repo_id'], device=self.device)
+
+    def set_model(self, name: str) -> None:
+        if name == self.model_name:
+            return
+        self.model_name = name
+        self.model = self._load_model(name)
+
+    def detect_and_visualize(
+        self, image: np.ndarray, score_threshold: float
+    ) -> tuple[list[np.ndarray] | tuple[list[np.ndarray],
+                                        list[list[np.ndarray]]]
+               | dict[str, np.ndarray], np.ndarray]:
+        out = self.detect(image)
+        vis = self.visualize_detection_results(image, out, score_threshold)
+        return out, vis
+
+    def detect(
+        self, image: np.ndarray
+    ) -> list[np.ndarray] | tuple[
+            list[np.ndarray], list[list[np.ndarray]]] | dict[str, np.ndarray]:
+        image = image[:, :, ::-1]  # RGB -> BGR
+        out = inference_detector(self.model, image)
+        return out
+
+    def visualize_detection_results(
+            self,
+            image: np.ndarray,
+            detection_results: list[np.ndarray]
+        | tuple[list[np.ndarray], list[list[np.ndarray]]]
+        | dict[str, np.ndarray],
+            score_threshold: float = 0.3) -> np.ndarray:
+        image = image[:, :, ::-1]  # RGB -> BGR
+        vis = self.model.show_result(image,
+                                     detection_results,
+                                     score_thr=score_threshold,
+                                     bbox_color=None,
+                                     text_color=(200, 200, 200),
+                                     mask_color=None)
+        return vis[:, :, ::-1]  # BGR -> RGB
+
+
+class AppModel(Model):
+    def run(
+        self, model_name: str, image: np.ndarray, score_threshold: float
+    ) -> tuple[list[np.ndarray] | tuple[list[np.ndarray],
+                                        list[list[np.ndarray]]]
+               | dict[str, np.ndarray], np.ndarray]:
+        self.set_model(model_name)
+        return self.detect_and_visualize(image, score_threshold)
+    
+    def model_list(self) -> list[str]:
+        return list(MODEL_DICT.keys())

requirements.txt

@@ -0,0 +1,6 @@
+mmcv-full==1.6.1
+package/mmdet_huggingface-2.25.1.tar.gz
+opencv-python-headless==4.5.5.64
+openmim==0.1.5
+torch==1.11.0
+torchvision==0.12.0

style.css

@@ -0,0 +1,15 @@
+h1 {
+    text-align: center;
+}
+
+img#overview {
+    display: block;
+    margin: auto;
+    max-width: 1000px;
+    max-height: 600px;
+}
+
+img#visitor-badge {
+    display: block;
+    margin: auto;
+}