Refactor

app.py

@@ -3,33 +3,11 @@
 from __future__ import annotations
 
 import argparse
-import os
 import pathlib
-import subprocess
-import sys
-
-if os.getenv('SYSTEM') == 'spaces':
-    import mim
-
-    mim.uninstall('mmcv-full', confirm_yes=True)
-    mim.install('mmcv-full==1.5.0', is_yes=True)
-
-    subprocess.run('pip uninstall -y opencv-python'.split())
-    subprocess.run('pip uninstall -y opencv-python-headless'.split())
-    subprocess.run('pip install opencv-python-headless==4.5.5.64'.split())
-
-    with open('patch') as f:
-        subprocess.run('patch -p1'.split(), cwd='CBNetV2', stdin=f)
-    subprocess.run('mv palette.py CBNetV2/mmdet/core/visualization/'.split())
 
 import gradio as gr
-import numpy as np
-import torch
-import torch.nn as nn
 
-sys.path.insert(0, 'CBNetV2/')
-
-from mmdet.apis import inference_detector, init_detector
+from model import Model
 
 DESCRIPTION = '''# CBNetV2
 
@@ -49,117 +27,6 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
-class Model:
-    def __init__(self, device: str | torch.device):
-        self.device = torch.device(device)
-        self.models = self._load_models()
-        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',
-            },
-            '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':
-            #                'CBNetV2/configs/cbnet/cascade_mask_rcnn_cbv2_swin_small_patch4_window7_mstrain_400-1400_adamw_3x_coco.py',
-            #            '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-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',
-            },
-        }
-
-        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]
-
-            out_path = out_dir / zip_name
-            if out_path.exists():
-                return
-            torch.hub.download_url_to_file(model_url, out_path)
-
-            with zipfile.ZipFile(out_path) as f:
-                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]
-            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)
-            for key, dic in model_dict.items()
-        }
-        return models
-
-    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]:
-        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]:
-        image = image[:, :, ::-1]  # RGB -> BGR
-        model = self.models[self.model_name]
-        out = inference_detector(model, image)
-        return out
-
-    def visualize_detection_results(
-            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)
-        return vis[:, :, ::-1]  # BGR -> RGB
-
-
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
 

model.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import os
+import pathlib
+import subprocess
+import sys
+
+if os.getenv('SYSTEM') == 'spaces':
+    import mim
+
+    mim.uninstall('mmcv-full', confirm_yes=True)
+    mim.install('mmcv-full==1.5.0', is_yes=True)
+
+    subprocess.run('pip uninstall -y opencv-python'.split())
+    subprocess.run('pip uninstall -y opencv-python-headless'.split())
+    subprocess.run('pip install opencv-python-headless==4.5.5.64'.split())
+
+    with open('patch') as f:
+        subprocess.run('patch -p1'.split(), 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
+
+sys.path.insert(0, 'CBNetV2/')
+
+from mmdet.apis import inference_detector, init_detector
+
+
+class Model:
+    def __init__(self, device: str | torch.device):
+        self.device = torch.device(device)
+        self.models = self._load_models()
+        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',
+            },
+            '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':
+            #                'CBNetV2/configs/cbnet/cascade_mask_rcnn_cbv2_swin_small_patch4_window7_mstrain_400-1400_adamw_3x_coco.py',
+            #            '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-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',
+            },
+        }
+
+        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]
+
+            out_path = out_dir / zip_name
+            if out_path.exists():
+                return
+            torch.hub.download_url_to_file(model_url, out_path)
+
+            with zipfile.ZipFile(out_path) as f:
+                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]
+            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)
+            for key, dic in model_dict.items()
+        }
+        return models
+
+    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]:
+        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]:
+        image = image[:, :, ::-1]  # RGB -> BGR
+        model = self.models[self.model_name]
+        out = inference_detector(model, image)
+        return out
+
+    def visualize_detection_results(
+            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)
+        return vis[:, :, ::-1]  # BGR -> RGB