Refactor raw logits processing in DFine and RTDETR models by introduc…

…ing a dedicated method for splitting and reshaping logits in explain mode.

src/otx/algo/detection/d_fine.py

@@ -347,19 +347,13 @@ def _forward_explain_detection(
         backbone_feats = self.encoder(self.backbone(entity.images))
         predictions = self.decoder(backbone_feats, explain_mode=True)
 
-        feature_vector = self.feature_vector_fn(backbone_feats)
-
-        splits = [f.shape[-2] * f.shape[-1] for f in backbone_feats]
-
-        # Permute and split logits in one line
-        raw_logits = torch.split(predictions["raw_logits"].permute(0, 2, 1), splits, dim=-1)
-
-        # Reshape each split in a list comprehension
-        raw_logits = [
-            logits.reshape(f.shape[0], -1, f.shape[-2], f.shape[-1]) for logits, f in zip(raw_logits, backbone_feats)
-        ]
+        raw_logits = DETR.split_and_reshape_logits(
+            backbone_feats,
+            predictions["raw_logits"],
+        )
 
         saliency_map = self.explain_fn(raw_logits)
+        feature_vector = self.feature_vector_fn(backbone_feats)
         predictions.update(
             {
                 "feature_vector": feature_vector,

src/otx/algo/detection/detectors/detection_transformer.py

@@ -103,16 +103,8 @@ def export(
             deploy_mode=True,
         )
         if explain_mode:
+            raw_logits = self.split_and_reshape_logits(backbone_feats, predictions["raw_logits"])
             feature_vector = self.feature_vector_fn(backbone_feats)
-            splits = [f.shape[-2] * f.shape[-1] for f in backbone_feats]
-            # Permute and split logits in one line
-            raw_logits = torch.split(predictions["raw_logits"].permute(0, 2, 1), splits, dim=-1)
-
-            # Reshape each split in a list comprehension
-            raw_logits = [
-                logits.reshape(f.shape[0], -1, f.shape[-2], f.shape[-1])
-                for logits, f in zip(raw_logits, backbone_feats)
-            ]
             saliency_map = self.explain_fn(raw_logits)
             xai_output = {
                 "feature_vector": feature_vector,
@@ -121,6 +113,29 @@ def export(
             results.update(xai_output)  # type: ignore[union-attr]
         return results
 
+    @staticmethod
+    def split_and_reshape_logits(
+        backbone_feats: tuple[Tensor, ...],
+        raw_logits: Tensor,
+    ) -> tuple[Tensor, ...]:
+        """Splits and reshapes raw logits for explain mode.
+
+        Args:
+            backbone_feats (tuple[Tensor,...]): Tuple of backbone features.
+            raw_logits (Tensor): Raw logits.
+
+        Returns:
+            tuple[Tensor,...]: The reshaped logits.
+        """
+        splits = [f.shape[-2] * f.shape[-1] for f in backbone_feats]
+        # Permute and split logits in one line
+        raw_logits = torch.split(raw_logits.permute(0, 2, 1), splits, dim=-1)
+
+        # Reshape each split in a list comprehension
+        return tuple(
+            logits.reshape(f.shape[0], -1, f.shape[-2], f.shape[-1]) for logits, f in zip(raw_logits, backbone_feats)
+        )
+
     def postprocess(
         self,
         outputs: dict[str, Tensor],

src/otx/algo/detection/rtdetr.py

@@ -312,19 +312,13 @@ def _forward_explain_detection(
         backbone_feats = self.encoder(self.backbone(entity.images))
         predictions = self.decoder(backbone_feats, explain_mode=True)
 
-        feature_vector = self.feature_vector_fn(backbone_feats)
-
-        splits = [f.shape[-2] * f.shape[-1] for f in backbone_feats]
-
-        # Permute and split logits in one line
-        raw_logits = torch.split(predictions["raw_logits"].permute(0, 2, 1), splits, dim=-1)
-
-        # Reshape each split in a list comprehension
-        raw_logits = [
-            logits.reshape(f.shape[0], -1, f.shape[-2], f.shape[-1]) for logits, f in zip(raw_logits, backbone_feats)
-        ]
+        raw_logits = DETR.split_and_reshape_logits(
+            backbone_feats,
+            predictions["raw_logits"],
+        )
 
         saliency_map = self.explain_fn(raw_logits)
+        feature_vector = self.feature_vector_fn(backbone_feats)
         predictions.update(
             {
                 "feature_vector": feature_vector,