Moved forward tiling patch to Model API

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(

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

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

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:

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()

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()

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