diff --git a/model_api/src/model_api/tilers/detection.py b/model_api/src/model_api/tilers/detection.py index 917ab294..d32bb086 100644 --- a/model_api/src/model_api/tilers/detection.py +++ b/model_api/src/model_api/tilers/detection.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 # +from typing import TypeVar, cast + import cv2 as cv import numpy as np @@ -12,8 +14,12 @@ from .tiler import Tiler +# Bound to DetectionResult so subclasses (e.g. instance segmentation) can specialize +# the tiler's result type while reusing this detection tiling/merging implementation. +DetResultT = TypeVar("DetResultT", bound=DetectionResult) + -class DetectionTiler(Tiler): +class DetectionTiler(Tiler[DetResultT]): """Tiler for object detection models. This tiler expects model to output a lsit of `Detection` objects or one `DetectionResult` object. @@ -65,7 +71,7 @@ def _postprocess_tile( return output_dict - def _merge_results(self, results: list[dict], shape: tuple[int, int, int]) -> DetectionResult: + def _merge_results(self, results: list[dict], shape: tuple[int, int, int]) -> DetResultT: """Merge results from all tiles. To merge detections, per-class NMS is applied. @@ -101,13 +107,16 @@ def _merge_results(self, results: list[dict], shape: tuple[int, int, int]) -> De saliency_map = self._merge_saliency_maps(saliency_maps, shape, tiles_coords) if saliency_maps else np.ndarray(0) label_names = [self.model.get_label_name(int(label_idx)) for label_idx in detections_array[:, 0]] - return DetectionResult( - bboxes=detections_array[:, 2:].astype(np.int32), - labels=detections_array[:, 0].astype(np.int32), - scores=detections_array[:, 1], - label_names=label_names, - saliency_map=saliency_map, - feature_vector=merged_vector, + return cast( + "DetResultT", + DetectionResult( + bboxes=detections_array[:, 2:].astype(np.int32), + labels=detections_array[:, 0].astype(np.int32), + scores=detections_array[:, 1], + label_names=label_names, + saliency_map=saliency_map, + feature_vector=merged_vector, + ), ) def _merge_saliency_maps( diff --git a/model_api/src/model_api/tilers/instance_segmentation.py b/model_api/src/model_api/tilers/instance_segmentation.py index 5a4ff895..417a96a9 100644 --- a/model_api/src/model_api/tilers/instance_segmentation.py +++ b/model_api/src/model_api/tilers/instance_segmentation.py @@ -15,7 +15,7 @@ from .detection import DetectionTiler -class InstanceSegmentationTiler(DetectionTiler): +class InstanceSegmentationTiler(DetectionTiler[InstanceSegmentationResult]): """Tiler for object instance segmentation models. This tiler expects model to output a list of `SegmentedObject` objects. @@ -190,18 +190,21 @@ def _merge_saliency_maps(self, saliency_maps, shape, tiles_coords): return merged_map - def __call__(self, inputs): - @contextmanager - def setup_maskrcnn(*args, **kwds): - postprocess_state = None + @contextmanager + def _setup_model(self): + """Disable semantic-mask postprocessing on the underlying model while tiling. + + Instance-segmentation masks are merged back to the full image by the tiler + (:meth:`_merge_results`), so per-tile semantic-mask postprocessing must be + turned off for both the full-image (:meth:`__call__`) and pre-tiled + (:meth:`predict_tiles`) inference paths. + """ + postprocess_state = None + if isinstance(self.model, InstanceSegmentationModel): + postprocess_state = self.model.params.postprocess_semantic_masks + self.model._postprocess_semantic_masks = False # noqa: SLF001 + try: + yield + finally: if isinstance(self.model, InstanceSegmentationModel): - postprocess_state = self.model.params.postprocess_semantic_masks - self.model._postprocess_semantic_masks = False # noqa: SLF001 - try: - yield - finally: - if isinstance(self.model, InstanceSegmentationModel): - self.model._postprocess_semantic_masks = postprocess_state # noqa: SLF001 - - with setup_maskrcnn(): - return super().__call__(inputs) + self.model._postprocess_semantic_masks = postprocess_state # noqa: SLF001 diff --git a/model_api/src/model_api/tilers/semantic_segmentation.py b/model_api/src/model_api/tilers/semantic_segmentation.py index 32f40990..37328aeb 100644 --- a/model_api/src/model_api/tilers/semantic_segmentation.py +++ b/model_api/src/model_api/tilers/semantic_segmentation.py @@ -14,7 +14,7 @@ from .tiler import Tiler -class SemanticSegmentationTiler(Tiler): +class SemanticSegmentationTiler(Tiler[ImageResultWithSoftPrediction]): """Tiler for segmentation models.""" def _postprocess_tile( @@ -69,18 +69,21 @@ def _merge_results( saliency_map=np.array([]), ) - def __call__(self, inputs): - @contextmanager - def setup_segm_model(): - return_soft_prediction_state = None + @contextmanager + def _setup_model(self): + """Enable soft predictions on the underlying model while tiling. + + Overlapping tiles are merged by accumulating per-class soft logits + (:meth:`_merge_results`), so the underlying segmentation model must return + soft predictions for both the full-image (:meth:`__call__`) and pre-tiled + (:meth:`predict_tiles`) inference paths. + """ + return_soft_prediction_state = None + if isinstance(self.model, SegmentationModel): + return_soft_prediction_state = self.model.params.return_soft_prediction + self.model._return_soft_prediction = True # noqa: SLF001 + try: + yield + finally: if isinstance(self.model, SegmentationModel): - return_soft_prediction_state = self.model.params.return_soft_prediction - self.model._return_soft_prediction = True # noqa: SLF001 - try: - yield - finally: - if isinstance(self.model, SegmentationModel): - self.model._return_soft_prediction = return_soft_prediction_state # noqa: SLF001 - - with setup_segm_model(): - return super().__call__(inputs) + self.model._return_soft_prediction = return_soft_prediction_state # noqa: SLF001 diff --git a/model_api/src/model_api/tilers/tiler.py b/model_api/src/model_api/tilers/tiler.py index e3398a26..788f3c8b 100644 --- a/model_api/src/model_api/tilers/tiler.py +++ b/model_api/src/model_api/tilers/tiler.py @@ -5,13 +5,18 @@ import abc import logging as log +from contextlib import contextmanager from itertools import product +from typing import Generic, TypeVar from model_api.models.parameters import ParameterRegistry +from model_api.models.result import Result from model_api.pipelines import AsyncPipeline +ResultT = TypeVar("ResultT", bound=Result) -class Tiler(abc.ABC): + +class Tiler(abc.ABC, Generic[ResultT]): EXECUTION_MODES = ("async", "sync") """ An abstract tiler @@ -119,7 +124,7 @@ def _load_config(self, config): f'The parameter "{name}" not found in tiler, will be omitted', ) - def __call__(self, inputs): + def __call__(self, inputs) -> ResultT: """Applies full pipeline of tiling inference in one call. Args: @@ -131,9 +136,70 @@ def __call__(self, inputs): tile_coords = self._tile(inputs) tile_coords = self._filter_tiles(inputs, tile_coords) - if self.execution_mode == "sync": - return self._predict_sync(inputs, tile_coords) - return self._predict_async(inputs, tile_coords) + with self._setup_model(): + if self.execution_mode == "sync": + return self._predict_sync(inputs, tile_coords) + return self._predict_async(inputs, tile_coords) + + def predict_tiles(self, tiles, tile_coords, shape) -> ResultT: + """Run tiled inference on externally provided (already cropped) tiles and merge. + + This mirrors :meth:`__call__` but skips the internal tiling stage + (``_tile`` / ``_filter_tiles`` / ``_crop_tile``): the caller supplies the tile + crops together with their absolute coordinates in the full image. This supports + pipelines that perform tiling upstream (e.g. inside a data loader) while still + reusing the tiler's per-task ``_postprocess_tile`` and ``_merge_results`` logic. + + The underlying model runs once per provided tile, so callers must not pass tiles + through a model that itself tiles internally (that would tile each tile a second + time). + + Args: + tiles: sequence of pre-cropped tile images, each in the layout expected by + the underlying model wrapper (typically ``(H, W, C)`` arrays). + tile_coords: sequence of ``[x1, y1, x2, y2]`` absolute tile coordinates in the + full image, aligned with ``tiles``. + shape: full-resolution image shape (``(H, W, C)`` or ``(H, W)``) used to place + and merge the per-tile predictions. + + Returns: + Merged prediction in the format produced by ``_merge_results``. + + Raises: + ValueError: if ``tiles`` and ``tile_coords`` have different lengths. + """ + tiles = list(tiles) + tile_coords = list(tile_coords) + if len(tiles) != len(tile_coords): + msg = "tiles and tile_coords must have the same length" + raise ValueError(msg) + + with self._setup_model(): + if self.execution_mode == "sync": + tile_results = [ + self._postprocess_tile(self.model(tile), coord) for tile, coord in zip(tiles, tile_coords) + ] + else: + for i, tile in enumerate(tiles): + self.async_pipeline.submit_data(tile, i) + self.async_pipeline.await_all() + tile_results = [ + self._postprocess_tile(self.async_pipeline.get_result(j)[0], tile_coords[j]) + for j in range(len(tiles)) + ] + return self._merge_results(tile_results, shape) + + @contextmanager + def _setup_model(self): + """Temporarily reconfigure the underlying model for tiled inference. + + The base implementation is a no-op. Subclasses override this to toggle + model-specific flags for the duration of tiled inference and merging (for + example, disabling semantic-mask postprocessing for instance segmentation or + enabling soft predictions for semantic segmentation). The reconfiguration is + applied consistently by both :meth:`__call__` and :meth:`predict_tiles`. + """ + yield def _tile(self, image): """Tiles an input image to overlapping or non-overlapping patches. @@ -173,7 +239,7 @@ def _filter_tiles(self, image, tile_coords): """ return tile_coords - def _predict_sync(self, image, tile_coords): + def _predict_sync(self, image, tile_coords) -> ResultT: """Makes prediction by splitting the input image into tiles in synchronous mode. Args: @@ -192,7 +258,7 @@ def _predict_sync(self, image, tile_coords): return self._merge_results(tile_results, image.shape) - def _predict_async(self, image, tile_coords): + def _predict_async(self, image, tile_coords) -> ResultT: """Makes prediction by splitting the input image into tiles in asynchronous mode. Args: @@ -228,7 +294,7 @@ def _postprocess_tile(self, predictions, coord): """ @abc.abstractmethod - def _merge_results(self, results, shape): + def _merge_results(self, results, shape) -> ResultT: """Merge results from all tiles. Args: diff --git a/model_api/tests/unit/tilers/test_detection_tiler.py b/model_api/tests/unit/tilers/test_detection_tiler.py index 55b3bbc8..fd643bf4 100644 --- a/model_api/tests/unit/tilers/test_detection_tiler.py +++ b/model_api/tests/unit/tilers/test_detection_tiler.py @@ -4,6 +4,7 @@ from unittest.mock import MagicMock, patch import numpy as np +import pytest from model_api.models import DetectionResult from model_api.tilers.detection import DetectionTiler, _non_linear_normalization @@ -128,6 +129,28 @@ def test_merge_results_multiple_tiles(self, mock_nms): assert isinstance(merged, DetectionResult) +class TestDetectionTilerPredictTiles: + @patch("model_api.tilers.detection.multiclass_nms") + def test_predict_tiles_sync_end_to_end(self, mock_nms): + """predict_tiles infers each pre-cropped tile once and merges to a DetectionResult.""" + model = _make_model() + mock_nms.return_value = np.array([0, 1]) + tiler = DetectionTiler(model, execution_mode="sync") + # One detection per tile; the model is called once per provided tile. + model.side_effect = [_make_detection_result(n=1), _make_detection_result(n=1)] + tiles = [np.zeros((50, 50, 3), dtype=np.uint8), np.zeros((50, 50, 3), dtype=np.uint8)] + coords = [[0, 0, 50, 50], [50, 0, 100, 50]] + merged = tiler.predict_tiles(tiles, coords, (50, 100, 3)) + assert isinstance(merged, DetectionResult) + assert model.call_count == 2 + + def test_predict_tiles_length_mismatch(self): + model = _make_model() + tiler = DetectionTiler(model, execution_mode="sync") + with pytest.raises(ValueError, match="same length"): + tiler.predict_tiles([np.zeros((5, 5, 3), dtype=np.uint8)], [], (5, 5, 3)) + + class TestDetectionTilerMergeSaliencyMaps: def test_merge_saliency_maps_empty(self): model = _make_model() diff --git a/model_api/tests/unit/tilers/test_instance_segmentation_tiler.py b/model_api/tests/unit/tilers/test_instance_segmentation_tiler.py index 02384a08..38aa27c6 100644 --- a/model_api/tests/unit/tilers/test_instance_segmentation_tiler.py +++ b/model_api/tests/unit/tilers/test_instance_segmentation_tiler.py @@ -189,6 +189,43 @@ def test_merge_saliency_maps_zero_sum(self): assert result[0].shape == (0,) +class TestInstanceSegmentationTilerPredictTiles: + @patch("model_api.tilers.instance_segmentation.multiclass_nms") + @patch("model_api.tilers.instance_segmentation._segm_postprocess") + def test_predict_tiles_sync_end_to_end(self, mock_segm, mock_nms): + """predict_tiles infers each pre-cropped tile once and merges masks + boxes.""" + model = _make_model() + mock_nms.return_value = np.array([0]) + mock_segm.return_value = np.ones((50, 100), dtype=np.uint8) + tiler = InstanceSegmentationTiler(model, execution_mode="sync") + model.side_effect = [_make_instance_seg_result(n=1), _make_instance_seg_result(n=1)] + tiles = [np.zeros((50, 50, 3), dtype=np.uint8), np.zeros((50, 50, 3), dtype=np.uint8)] + coords = [[0, 0, 50, 50], [50, 0, 100, 50]] + # Saliency-map merging is exercised separately; keep this focused on box/mask merging. + with patch.object(tiler, "_merge_saliency_maps", return_value=[]): + merged = tiler.predict_tiles(tiles, coords, (50, 100, 3)) + assert isinstance(merged, InstanceSegmentationResult) + assert model.call_count == 2 + + def test_setup_model_toggles_postprocess_flag(self): + """_setup_model disables semantic-mask postprocessing while tiling, then restores it.""" + from model_api.models.instance_segmentation import MaskRCNNModel + + model = MagicMock(spec=MaskRCNNModel) + model.load = MagicMock() + model.inference_adapter = MagicMock() + model.inference_adapter.get_rt_info.side_effect = RuntimeError( + "Cannot get runtime attribute. Path to runtime attribute is incorrect.", + ) + model.params = MagicMock() + model.params.postprocess_semantic_masks = True + + tiler = InstanceSegmentationTiler(model, execution_mode="sync") + with tiler._setup_model(): # noqa: SLF001 + assert model._postprocess_semantic_masks is False # noqa: SLF001 + assert model._postprocess_semantic_masks is True # noqa: SLF001 + + class TestInstanceSegmentationTilerCall: def test_call_with_maskrcnn(self): model = _make_model() diff --git a/model_api/tests/unit/tilers/test_semantic_segmentation_tiler.py b/model_api/tests/unit/tilers/test_semantic_segmentation_tiler.py index e639ff77..72f8dc8c 100644 --- a/model_api/tests/unit/tilers/test_semantic_segmentation_tiler.py +++ b/model_api/tests/unit/tilers/test_semantic_segmentation_tiler.py @@ -75,6 +75,39 @@ def test_merge_results_single_tile(self): np.testing.assert_array_almost_equal(merged.soft_prediction, soft) +class TestSemanticSegmentationTilerPredictTiles: + def test_predict_tiles_sync_end_to_end(self): + """predict_tiles infers each pre-cropped tile once and merges soft predictions.""" + model = _make_model(num_labels=2) + tiler = SemanticSegmentationTiler(model, execution_mode="sync") + model.side_effect = [_make_seg_result(50, 50, num_classes=2), _make_seg_result(50, 50, num_classes=2)] + tiles = [np.zeros((50, 50, 3), dtype=np.uint8), np.zeros((50, 50, 3), dtype=np.uint8)] + coords = [[0, 0, 50, 50], [50, 0, 100, 50]] + merged = tiler.predict_tiles(tiles, coords, (50, 100, 3)) + assert isinstance(merged, ImageResultWithSoftPrediction) + assert merged.soft_prediction.shape == (50, 100, 2) + assert model.call_count == 2 + + def test_setup_model_toggles_soft_prediction_flag(self): + """_setup_model enables soft predictions while tiling, then restores the prior state.""" + from model_api.models.segmentation import SegmentationModel + + model = MagicMock(spec=SegmentationModel) + model.load = MagicMock() + model.inference_adapter = MagicMock() + model.inference_adapter.get_rt_info.side_effect = RuntimeError( + "Cannot get runtime attribute. Path to runtime attribute is incorrect.", + ) + model.params = MagicMock() + model.params.labels = ["a", "b"] + model.params.return_soft_prediction = False + + tiler = SemanticSegmentationTiler(model, execution_mode="sync") + with tiler._setup_model(): # noqa: SLF001 + assert model._return_soft_prediction is True # noqa: SLF001 + assert model._return_soft_prediction is False # noqa: SLF001 + + class TestSemanticSegmentationTilerCall: def test_call_with_segmentation_model(self): from model_api.models.segmentation import SegmentationModel diff --git a/model_api/tests/unit/tilers/test_tiler.py b/model_api/tests/unit/tilers/test_tiler.py index 8be0a7ca..15488b95 100644 --- a/model_api/tests/unit/tilers/test_tiler.py +++ b/model_api/tests/unit/tilers/test_tiler.py @@ -1,6 +1,7 @@ # Copyright (C) 2025 Intel Corporation # SPDX-License-Identifier: Apache-2.0 +from contextlib import contextmanager from unittest.mock import MagicMock, patch import numpy as np @@ -233,3 +234,86 @@ def test_predict_async(self, mock_pipeline_cls): assert mock_pipeline.submit_data.call_count == 2 mock_pipeline.await_all.assert_called_once() assert len(results) == 2 + + +class TestTilerPredictTiles: + def test_predict_tiles_sync(self): + """Pre-cropped tiles are inferred once each and merged, skipping internal tiling.""" + model = _make_model() + model.return_value = "pred" + tiler = ConcreteTiler(model, execution_mode="sync") + tiles = [np.zeros((10, 10, 3)), np.zeros((10, 10, 3))] + coords = [[0, 0, 10, 10], [10, 0, 20, 10]] + results = tiler.predict_tiles(tiles, coords, (10, 20, 3)) + assert model.call_count == 2 + assert len(results) == 2 + assert results[0]["coord"] == coords[0] + + @patch("model_api.tilers.tiler.AsyncPipeline") + def test_predict_tiles_async(self, mock_pipeline_cls): + model = _make_model() + mock_pipeline = MagicMock() + mock_pipeline_cls.return_value = mock_pipeline + mock_pipeline.get_result.return_value = ("pred", {}) + tiler = ConcreteTiler(model, execution_mode="async") + tiles = [np.zeros((10, 10, 3)), np.zeros((10, 10, 3))] + coords = [[0, 0, 10, 10], [10, 0, 20, 10]] + results = tiler.predict_tiles(tiles, coords, (10, 20, 3)) + assert mock_pipeline.submit_data.call_count == 2 + mock_pipeline.await_all.assert_called_once() + assert len(results) == 2 + + def test_predict_tiles_length_mismatch(self): + model = _make_model() + tiler = ConcreteTiler(model, execution_mode="sync") + with pytest.raises(ValueError, match="same length"): + tiler.predict_tiles([np.zeros((5, 5, 3))], [], (5, 5, 3)) + + def test_predict_tiles_invokes_setup_model(self): + """predict_tiles must apply the _setup_model reconfiguration, like __call__.""" + model = _make_model() + model.return_value = "pred" + + class RecordingTiler(ConcreteTiler): + def __init__(self, *args, **kwargs): + self.setup_called = 0 + super().__init__(*args, **kwargs) + + @contextmanager + def _setup_model(self): + self.setup_called += 1 + yield + + tiler = RecordingTiler(model, execution_mode="sync") + tiler.predict_tiles([np.zeros((5, 5, 3))], [[0, 0, 5, 5]], (5, 5, 3)) + assert tiler.setup_called == 1 + + +class TestTilerSetupModel: + def test_default_setup_model_is_noop(self): + model = _make_model() + tiler = ConcreteTiler(model) + with tiler._setup_model(): # noqa: SLF001 + pass + + def test_call_invokes_setup_model(self): + model = _make_model() + model.return_value = "pred" + + class RecordingTiler(ConcreteTiler): + def __init__(self, *args, **kwargs): + self.setup_called = 0 + super().__init__(*args, **kwargs) + + @contextmanager + def _setup_model(self): + self.setup_called += 1 + yield + + tiler = RecordingTiler( + model, + configuration={"tile_size": 100, "tiles_overlap": 0.0, "tile_with_full_img": False}, + execution_mode="sync", + ) + tiler(np.zeros((100, 100, 3), dtype=np.uint8)) + assert tiler.setup_called == 1