Cover ImageModel

Author: sovrasov

model_api/python/model_api/adapters/inference_adapter.py

@@ -169,7 +169,7 @@ def embed_preprocessing(
         layout,
         resize_mode: str,
         interpolation_mode,
-        target_shape: tuple[int],
+        target_shape: tuple[int, ...],
         pad_value,
         dtype: type = int,
         brg2rgb=False,

model_api/python/model_api/adapters/openvino_adapter.py

@@ -338,7 +338,7 @@ def embed_preprocessing(
         layout: str,
         resize_mode: str,
         interpolation_mode: str,
-        target_shape: tuple[int],
+        target_shape: tuple[int, ...],
         pad_value: int,
         dtype: type = int,
         brg2rgb: bool = False,

model_api/python/model_api/models/image_model.py

@@ -3,8 +3,11 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 
+from typing import Any
+
 import numpy as np
 
+from model_api.adapters.inference_adapter import InferenceAdapter
 from model_api.adapters.utils import RESIZE_TYPES, InputTransform
 from model_api.models.model import Model
 from model_api.models.types import BooleanValue, ListValue, NumericalValue, StringValue
@@ -33,7 +36,7 @@ class ImageModel(Model):
 
     __model__ = "ImageModel"
 
-    def __init__(self, inference_adapter, configuration: dict = {}, preload=False):
+    def __init__(self, inference_adapter: InferenceAdapter, configuration: dict = {}, preload: bool = False) -> None:
         """Image model constructor
 
         It extends the `Model` constructor.
@@ -59,6 +62,7 @@ def __init__(self, inference_adapter, configuration: dict = {}, preload=False):
         self.scale_values: list
         self.reverse_input_channels: bool
         self.embedded_processing: bool
+        self.labels: list[str]
 
         self.nchw_layout = self.inputs[self.image_blob_name].layout == "NCHW"
         if self.nchw_layout:
@@ -90,7 +94,7 @@ def __init__(self, inference_adapter, configuration: dict = {}, preload=False):
             self.orig_height, self.orig_width = self.h, self.w
 
     @classmethod
-    def parameters(cls) -> dict:
+    def parameters(cls) -> dict[str, Any]:
         parameters = super().parameters()
         parameters.update(
             {
@@ -137,14 +141,14 @@ def parameters(cls) -> dict:
         )
         return parameters
 
-    def get_label_name(self, label_id):
+    def get_label_name(self, label_id: int) -> str:
         if self.labels is None:
             return f"#{label_id}"
         if label_id >= len(self.labels):
             return f"#{label_id}"
         return self.labels[label_id]
 
-    def _get_inputs(self):
+    def _get_inputs(self) -> tuple[list[str], ...]:
         """Defines the model inputs for images and additional info.
 
         Raises:
@@ -170,7 +174,7 @@ def _get_inputs(self):
             )
         return image_blob_names, image_info_blob_names
 
-    def preprocess(self, inputs) -> list[dict]:
+    def preprocess(self, inputs: np.ndarray) -> list[dict]:
         """Data preprocess method
 
         It performs basic preprocessing of a single image:
@@ -203,7 +207,7 @@ def preprocess(self, inputs) -> list[dict]:
             },
         ]
 
-    def _change_layout(self, image):
+    def _change_layout(self, image: np.ndarray) -> np.ndarray:
         """Changes the input image layout to fit the layout of the model input layer.
 
         Args: