AIKey/mmdetection1

1from __future__ import annotations2 3import os4 5import huggingface_hub6import numpy as np7import torch8import torch.nn as nn9import yaml10from mmdet.apis import inference_detector, init_detector11 12 13def _load_model_dict(path: str) -> dict[str, dict[str, str]]:14    with open(path) as f:15        dic = yaml.safe_load(f)16    _update_config_path(dic)17    _update_model_dict_if_hf_token_is_given(dic)18    return dic19 20 21def _update_config_path(model_dict: dict[str, dict[str, str]]) -> None:22    for dic in model_dict.values():23        dic['config'] = dic['config'].replace(24            'https://github.com/open-mmlab/mmdetection/tree/master',25            'mmdet_configs')26 27 28def _update_model_dict_if_hf_token_is_given(29        model_dict: dict[str, dict[str, str]]) -> None:30    token = os.getenv('HF_TOKEN')31    if token is None:32        return33 34    for dic in model_dict.values():35        ckpt_path = dic['model']36        name = ckpt_path.split('/')[-1]37        ckpt_path = huggingface_hub.hf_hub_download('hysts/mmdetection',38                                                    f'models/{name}',39                                                    use_auth_token=token)40        dic['model'] = ckpt_path41 42 43class Model:44    DETECTION_MODEL_DICT = _load_model_dict('model_dict/detection.yaml')45    INSTANCE_SEGMENTATION_MODEL_DICT = _load_model_dict(46        'model_dict/instance_segmentation.yaml')47    PANOPTIC_SEGMENTATION_MODEL_DICT = _load_model_dict(48        'model_dict/panoptic_segmentation.yaml')49    MODEL_DICT = DETECTION_MODEL_DICT | INSTANCE_SEGMENTATION_MODEL_DICT | PANOPTIC_SEGMENTATION_MODEL_DICT50 51    def __init__(self, model_name: str, device: str | torch.device):52        self.device = torch.device(device)53        self._load_all_models_once()54        self.model_name = model_name55        self.model = self._load_model(model_name)56 57    def _load_all_models_once(self) -> None:58        for name in self.MODEL_DICT:59            self._load_model(name)60 61    def _load_model(self, name: str) -> nn.Module:62        dic = self.MODEL_DICT[name]63        return init_detector(dic['config'], dic['model'], device=self.device)64 65    def set_model(self, name: str) -> None:66        if name == self.model_name:67            return68        self.model_name = name69        self.model = self._load_model(name)70 71    def detect_and_visualize(72        self, image: np.ndarray, score_threshold: float73    ) -> tuple[list[np.ndarray] | tuple[list[np.ndarray],74                                        list[list[np.ndarray]]]75               | dict[str, np.ndarray], np.ndarray]:76        out = self.detect(image)77        vis = self.visualize_detection_results(image, out, score_threshold)78        return out, vis79 80    def detect(81        self, image: np.ndarray82    ) -> list[np.ndarray] | tuple[83            list[np.ndarray], list[list[np.ndarray]]] | dict[str, np.ndarray]:84        image = image[:, :, ::-1]  # RGB -> BGR85        out = inference_detector(self.model, image)86        return out87 88    def visualize_detection_results(89            self,90            image: np.ndarray,91            detection_results: list[np.ndarray]92        | tuple[list[np.ndarray], list[list[np.ndarray]]]93        | dict[str, np.ndarray],94            score_threshold: float = 0.3) -> np.ndarray:95        image = image[:, :, ::-1]  # RGB -> BGR96        vis = self.model.show_result(image,97                                     detection_results,98                                     score_thr=score_threshold,99                                     bbox_color=None,100                                     text_color=(200, 200, 200),101                                     mask_color=None)102        return vis[:, :, ::-1]  # BGR -> RGB103 104 105class AppModel(Model):106    def run(107        self, model_name: str, image: np.ndarray, score_threshold: float108    ) -> tuple[list[np.ndarray] | tuple[list[np.ndarray],109                                        list[list[np.ndarray]]]110               | dict[str, np.ndarray], np.ndarray]:111        self.set_model(model_name)112        return self.detect_and_visualize(image, score_threshold)113