Update

.pre-commit-config.yaml

@@ -1,37 +1,61 @@
 exclude: ^patch
 repos:
-- repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v4.2.0
-  hooks:
-  - id: check-executables-have-shebangs
-  - id: check-json
-  - id: check-merge-conflict
-  - id: check-shebang-scripts-are-executable
-  - id: check-toml
-  - id: check-yaml
-  - id: double-quote-string-fixer
-  - id: end-of-file-fixer
-  - id: mixed-line-ending
-    args: ['--fix=lf']
-  - id: requirements-txt-fixer
-  - id: trailing-whitespace
-- repo: https://github.com/myint/docformatter
-  rev: v1.4
-  hooks:
-  - id: docformatter
-    args: ['--in-place']
-- repo: https://github.com/pycqa/isort
-  rev: 5.12.0
-  hooks:
-    - id: isort
-- repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.991
-  hooks:
-    - id: mypy
-      args: ['--ignore-missing-imports']
-      additional_dependencies: ['types-python-slugify']
-- repo: https://github.com/google/yapf
-  rev: v0.32.0
-  hooks:
-  - id: yapf
-    args: ['--parallel', '--in-place']
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0
+    hooks:
+      - id: check-executables-have-shebangs
+      - id: check-json
+      - id: check-merge-conflict
+      - id: check-shebang-scripts-are-executable
+      - id: check-toml
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+        args: ["--fix=lf"]
+      - id: requirements-txt-fixer
+      - id: trailing-whitespace
+  - repo: https://github.com/myint/docformatter
+    rev: v1.7.5
+    hooks:
+      - id: docformatter
+        args: ["--in-place"]
+  - repo: https://github.com/pycqa/isort
+    rev: 5.13.2
+    hooks:
+      - id: isort
+        args: ["--profile", "black"]
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.10.0
+    hooks:
+      - id: mypy
+        args: ["--ignore-missing-imports"]
+        additional_dependencies:
+          [
+            "types-python-slugify",
+            "types-requests",
+            "types-PyYAML",
+            "types-pytz",
+          ]
+  - repo: https://github.com/psf/black
+    rev: 24.4.2
+    hooks:
+      - id: black
+        language_version: python3.10
+        args: ["--line-length", "119"]
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.7.1
+    hooks:
+      - id: nbstripout
+        args:
+          [
+            "--extra-keys",
+            "metadata.interpreter metadata.kernelspec cell.metadata.pycharm",
+          ]
+  - repo: https://github.com/nbQA-dev/nbQA
+    rev: 1.8.5
+    hooks:
+      - id: nbqa-black
+      - id: nbqa-pyupgrade
+        args: ["--py37-plus"]
+      - id: nbqa-isort
+        args: ["--float-to-top"]

.vscode/settings.json

@@ -0,0 +1,30 @@
+{
+    "editor.formatOnSave": true,
+    "files.insertFinalNewline": false,
+    "[python]": {
+        "editor.defaultFormatter": "ms-python.black-formatter",
+        "editor.formatOnType": true,
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": "explicit"
+        }
+    },
+    "[jupyter]": {
+        "files.insertFinalNewline": false
+    },
+    "black-formatter.args": [
+        "--line-length=119"
+    ],
+    "isort.args": ["--profile", "black"],
+    "flake8.args": [
+        "--max-line-length=119"
+    ],
+    "ruff.lint.args": [
+        "--line-length=119"
+    ],
+    "notebook.output.scrolling": true,
+    "notebook.formatOnCellExecution": true,
+    "notebook.formatOnSave.enabled": true,
+        "notebook.codeActionsOnSave": {
+            "source.organizeImports": "explicit"
+        }
+}

app.py

@@ -8,60 +8,57 @@ import gradio as gr
 
 from model import Model
 
-DESCRIPTION = '# [CBNetV2](https://github.com/VDIGPKU/CBNetV2)'
+DESCRIPTION = "# [CBNetV2](https://github.com/VDIGPKU/CBNetV2)"
 
 model = Model()
 
-with gr.Blocks(css='style.css') as demo:
+with gr.Blocks(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')
+                input_image = gr.Image(label="Input Image", type="numpy")
             with gr.Row():
-                detector_name = gr.Dropdown(label='Detector',
-                                            choices=list(model.models.keys()),
-                                            value=model.model_name)
+                detector_name = gr.Dropdown(
+                    label="Detector", choices=list(model.models.keys()), value=model.model_name
+                )
             with gr.Row():
-                detect_button = gr.Button('Detect')
+                detect_button = gr.Button("Detect")
                 detection_results = gr.Variable()
         with gr.Column():
             with gr.Row():
-                detection_visualization = gr.Image(label='Detection Result',
-                                                   type='numpy')
+                detection_visualization = gr.Image(label="Detection Result", type="numpy")
             with gr.Row():
                 visualization_score_threshold = gr.Slider(
-                    label='Visualization Score Threshold',
-                    minimum=0,
-                    maximum=1,
-                    step=0.05,
-                    value=0.3)
+                    label="Visualization Score Threshold", minimum=0, maximum=1, step=0.05, value=0.3
+                )
             with gr.Row():
-                redraw_button = gr.Button('Redraw')
+                redraw_button = gr.Button("Redraw")
 
     with gr.Row():
-        paths = sorted(pathlib.Path('images').rglob('*.jpg'))
-        gr.Examples(examples=[[path.as_posix()] for path in paths],
-                    inputs=input_image)
+        paths = sorted(pathlib.Path("images").rglob("*.jpg"))
+        gr.Examples(examples=[[path.as_posix()] for path in paths], inputs=input_image)
 
-    detector_name.change(fn=model.set_model_name,
-                         inputs=[detector_name],
-                         outputs=None)
-    detect_button.click(fn=model.detect_and_visualize,
-                        inputs=[
-                            input_image,
-                            visualization_score_threshold,
-                        ],
-                        outputs=[
-                            detection_results,
-                            detection_visualization,
-                        ])
-    redraw_button.click(fn=model.visualize_detection_results,
-                        inputs=[
-                            input_image,
-                            detection_results,
-                            visualization_score_threshold,
-                        ],
-                        outputs=[detection_visualization])
+    detector_name.change(fn=model.set_model_name, inputs=[detector_name], outputs=None)
+    detect_button.click(
+        fn=model.detect_and_visualize,
+        inputs=[
+            input_image,
+            visualization_score_threshold,
+        ],
+        outputs=[
+            detection_results,
+            detection_visualization,
+        ],
+    )
+    redraw_button.click(
+        fn=model.visualize_detection_results,
+        inputs=[
+            input_image,
+            detection_results,
+            visualization_score_threshold,
+        ],
+        outputs=[detection_visualization],
+    )
 demo.queue(max_size=10).launch()

model.py

@@ -6,26 +6,26 @@ import shlex
 import subprocess
 import sys
 
-if os.getenv('SYSTEM') == 'spaces':
+if os.getenv("SYSTEM") == "spaces":
     import mim
 
-    mim.uninstall('mmcv-full', confirm_yes=True)
-    mim.install('mmcv-full==1.5.0', is_yes=True)
+    mim.uninstall("mmcv-full", confirm_yes=True)
+    mim.install("mmcv-full==1.5.0", is_yes=True)
 
-    subprocess.run(shlex.split('pip uninstall -y opencv-python'))
-    subprocess.run(shlex.split('pip uninstall -y opencv-python-headless'))
-    subprocess.run(shlex.split('pip install opencv-python-headless==4.8.0.74'))
+    subprocess.run(shlex.split("pip uninstall -y opencv-python"))
+    subprocess.run(shlex.split("pip uninstall -y opencv-python-headless"))
+    subprocess.run(shlex.split("pip install opencv-python-headless==4.8.0.74"))
 
-    with open('patch') as f:
-        subprocess.run(shlex.split('patch -p1'), cwd='CBNetV2', stdin=f)
-    subprocess.run('mv palette.py CBNetV2/mmdet/core/visualization/'.split())
+    with open("patch") as f:
+        subprocess.run(shlex.split("patch -p1"), cwd="CBNetV2", stdin=f)
+    subprocess.run("mv palette.py CBNetV2/mmdet/core/visualization/".split())
 
 import numpy as np
 import torch
 import torch.nn as nn
 
 app_dir = pathlib.Path(__file__).parent
-submodule_dir = app_dir / 'CBNetV2/'
+submodule_dir = app_dir / "CBNetV2/"
 sys.path.insert(0, submodule_dir.as_posix())
 
 from mmdet.apis import inference_detector, init_detector
@@ -33,24 +33,19 @@ from mmdet.apis import inference_detector, init_detector
 
 class Model:
     def __init__(self):
-        self.device = torch.device(
-            'cuda:0' if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         self.models = self._load_models()
-        self.model_name = 'Improved HTC (DB-Swin-B)'
+        self.model_name = "Improved HTC (DB-Swin-B)"
 
     def _load_models(self) -> dict[str, nn.Module]:
         model_dict = {
-            'Faster R-CNN (DB-ResNet50)': {
-                'config':
-                'CBNetV2/configs/cbnet/faster_rcnn_cbv2d1_r50_fpn_1x_coco.py',
-                'model':
-                'https://github.com/CBNetwork/storage/releases/download/v1.0.0/faster_rcnn_cbv2d1_r50_fpn_1x_coco.pth.zip',
+            "Faster R-CNN (DB-ResNet50)": {
+                "config": "CBNetV2/configs/cbnet/faster_rcnn_cbv2d1_r50_fpn_1x_coco.py",
+                "model": "https://github.com/CBNetwork/storage/releases/download/v1.0.0/faster_rcnn_cbv2d1_r50_fpn_1x_coco.pth.zip",
             },
-            'Mask R-CNN (DB-Swin-T)': {
-                'config':
-                'CBNetV2/configs/cbnet/mask_rcnn_cbv2_swin_tiny_patch4_window7_mstrain_480-800_adamw_3x_coco.py',
-                'model':
-                'https://github.com/CBNetwork/storage/releases/download/v1.0.0/mask_rcnn_cbv2_swin_tiny_patch4_window7_mstrain_480-800_adamw_3x_coco.pth.zip',
+            "Mask R-CNN (DB-Swin-T)": {
+                "config": "CBNetV2/configs/cbnet/mask_rcnn_cbv2_swin_tiny_patch4_window7_mstrain_480-800_adamw_3x_coco.py",
+                "model": "https://github.com/CBNetwork/storage/releases/download/v1.0.0/mask_rcnn_cbv2_swin_tiny_patch4_window7_mstrain_480-800_adamw_3x_coco.pth.zip",
             },
             #        'Cascade Mask R-CNN (DB-Swin-S)': {
             #            'config':
@@ -58,34 +53,28 @@ class Model:
             #            'model':
             #                'https://github.com/CBNetwork/storage/releases/download/v1.0.0/cascade_mask_rcnn_cbv2_swin_small_patch4_window7_mstrain_400-1400_adamw_3x_coco.pth.zip',
             #        },
-            'Improved HTC (DB-Swin-B)': {
-                'config':
-                'CBNetV2/configs/cbnet/htc_cbv2_swin_base_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_20e_coco.py',
-                'model':
-                'https://github.com/CBNetwork/storage/releases/download/v1.0.0/htc_cbv2_swin_base22k_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_20e_coco.pth.zip',
+            "Improved HTC (DB-Swin-B)": {
+                "config": "CBNetV2/configs/cbnet/htc_cbv2_swin_base_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_20e_coco.py",
+                "model": "https://github.com/CBNetwork/storage/releases/download/v1.0.0/htc_cbv2_swin_base22k_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_20e_coco.pth.zip",
             },
-            'Improved HTC (DB-Swin-L)': {
-                'config':
-                'CBNetV2/configs/cbnet/htc_cbv2_swin_large_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.py',
-                'model':
-                'https://github.com/CBNetwork/storage/releases/download/v1.0.0/htc_cbv2_swin_large22k_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.pth.zip',
+            "Improved HTC (DB-Swin-L)": {
+                "config": "CBNetV2/configs/cbnet/htc_cbv2_swin_large_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.py",
+                "model": "https://github.com/CBNetwork/storage/releases/download/v1.0.0/htc_cbv2_swin_large22k_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.pth.zip",
             },
-            'Improved HTC (DB-Swin-L (TTA))': {
-                'config':
-                'CBNetV2/configs/cbnet/htc_cbv2_swin_large_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.py',
-                'model':
-                'https://github.com/CBNetwork/storage/releases/download/v1.0.0/htc_cbv2_swin_large22k_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.pth.zip',
+            "Improved HTC (DB-Swin-L (TTA))": {
+                "config": "CBNetV2/configs/cbnet/htc_cbv2_swin_large_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.py",
+                "model": "https://github.com/CBNetwork/storage/releases/download/v1.0.0/htc_cbv2_swin_large22k_patch4_window7_mstrain_400-1400_giou_4conv1f_adamw_1x_coco.pth.zip",
             },
         }
 
-        weight_dir = pathlib.Path('weights')
+        weight_dir = pathlib.Path("weights")
         weight_dir.mkdir(exist_ok=True)
 
         def _download(model_name: str, out_dir: pathlib.Path) -> None:
             import zipfile
 
-            model_url = model_dict[model_name]['model']
-            zip_name = model_url.split('/')[-1]
+            model_url = model_dict[model_name]["model"]
+            zip_name = model_url.split("/")[-1]
 
             out_path = out_dir / zip_name
             if out_path.exists():
@@ -96,17 +85,15 @@ class Model:
                 f.extractall(out_dir)
 
         def _get_model_path(model_name: str) -> str:
-            model_url = model_dict[model_name]['model']
-            model_name = model_url.split('/')[-1][:-4]
+            model_url = model_dict[model_name]["model"]
+            model_name = model_url.split("/")[-1][:-4]
             return (weight_dir / model_name).as_posix()
 
         for model_name in model_dict:
             _download(model_name, weight_dir)
 
         models = {
-            key: init_detector(dic['config'],
-                               _get_model_path(key),
-                               device=self.device)
+            key: init_detector(dic["config"], _get_model_path(key), device=self.device)
             for key, dic in model_dict.items()
         }
         return models
@@ -114,9 +101,7 @@ class Model:
     def set_model_name(self, name: str) -> None:
         self.model_name = name
 
-    def detect_and_visualize(
-            self, image: np.ndarray,
-            score_threshold: float) -> tuple[list[np.ndarray], np.ndarray]:
+    def detect_and_visualize(self, image: np.ndarray, score_threshold: float) -> tuple[list[np.ndarray], np.ndarray]:
         out = self.detect(image)
         vis = self.visualize_detection_results(image, out, score_threshold)
         return out, vis
@@ -128,16 +113,16 @@ class Model:
         return out
 
     def visualize_detection_results(
-            self,
-            image: np.ndarray,
-            detection_results: list[np.ndarray],
-            score_threshold: float = 0.3) -> np.ndarray:
+        self, image: np.ndarray, detection_results: list[np.ndarray], score_threshold: float = 0.3
+    ) -> np.ndarray:
         image = image[:, :, ::-1]  # RGB -> BGR
         model = self.models[self.model_name]
-        vis = model.show_result(image,
-                                detection_results,
-                                score_thr=score_threshold,
-                                bbox_color=None,
-                                text_color=(200, 200, 200),
-                                mask_color=None)
+        vis = 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

palette.py

@@ -208,6 +208,7 @@ Copyright 2018-2023 OpenMMLab. All rights reserved.
    limitations under the License.
 ```
 """
+
 # Copyright (c) OpenMMLab. All rights reserved.
 import mmcv
 import numpy as np
@@ -245,29 +246,31 @@ def get_palette(palette, num_classes):
         dataset_palette = palette
     elif isinstance(palette, tuple):
         dataset_palette = [palette] * num_classes
-    elif palette == 'random' or palette is None:
+    elif palette == "random" or palette is None:
         state = np.random.get_state()
         # random color
         np.random.seed(42)
         palette = np.random.randint(0, 256, size=(num_classes, 3))
         np.random.set_state(state)
         dataset_palette = [tuple(c) for c in palette]
-    elif palette == 'coco':
+    elif palette == "coco":
         from mmdet.datasets import CocoDataset, CocoPanopticDataset
+
         dataset_palette = CocoDataset.PALETTE
         if len(dataset_palette) < num_classes:
             dataset_palette = CocoPanopticDataset.PALETTE
-    elif palette == 'citys':
+    elif palette == "citys":
         from mmdet.datasets import CityscapesDataset
+
         dataset_palette = CityscapesDataset.PALETTE
-    elif palette == 'voc':
+    elif palette == "voc":
         from mmdet.datasets import VOCDataset
+
         dataset_palette = VOCDataset.PALETTE
     elif mmcv.is_str(palette):
         dataset_palette = [mmcv.color_val(palette)[::-1]] * num_classes
     else:
-        raise TypeError(f'Invalid type for palette: {type(palette)}')
+        raise TypeError(f"Invalid type for palette: {type(palette)}")
 
-    assert len(dataset_palette) >= num_classes, \
-        'The length of palette should not be less than `num_classes`.'
+    assert len(dataset_palette) >= num_classes, "The length of palette should not be less than `num_classes`."
     return dataset_palette

style.css

@@ -1,7 +1,11 @@
 h1 {
   text-align: center;
-}
-img#visitor-badge {
   display: block;
+}
+
+#duplicate-button {
   margin: auto;
+  color: #fff;
+  background: #1565c0;
+  border-radius: 100vh;
 }