Spaces:
Running
Running
File size: 4,931 Bytes
7fd17d1 d25c07c 293637a d25c07c 7fd17d1 9f992a7 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 293637a 7fd17d1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
from __future__ import annotations
import pathlib
import sys
import numpy as np
import torch
import torch.nn as nn
app_dir = pathlib.Path(__file__).parent
submodule_dir = app_dir / "CBNetV2/"
sys.path.insert(0, submodule_dir.as_posix())
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.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
|