update explain_mode

Author: eugene123tw

src/otx/core/model/detection.py

@@ -241,6 +241,7 @@ def forward_tiles(self, inputs: OTXTileBatchDataEntity[DetBatchDataEntity]) -> D
             inputs.imgs_info,
             self.num_classes,
             self.tile_config,
+            self.explain_mode,
         )
         for batch_tile_attrs, batch_tile_input in inputs.unbind():
             output = self.forward_explain(batch_tile_input) if self.explain_mode else self.forward(batch_tile_input)

src/otx/core/model/instance_segmentation.py

@@ -232,6 +232,7 @@ def forward_tiles(self, inputs: OTXTileBatchDataEntity[InstanceSegBatchDataEntit
             inputs.imgs_info,
             self.num_classes,
             self.tile_config,
+            self.explain_mode,
         )
         for batch_tile_attrs, batch_tile_input in inputs.unbind():
             output = self.forward_explain(batch_tile_input) if self.explain_mode else self.forward(batch_tile_input)

src/otx/core/model/segmentation.py

@@ -245,6 +245,7 @@ def forward_tiles(self, inputs: OTXTileBatchDataEntity[SegBatchDataEntity]) -> S
             inputs.imgs_info,
             self.num_classes,
             self.tile_config,
+            self.explain_mode,
         )
         for batch_tile_attrs, batch_tile_input in inputs.unbind():
             tile_size = batch_tile_attrs[0]["tile_size"]

src/otx/core/utils/tile_merge.py

@@ -28,23 +28,25 @@ class TileMerge(Generic[T_OTXDataEntity, T_OTXBatchPredEntity]):
 
     Args:
         img_infos (list[ImageInfo]): Original image information before tiling.
-        iou_threshold (float, optional): IoU threshold for non-maximum suppression. Defaults to 0.45.
-        max_num_instances (int, optional): Maximum number of instances to keep. Defaults to 500.
-
+        num_classes (int): Number of classes.
+        tile_config (TileConfig): Tile configuration.
+        explain_mode (bool, optional): Whether or not tiles have explain features. Default: False.
     """
 
     def __init__(
         self,
         img_infos: list[ImageInfo],
         num_classes: int,
         tile_config: TileConfig,
+        explain_mode: bool = False,
     ) -> None:
         self.img_infos = img_infos
         self.num_classes = num_classes
         self.tile_size = tile_config.tile_size
         self.iou_threshold = tile_config.iou_threshold
         self.max_num_instances = tile_config.max_num_instances
         self.with_full_img = tile_config.with_full_img
+        self.explain_mode = explain_mode
 
     @abstractmethod
     def _merge_entities(
@@ -116,10 +118,10 @@ def merge(
         """
         entities_to_merge = defaultdict(list)
         img_ids = []
-        explain_mode = len(batch_tile_preds[0].feature_vector) > 0
+        explain_mode = self.explain_mode
 
-        for tile_preds, tile_attrs in zip(batch_tile_preds, batch_tile_attrs):
-            batch_size = tile_preds.batch_size
+        for tile_preds, tile_attrs in zip(batch_tile_preds, batch_tile_attrs, strict=True):
+            batch_size = len(tile_attrs)
             saliency_maps = tile_preds.saliency_map if explain_mode else [[] for _ in range(batch_size)]
             feature_vectors = tile_preds.feature_vector if explain_mode else [[] for _ in range(batch_size)]
             for tile_attr, tile_img_info, tile_bboxes, tile_labels, tile_scores, tile_s_map, tile_f_vect in zip(
@@ -130,6 +132,7 @@ def merge(
                 tile_preds.scores,
                 saliency_maps,
                 feature_vectors,
+                strict=True,
             ):
                 offset_x, offset_y, _, _ = tile_attr["roi"]
                 tile_bboxes[:, 0::2] += offset_x
@@ -155,7 +158,7 @@ def merge(
 
         return [
             self._merge_entities(image_info, entities_to_merge[img_id], explain_mode)
-            for img_id, image_info in zip(img_ids, self.img_infos)
+            for img_id, image_info in zip(img_ids, self.img_infos, strict=True)
         ]
 
     def _merge_entities(
@@ -316,10 +319,10 @@ def merge(
         """
         entities_to_merge = defaultdict(list)
         img_ids = []
-        explain_mode = len(batch_tile_preds[0].feature_vector) > 0
+        explain_mode = self.explain_mode
 
-        for tile_preds, tile_attrs in zip(batch_tile_preds, batch_tile_attrs):
-            feature_vectors = tile_preds.feature_vector if explain_mode else [[] for _ in range(tile_preds.batch_size)]
+        for tile_preds, tile_attrs in zip(batch_tile_preds, batch_tile_attrs, strict=True):
+            feature_vectors = tile_preds.feature_vector if explain_mode else [[] for _ in range(len(tile_attrs))]
             for tile_attr, tile_img_info, tile_bboxes, tile_labels, tile_scores, tile_masks, tile_f_vect in zip(
                 tile_attrs,
                 tile_preds.imgs_info,
@@ -328,6 +331,7 @@ def merge(
                 tile_preds.scores,
                 tile_preds.masks,
                 feature_vectors,
+                strict=True,
             ):
                 keep_indices = tile_masks.to_sparse().sum((1, 2)).to_dense() > 0
                 keep_indices = keep_indices.nonzero(as_tuple=True)[0]
@@ -362,7 +366,7 @@ def merge(
 
         return [
             self._merge_entities(image_info, entities_to_merge[img_id], explain_mode)
-            for img_id, image_info in zip(img_ids, self.img_infos)
+            for img_id, image_info in zip(img_ids, self.img_infos, strict=True)
         ]
 
     def _merge_entities(
@@ -454,6 +458,18 @@ def get_saliency_maps_from_masks(
 class SegmentationTileMerge(TileMerge):
     """Semantic segmentation tile merge."""
 
+    def __init__(
+        self,
+        img_infos: list[ImageInfo],
+        num_classes: int,
+        tile_config: TileConfig,
+        explain_mode: bool = False,
+    ) -> None:
+        super().__init__(img_infos, num_classes, tile_config, explain_mode)
+        if explain_mode:
+            msg = "Explain mode is not supported for segmentation"
+            raise ValueError(msg)
+
     def merge(
         self,
         batch_tile_preds: list[SegBatchPredEntity],
@@ -470,7 +486,7 @@ def merge(
         """
         entities_to_merge = defaultdict(list)
         img_ids = []
-        explain_mode = len(batch_tile_preds[0].feature_vector) > 0
+        explain_mode = self.explain_mode
 
         for tile_preds, tile_attrs in zip(batch_tile_preds, batch_tile_attrs):
             batch_size = tile_preds.batch_size
@@ -538,15 +554,9 @@ def _merge_entities(
             ]
         full_logits_mask = full_logits_mask / vote_mask.unsqueeze(0)
 
-        seg_pred_entity = SegPredEntity(
+        return SegPredEntity(
             image=torch.empty(img_size),
             img_info=img_info,
             masks=full_logits_mask.argmax(0).unsqueeze(0),
             score=[],
         )
-
-        if explain_mode:
-            msg = "Explain mode is not supported for segmentation task."
-            raise NotImplementedError(msg)
-
-        return seg_pred_entity