MMSegmentation

import torch
from mmseg.apis import init_segmentor
import mmcv

# 모델과 구성 파일 경로
config_file = 'configs/fcn/fcn_r50-d8_512x512_40k_voc12aug.py'
checkpoint_file = 'checkpoints/fcn_r50-d8_512x512_40k_voc12aug_20200617_071005-41a67fa4.pth'

# 모델 초기화
model = init_segmentor(config_file, checkpoint_file, device='cuda:0')

# 입력 데이터 생성 (예시 입력)
dummy_input = torch.randn(1, 3, 512, 512).cuda()

# 모델을 평가 모드로 전환
model.eval()

# 모델 변환
torch.onnx.export(model, dummy_input, "mmsegmentation_model.onnx", 
                  export_params=True, opset_version=11, 
                  input_names=['input'], output_names=['output'])

import tritonclient.http as httpclient
import tritonclient.utils.shared_memory as shm
import numpy as np

# Triton Inference Server URL
url = "localhost:8000"

# 클라이언트 초기화
client = httpclient.InferenceServerClient(url=url)

# 모델 정보
model_name = "mmsegmentation_model"
input_name = "input"
output_name = "output"

# 입력 데이터 준비
input_data = np.random.rand(1, 3, 512, 512).astype(np.float32)

# Shared memory 설정
input_byte_size = input_data.nbytes
output_byte_size = input_byte_size  # 예제에서는 입력과 출력 크기가 같다고 가정

# Shared memory 생성
input_shm_handle = shm.create_shared_memory_region("input_data", "/dev/shm/input", input_byte_size)
output_shm_handle = shm.create_shared_memory_region("output_data", "/dev/shm/output", output_byte_size)

# Shared memory에 데이터 쓰기
shm.set_shared_memory_region(input_shm_handle, [input_data])

# Shared memory 등록
client.register_system_shared_memory("input_data", "/dev/shm/input", input_byte_size)
client.register_system_shared_memory("output_data", "/dev/shm/output", output_byte_size)

# 입력 및 출력 설정
inputs = []
outputs = []

inputs.append(httpclient.InferInput(input_name, input_data.shape, "FP32"))
inputs[-1].set_shared_memory("input_data", input_byte_size)

outputs.append(httpclient.InferRequestedOutput(output_name, binary_data=True))
outputs[-1].set_shared_memory("output_data", output_byte_size)

# 추론 요청
results = client.infer(model_name, inputs, outputs=outputs)

# 결과 읽기
output_data = shm.get_shared_memory_region(output_shm_handle, output_byte_size, np.float32)
print(output_data)

# Shared memory 해제
client.unregister_system_shared_memory("input_data")
client.unregister_system_shared_memory("output_data")
shm.destroy_shared_memory_region(input_shm_handle)
shm.destroy_shared_memory_region(output_shm_handle)

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

from mmseg.datasets.coco import CocoDataset  # noqa
from mmseg.registry import DATASETS  # noqa


@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 -*-

# [IMPORTANT]
# Do not import the 'mmdet' packages !!

from typing import List, Optional, Tuple

from cvlayer.geometry.iou import calculate_iou
from cvlayer.typing import RectT
from numpy.typing import NDArray


class Sample:
    def __init__(
        self,
        index: int,
        label: str,
        score: float,
        x1: float,
        y1: float,
        x2: float,
        y2: float,
        color: Tuple[int, int, int],
        mask: Optional[NDArray] = None,
    ):
        self.index = index
        self.label = label
        self.score = score
        self.x1 = x1
        self.y1 = y1
        self.x2 = x2
        self.y2 = y2
        self.color = color
        self.mask = mask

    def __repr__(self):
        return f"<Sample index={self.index},label={self.label},score={self.score:.2f}>"

    def __str__(self):
        return f"{self.label}({self.score:.2f})"

    @property
    def shape(self):
        return self.mask.shape

    @property
    def p1(self):
        return self.x1, self.y1

    @property
    def p2(self):
        return self.x2, self.y2

    @property
    def width(self):
        return self.x2 - self.x1

    @property
    def height(self):
        return self.y2 - self.y1

    @property
    def center(self):
        return self.x1 + self.width / 2, self.y1 + self.height / 2

    @property
    def rect(self):
        return self.x1, self.y1, self.x2, self.y2

    def intersection(self, roi: RectT) -> bool:
        left = max(self.x1, roi[0])
        top = max(self.y1, roi[1])
        right = min(self.x2, roi[2])
        bottom = min(self.y2, roi[3])
        return right > left and bottom > top

    def find_nearest_sample(self, slots: List["Sample"]) -> Optional["Sample"]:
        ss = list(filter(lambda s: s.intersection(self.rect), slots))
        if not ss:
            return None
        ious = [calculate_iou(self.rect, s.rect) for s in ss]
        largest_iou = max(ious)
        return ss[ious.index(largest_iou)]

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

from abc import ABCMeta, abstractmethod
from typing import Any, List, Optional, Tuple

from numpy.typing import NDArray

from ddrm.models.sample import Sample


class ModelInterface(metaclass=ABCMeta):
    @abstractmethod
    def inference_raw(self, image: NDArray, **kwargs):
        raise NotImplementedError

    @abstractmethod
    def inference_samples(self, image: NDArray, **kwargs) -> List[Sample]:
        raise NotImplementedError

    @abstractmethod
    def get_extra(self) -> Any:
        raise NotImplementedError

    @abstractmethod
    def set_extra(self, extra: Any) -> None:
        raise NotImplementedError

    @abstractmethod
    def run(
        self,
        image: NDArray,
        preview: Optional[NDArray] = None,
        **kwargs,
    ) -> Tuple[List[Sample], Optional[NDArray]]:
        raise NotImplementedError

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

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

import numpy as np
from cvlayer.cv.cvt_color import CvtColorCode, cvt_color
from mmseg.apis import inference_model, init_model  # noqa
from mmseg.registry import VISUALIZERS  # noqa
from mmseg.structures.seg_data_sample import SegDataSample  # 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.sample import Sample


def get_bounding_box(mask):
    y_indices, x_indices = np.where(mask > 0)

    if len(x_indices) == 0 or len(y_indices) == 0:
        return None

    x1, x2 = np.min(x_indices), np.max(x_indices)
    y1, y2 = np.min(y_indices), np.max(y_indices)

    return int(x1), int(y1), int(x2), int(y2)


class MMSegModel(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_model(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_seg_to_samples(self, seg: SegDataSample) -> List[Sample]:
        result = list()
        num_classes = len(self.classes)
        sem_seg = seg.pred_sem_seg.cpu().data  # noqa
        ids = np.unique(sem_seg)[::-1]
        legal_indices = ids < num_classes
        ids = ids[legal_indices]
        labels = np.array(ids, dtype=np.int64)
        colors = [self.palette[label] for label in labels]
        height = sem_seg.shape[1]
        width = sem_seg.shape[2]
        shape = height, width
        mask = zeros(shape, dtype=uint8)  # noqa
        for label, color in zip(labels, colors):
            label_index = int(label)
            label_name = self.classes[label_index]
            if label_index == 0:
                assert label_name == "background_bg"
                continue
            if self._ignore_labels and label_name in self._ignore_labels:
                continue
            score = 1.0
            mask[sem_seg[0] == label] = 255
            x1, y1, x2, y2 = get_bounding_box(mask)
            sample = Sample(label_index, label_name, score, x1, y1, x2, y2, color, mask)
            result.append(sample)
        return result

    def inference(self, image: NDArray) -> SegDataSample:
        return inference_model(self._model, image)

    def visualize(self, image: NDArray, seg: SegDataSample) -> NDArray:
        self._visualizer.add_datasample(
            name=self._name,
            image=cvt_color(image, CvtColorCode.BGR2RGB),
            data_sample=seg,
            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_seg_to_samples(self.inference_raw(image, **kwargs))

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