Source code for datumaro.plugins.openvino_plugin.samples.otx_custom_object_detection_gen3_ssd_interp
# Copyright (C) 2023 Intel Corporation
#
# SPDX-License-Identifier: MIT
from typing import List, Tuple
import cv2
from datumaro.components.abstracts import IModelInterpreter
from datumaro.components.abstracts.model_interpreter import LauncherInputType, ModelPred, PrepInfo
from datumaro.components.annotation import Annotation
from datumaro.components.dataset_base import DatasetItem
from datumaro.components.media import Image
from datumaro.plugins.openvino_plugin.samples.utils import (
create_bboxes_with_rescaling,
rescale_img_keeping_aspect_ratio,
)
__all__ = ["OTXSSDModelInterpreter"]
[docs]
class OTXSSDModelInterpreter(IModelInterpreter):
h_model = 864
w_model = 864
[docs]
def preprocess(self, inp: DatasetItem) -> Tuple[LauncherInputType, PrepInfo]:
img = inp.media_as(Image).data
output = rescale_img_keeping_aspect_ratio(img, self.h_model, self.w_model)
# From BGR to RGB
output.image = cv2.cvtColor(output.image, cv2.COLOR_BGR2RGB)
# From HWC to CHW
output.image = output.image.transpose(2, 0, 1)
return output.image, output.scale
[docs]
def postprocess(self, pred: ModelPred, info: PrepInfo) -> List[Annotation]:
scale = info
r_scale = 1 / scale
return create_bboxes_with_rescaling(pred["boxes"], pred["labels"], r_scale)
[docs]
def get_categories(self):
return None