MMDetection

from mmengine.registry import OPTIMIZERS

@OPTIMIZERS.register_module()
class CustomOptim:
    pass

# -*- coding: utf-8 -*-

from mmdet.datasets.coco import CocoDataset
from mmdet.registry import DATASETS


@DATASETS.register_module()
class PlateCocoDataset(CocoDataset):
    """Dataset for COCO"""

    METAINFO = {
        "classes": ("Plate",),
        # palette is a list of color tuples, which is used for visualization.
        "palette": [
            (106, 0, 228),
        ],
    }

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._metainfo["dataset_type"] = None

# -*- coding: utf-8 -*-

from mmdet.datasets.xml_style import XMLDataset
from mmdet.registry import DATASETS


@DATASETS.register_module()
class SeahVOCDataset(XMLDataset):
    """Dataset for PASCAL VOC."""

    METAINFO = {
        "classes": ("Bloom", "Ingot", "Billet", "Slab", "Rod", "Forklift", "Support"),
        # palette is a list of color tuples, which is used for visualization.
        "palette": [
            (106, 0, 228),
            (119, 11, 32),
            (165, 42, 42),
            (0, 0, 192),
            (197, 226, 255),
            (0, 60, 100),
            (0, 0, 142),
        ],
    }

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._metainfo["dataset_type"] = None

# -*- coding: utf-8 -*-

import os
from datetime import datetime
from glob import glob
from pathlib import Path
from sys import exit as sys_exit
from typing import List, Tuple, Union

import cv2
from mmdet.apis import init_detector, inference_detector
from mmdet.datasets.coco import CocoDataset
from mmdet.registry import DATASETS, VISUALIZERS
from mmdet.structures.det_data_sample import DetDataSample
from numpy import uint8, ndarray
from numpy.typing import NDArray

NumberType = Union[int, float]


@DATASETS.register_module()
class PlateCocoDataset(CocoDataset):
    """Dataset for COCO"""

    METAINFO = {
        "classes": ("Plate",),
        # palette is a list of color tuples, which is used for visualization.
        "palette": [
            (106, 0, 228),
        ],
    }

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._metainfo["dataset_type"] = None


class ObjectSample:
    def __init__(
        self,
        index: int,
        label: str,
        score: float,
        x1: NumberType,
        y1: NumberType,
        x2: NumberType,
        y2: NumberType,
        color: Tuple[int, int, int],
    ):
        self.index = index
        self.label = label
        self.score = score
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2
        self.color = color

def convert_det_to_obj(classes: List[str], palette: List[Tuple[int, int, int]], sample: DetDataSample) -> List[ObjectSample]:
    result = list()
    pred = sample.pred_instances.numpy()
    for i in range(len(pred)):
        x1, y1, x2, y2 = pred.bboxes[i].tolist()  # noqa
        label_index = pred.labels[i].tolist()  # noqa
        label_name = classes[label_index]
        color = palette[label_index]
        score = pred.scores[i].tolist()  # noqa
        obj = ObjectSample(label_index, label_name, score, x1, y1, x2, y2, color)
        result.append(obj)
    return result

def visualize(visualizer, score_threshold: float, image: NDArray[uint8], data_sample: DetDataSample) -> NDArray[uint8]:
    rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    visualizer.add_datasample(
        name="Preview",
        image=rgb_image,
        data_sample=data_sample,
        draw_gt=False,
        pred_score_thr=score_threshold,
        show=False,
    )
    visualized_image = visualizer.get_image()
    return cv2.cvtColor(visualized_image, cv2.COLOR_RGB2BGR)

def main() -> None:
    config_file = "./rtmdet-ins_x_8xb16-300e_coco_dongkuk_rolling_plate_cam1_2nd.py"
    checkpoint_file = "./epoch_300.pth"
    model = init_detector(config_file, checkpoint_file, device="cuda:0")  # or device="cpu"

    files = glob("./input/*.jpg")
    files.sort()

    output_dir = Path("output")
    if not output_dir.is_dir():
        output_dir.mkdir()

    visualizer = VISUALIZERS.build(model.cfg.visualizer)
    visualizer.dataset_meta = model.dataset_meta

    for filename in files:
        frame = cv2.imread(filename)

        begin = datetime.now()
        result = inference_detector(model, frame)
        assert isinstance(result, DetDataSample)
        duration = (datetime.now() - begin).total_seconds()
        print(f"Predict duration: {duration:.04f}s")

        # print(result)

        begin = datetime.now()
        try:
            visualized_image = visualize(visualizer, 0.5, frame, result)
        except BaseException as e:
            print(f"Visualizer error: {e}")
            continue
        duration = (datetime.now() - begin).total_seconds()
        print(f"Visualizer duration: {duration:.04f}s")

        file_prefix, ext = os.path.splitext(os.path.basename(filename))
        output_filename = file_prefix + ".png"
        output_filepath = str(output_dir / output_filename)
        print(output_filepath)
        cv2.imwrite(output_filepath, visualized_image)

        classes = visualizer.dataset_meta["classes"]
        palette = visualizer.dataset_meta["palette"]
        objs = convert_det_to_obj(classes, palette, result)
        # print(objs)

if __name__ == "__main__":
    sys_exit(main())

# -*- coding: utf-8 -*-

from os import path
from typing import Any, List, Optional, Sequence, Tuple

from cvlayer.cv.cvt_color import CvtColorCode, cvt_color
from mmdet.apis import inference_detector, init_detector  # noqa
from mmdet.registry import VISUALIZERS  # noqa
from mmdet.structures.det_data_sample import DetDataSample  # noqa
from numpy import uint8, zeros
from numpy.typing import NDArray
from overrides import override

from ddrm.arguments import DEFAULT_DEVICE, DEFAULT_SCORE_THRESHOLD
from ddrm.models.interface import ModelInterface
from ddrm.models.mmdet.datasets.plate_coco import PlateCocoDataset  # noqa
from ddrm.models.sample import Sample


class MMDetModel(ModelInterface):
    def __init__(
        self,
        config: str,
        checkpoint: str,
        device=DEFAULT_DEVICE,
        threshold=DEFAULT_SCORE_THRESHOLD,
        ignore_labels: Optional[Sequence[str]] = None,
        extra: Optional[str] = None,
        name: Optional[str] = None,
    ):
        if not path.isfile(config):
            raise FileNotFoundError(f"Config file not found: '{config}'")
        if not path.isfile(checkpoint):
            raise FileNotFoundError(f"Checkpoint file not found: '{checkpoint}'")

        self._model = init_detector(config, checkpoint, device=device)
        self._ignore_labels = ignore_labels
        self._extra = extra if extra else str()
        self._name = name if name else type(self).__class__.__name__
        self._visualizer = VISUALIZERS.build(self._model.cfg.visualizer)
        self._threshold = threshold

        # the dataset_meta is loaded from the checkpoint and
        # then pass to the model in init_detector
        self._visualizer.dataset_meta = self._model.dataset_meta

    @property
    def name(self) -> str:
        return self._name

    @property
    def threshold(self) -> float:
        return self._threshold

    @property
    def classes(self) -> List[str]:
        return self._visualizer.dataset_meta["classes"]

    @property
    def palette(self) -> List[Tuple[int, int, int]]:
        return self._visualizer.dataset_meta["palette"]

    def convert_det_to_samples(self, det: DetDataSample) -> List[Sample]:
        result = list()
        pred = det.pred_instances.numpy()
        for i in range(len(pred)):
            label_index = pred.labels[i].tolist()  # noqa
            label_name = self.classes[label_index]
            if self._ignore_labels and label_name in self._ignore_labels:
                continue
            x1, y1, x2, y2 = pred.bboxes[i].tolist()  # noqa
            score = pred.scores[i].tolist()  # noqa
            color = self.palette[label_index]
            mask = zeros(pred.masks[i].shape, dtype=uint8)  # noqa
            mask[pred.masks[i]] = 255
            sample = Sample(label_index, label_name, score, x1, y1, x2, y2, color, mask)
            result.append(sample)
        return result

    def inference(self, image: NDArray) -> DetDataSample:
        return inference_detector(self._model, image)

    def visualize(self, image: NDArray, det: DetDataSample) -> NDArray:
        self._visualizer.add_datasample(
            name=self._name,
            image=cvt_color(image, CvtColorCode.BGR2RGB),
            data_sample=det,
            draw_gt=False,
            pred_score_thr=self._threshold,
            show=False,
        )
        return cvt_color(self._visualizer.get_image(), CvtColorCode.RGB2BGR)

    @override
    def get_extra(self) -> Any:
        return self._extra

    @override
    def set_extra(self, extra: Any) -> None:
        self._extra = extra

    @override
    def inference_raw(self, image: NDArray, **kwargs):
        return self.inference(image)

    @override
    def inference_samples(self, image: NDArray, **kwargs) -> List[Sample]:
        return self.convert_det_to_samples(self.inference_raw(image, **kwargs))

    @override
    def run(
        self,
        image: NDArray,
        preview: Optional[NDArray] = None,
        **kwargs,
    ) -> Tuple[List[Sample], Optional[NDArray]]:
        det = self.inference(image)
        samples = self.convert_det_to_samples(det)
        visualized = self.visualize(preview if preview is not None else image, det)
        return samples, visualized