birder 0.4.1__py3-none-any.whl → 0.4.4__py3-none-any.whl

birder/net/detection/plain_detr.py

@@ -301,14 +301,11 @@ class GlobalDecoderLayer(nn.Module):
 
 
 class GlobalDecoder(nn.Module):
-    def __init__(
-        self, decoder_layer: nn.Module, num_layers: int, norm: nn.Module, return_intermediate: bool, d_model: int
-    ) -> None:
+    def __init__(self, decoder_layer: nn.Module, num_layers: int, norm: nn.Module, d_model: int) -> None:
         super().__init__()
         self.layers = _get_clones(decoder_layer, num_layers)
         self.num_layers = num_layers
         self.norm = norm
-        self.return_intermediate = return_intermediate
         self.d_model = d_model
 
         self.bbox_embed: Optional[nn.ModuleList] = None
@@ -339,6 +336,7 @@ class GlobalDecoder(nn.Module):
         reference_points: torch.Tensor,
         spatial_shape: tuple[int, int],
         memory_key_padding_mask: Optional[torch.Tensor] = None,
+        return_intermediates: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         output = tgt
         intermediate = []
@@ -364,14 +362,14 @@ class GlobalDecoder(nn.Module):
                 new_reference_points = new_reference_points.sigmoid()
                 reference_points = new_reference_points.detach()
 
-                if self.return_intermediate is True:
+                if return_intermediates is True:
                     intermediate.append(output_for_pred)
                     intermediate_reference_points.append(new_reference_points)
 
-            if self.return_intermediate is True:
+            if return_intermediates is True:
                 return torch.stack(intermediate), torch.stack(intermediate_reference_points)
 
-            return output_for_pred.unsqueeze(0), new_reference_points.unsqueeze(0)
+            return output_for_pred, new_reference_points
 
         for layer in self.layers:
             reference_points_input = reference_points.detach().clamp(0, 1)
@@ -388,14 +386,14 @@ class GlobalDecoder(nn.Module):
 
             output_for_pred = self.norm(output)
 
-            if self.return_intermediate is True:
+            if return_intermediates is True:
                 intermediate.append(output_for_pred)
                 intermediate_reference_points.append(reference_points)
 
-        if self.return_intermediate is True:
+        if return_intermediates is True:
             return torch.stack(intermediate), torch.stack(intermediate_reference_points)
 
-        return output_for_pred.unsqueeze(0), reference_points.unsqueeze(0)
+        return output_for_pred, reference_points
 
 
 class TransformerEncoderLayer(nn.Module):
@@ -467,7 +465,6 @@ class Plain_DETR(DetectionBaseNet):
         hidden_dim = 256
         num_heads = 8
         dropout = 0.0
-        return_intermediate = True
         dim_feedforward: int = self.config.get("dim_feedforward", 2048)
         num_encoder_layers: int = self.config["num_encoder_layers"]
         num_decoder_layers: int = self.config["num_decoder_layers"]
@@ -516,19 +513,18 @@ class Plain_DETR(DetectionBaseNet):
             decoder_layer,
             num_decoder_layers,
             decoder_norm,
-            return_intermediate=return_intermediate,
             d_model=hidden_dim,
         )
 
         self.class_embed = nn.Linear(hidden_dim, self.num_classes)
         self.bbox_embed = MLP(hidden_dim, [hidden_dim, hidden_dim, 4], activation_layer=nn.ReLU)
-        self.query_embed = nn.Embedding(self.num_queries, hidden_dim * 2)
+        self.query_embed = nn.Parameter(torch.empty(self.num_queries, hidden_dim * 2))
         self.reference_point_head = MLP(hidden_dim, [hidden_dim, hidden_dim, 4], activation_layer=nn.ReLU)
         self.input_proj = nn.Conv2d(
             self.backbone.return_channels[-1], hidden_dim, kernel_size=(1, 1), stride=(1, 1), padding=(0, 0)
         )
         self.pos_enc = PositionEmbeddingSine(hidden_dim // 2, normalize=True)
-        self.matcher = HungarianMatcher(cost_class=2, cost_bbox=5, cost_giou=2)
+        self.matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0)
 
         if box_refine is True:
             self.class_embed = _get_clones(self.class_embed, num_decoder_layers)
@@ -554,6 +550,7 @@ class Plain_DETR(DetectionBaseNet):
             if idx == 0:
                 nn.init.constant_(last_linear.bias[2:], -2.0)  # Small initial wh
 
+        nn.init.normal_(self.query_embed)
         ref_last_linear = [m for m in self.reference_point_head.modules() if isinstance(m, nn.Linear)][-1]
         nn.init.zeros_(ref_last_linear.weight)
         nn.init.zeros_(ref_last_linear.bias)
@@ -576,7 +573,8 @@ class Plain_DETR(DetectionBaseNet):
             for param in self.class_embed.parameters():
                 param.requires_grad_(True)
 
-    def _get_src_permutation_idx(self, indices: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    @staticmethod
+    def _get_src_permutation_idx(indices: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
         batch_idx = torch.concat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
         src_idx = torch.concat([src for (src, _) in indices])
         return (batch_idx, src_idx)
@@ -585,7 +583,7 @@ class Plain_DETR(DetectionBaseNet):
         self,
         cls_logits: torch.Tensor,
         targets: list[dict[str, torch.Tensor]],
-        indices: list[torch.Tensor],
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
         num_boxes: int,
     ) -> torch.Tensor:
         idx = self._get_src_permutation_idx(indices)
@@ -610,7 +608,7 @@ class Plain_DETR(DetectionBaseNet):
         self,
         box_output: torch.Tensor,
         targets: list[dict[str, torch.Tensor]],
-        indices: list[torch.Tensor],
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
         num_boxes: int,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         idx = self._get_src_permutation_idx(indices)
@@ -646,7 +644,7 @@ class Plain_DETR(DetectionBaseNet):
         if training_utils.is_dist_available_and_initialized() is True:
             torch.distributed.all_reduce(num_boxes)
 
-        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1).item()
+        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1)
 
         loss_ce_list = []
         loss_bbox_list = []
@@ -697,20 +695,17 @@ class Plain_DETR(DetectionBaseNet):
         return losses
 
     def postprocess_detections(
-        self, class_logits: torch.Tensor, box_regression: torch.Tensor, image_shapes: list[tuple[int, int]]
+        self, class_logits: torch.Tensor, box_regression: torch.Tensor, image_sizes: torch.Tensor
     ) -> list[dict[str, torch.Tensor]]:
         prob = class_logits.sigmoid()
         scores, labels = prob.max(-1)
         labels = labels + 1  # Background offset
 
-        # TorchScript doesn't support creating tensor from tuples, convert everything to lists
-        target_sizes = torch.tensor([list(s) for s in image_shapes], device=class_logits.device)
-
         # Convert to [x0, y0, x1, y1] format
         boxes = box_ops.box_convert(box_regression, in_fmt="cxcywh", out_fmt="xyxy")
 
         # Convert from relative [0, 1] to absolute [0, height] coordinates
-        img_h, img_w = target_sizes.unbind(1)
+        img_h, img_w = image_sizes.unbind(1)
         scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
         boxes = boxes * scale_fct[:, None, :]
 
@@ -735,17 +730,7 @@ class Plain_DETR(DetectionBaseNet):
 
         return detections
 
-    # pylint: disable=too-many-locals
-    def forward(
-        self,
-        x: torch.Tensor,
-        targets: Optional[list[dict[str, torch.Tensor]]] = None,
-        masks: Optional[torch.Tensor] = None,
-        image_sizes: Optional[list[list[int]]] = None,
-    ) -> tuple[list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]:
-        self._input_check(targets)
-        images = self._to_img_list(x, image_sizes)
-
+    def forward_net(self, x: torch.Tensor, masks: Optional[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
         features: dict[str, torch.Tensor] = self.backbone.detection_features(x)
         src = features[self.backbone.return_stages[-1]]
         src = self.input_proj(src)
@@ -772,7 +757,7 @@ class Plain_DETR(DetectionBaseNet):
         else:
             num_queries_to_use = self.num_queries_one2one
 
-        query_embed = self.query_embed.weight[:num_queries_to_use]
+        query_embed = self.query_embed[:num_queries_to_use]
         query_embed, query_pos = torch.split(query_embed, self.hidden_dim, dim=1)
         query_embed = query_embed.unsqueeze(0).expand(B, -1, -1)
         query_pos = query_pos.unsqueeze(0).expand(B, -1, -1)
@@ -787,25 +772,52 @@ class Plain_DETR(DetectionBaseNet):
             reference_points=reference_points,
             spatial_shape=(H, W),
             memory_key_padding_mask=mask_flatten,
+            return_intermediates=self.training is True,
         )
 
-        outputs_classes = []
-        outputs_coords = []
-        for lvl, (class_embed, bbox_embed) in enumerate(zip(self.class_embed, self.bbox_embed)):
-            outputs_class = class_embed(hs[lvl])
-            outputs_classes.append(outputs_class)
+        if self.training is True:
+            outputs_classes = []
+            outputs_coords = []
+            for lvl, (class_embed, bbox_embed) in enumerate(zip(self.class_embed, self.bbox_embed)):
+                outputs_class = class_embed(hs[lvl])
+                outputs_classes.append(outputs_class)
+
+                if self.box_refine is True:
+                    outputs_coord = inter_references[lvl]
+                else:
+                    tmp = bbox_embed(hs[lvl])
+                    tmp = tmp + inverse_sigmoid(reference_points)
+                    outputs_coord = tmp.sigmoid()
+
+                outputs_coords.append(outputs_coord)
+
+            outputs_class = torch.stack(outputs_classes)
+            outputs_coord = torch.stack(outputs_coords)
+        else:
+            class_embed = self.class_embed[-1]
+            bbox_embed = self.bbox_embed[-1]
+            outputs_class = class_embed(hs)
 
             if self.box_refine is True:
-                outputs_coord = inter_references[lvl]
+                outputs_coord = inter_references
             else:
-                tmp = bbox_embed(hs[lvl])
+                tmp = bbox_embed(hs)
                 tmp = tmp + inverse_sigmoid(reference_points)
                 outputs_coord = tmp.sigmoid()
 
-            outputs_coords.append(outputs_coord)
+        return (outputs_class, outputs_coord)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        targets: Optional[list[dict[str, torch.Tensor]]] = None,
+        masks: Optional[torch.Tensor] = None,
+        image_sizes: Optional[list[tuple[int, int]]] = None,
+    ) -> tuple[list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]:
+        self._input_check(targets)
+        images = self._to_img_list(x, image_sizes)
 
-        outputs_class = torch.stack(outputs_classes)
-        outputs_coord = torch.stack(outputs_coords)
+        outputs_class, outputs_coord = self.forward_net(x, masks)
 
         losses = {}
         detections: list[dict[str, torch.Tensor]] = []
@@ -815,7 +827,8 @@ class Plain_DETR(DetectionBaseNet):
             for idx, target in enumerate(targets):
                 boxes = target["boxes"]
                 boxes = box_ops.box_convert(boxes, in_fmt="xyxy", out_fmt="cxcywh")
-                boxes = boxes / torch.tensor(images.image_sizes[idx][::-1] * 2, dtype=torch.float32, device=x.device)
+                scale = images.image_sizes[idx].flip(0).repeat(2).float()  # flip to [W, H], repeat to [W, H, W, H]
+                boxes = boxes / scale
                 targets[idx]["boxes"] = boxes
                 targets[idx]["labels"] = target["labels"] - 1  # No background
 
@@ -835,7 +848,7 @@ class Plain_DETR(DetectionBaseNet):
             )
 
         else:
-            detections = self.postprocess_detections(outputs_class[-1], outputs_coord[-1], images.image_sizes)
+            detections = self.postprocess_detections(outputs_class, outputs_coord, images.image_sizes)
 
         return (detections, losses)
 

birder/net/detection/retinanet.py

@@ -30,6 +30,7 @@ from birder.net.detection.base import BackboneWithSimpleFPN
 from birder.net.detection.base import BoxCoder
 from birder.net.detection.base import DetectionBaseNet
 from birder.net.detection.base import Matcher
+from birder.net.detection.base import clip_boxes_to_image
 from birder.ops.soft_nms import SoftNMS
 
 
@@ -63,7 +64,7 @@ class RetinaNetClassificationHead(nn.Module):
             if isinstance(layer, nn.Conv2d):
                 nn.init.normal_(layer.weight, std=0.01)
                 if layer.bias is not None:
-                    nn.init.constant_(layer.bias, 0)
+                    nn.init.zeros_(layer.bias)
 
         self.cls_logits = nn.Conv2d(in_channels, num_anchors * num_classes, kernel_size=3, stride=1, padding=1)
 
@@ -281,6 +282,11 @@ class RetinaNet(DetectionBaseNet):
         if soft_nms is True:
             self.soft_nms = SoftNMS()
 
+        self.score_thresh = score_thresh
+        self.nms_thresh = nms_thresh
+        self.detections_per_img = detections_per_img
+        self.topk_candidates = topk_candidates
+
         if feature_pyramid_type == "fpn":
             feature_pyramid: Callable[..., nn.Module] = BackboneWithFPN
             num_anchor_sizes = len(self.backbone.return_stages) + 2
@@ -314,10 +320,8 @@ class RetinaNet(DetectionBaseNet):
         self.proposal_matcher = Matcher(fg_iou_thresh, bg_iou_thresh, allow_low_quality_matches=True)
         self.box_coder = BoxCoder(weights=(1.0, 1.0, 1.0, 1.0))
 
-        self.score_thresh = score_thresh
-        self.nms_thresh = nms_thresh
-        self.detections_per_img = detections_per_img
-        self.topk_candidates = topk_candidates
+        if self.export_mode is False:
+            self.forward = torch.compiler.disable(recursive=False)(self.forward)  # type: ignore[method-assign]
 
     def reset_classifier(self, num_classes: int) -> None:
         self.num_classes = num_classes
@@ -341,10 +345,7 @@ class RetinaNet(DetectionBaseNet):
     @torch.jit.unused  # type: ignore[untyped-decorator]
     @torch.compiler.disable()  # type: ignore[untyped-decorator]
     def compute_loss(
-        self,
-        targets: list[dict[str, torch.Tensor]],
-        head_outputs: dict[str, torch.Tensor],
-        anchors: list[torch.Tensor],
+        self, targets: list[dict[str, torch.Tensor]], head_outputs: dict[str, torch.Tensor], anchors: list[torch.Tensor]
     ) -> dict[str, torch.Tensor]:
         matched_idxs = []
         for idx, (anchors_per_image, targets_per_image) in enumerate(zip(anchors, targets)):
@@ -362,22 +363,19 @@ class RetinaNet(DetectionBaseNet):
 
     # pylint: disable=too-many-locals
     def postprocess_detections(
-        self,
-        head_outputs: dict[str, list[torch.Tensor]],
-        anchors: list[list[torch.Tensor]],
-        image_shapes: list[tuple[int, int]],
+        self, head_outputs: dict[str, list[torch.Tensor]], anchors: list[list[torch.Tensor]], image_sizes: torch.Tensor
     ) -> list[dict[str, torch.Tensor]]:
         class_logits = head_outputs["cls_logits"]
         box_regression = head_outputs["bbox_regression"]
 
-        num_images = len(image_shapes)
+        num_images = image_sizes.size(0)
 
         detections: list[dict[str, torch.Tensor]] = []
         for index in range(num_images):
             box_regression_per_image = [br[index] for br in box_regression]
             logits_per_image = [cl[index] for cl in class_logits]
             anchors_per_image = anchors[index]
-            image_shape = image_shapes[index]
+            image_shape = image_sizes[index]
 
             image_boxes_list = []
             image_scores_list = []
@@ -394,7 +392,7 @@ class RetinaNet(DetectionBaseNet):
                 topk_idxs = torch.where(keep_idxs)[0]
 
                 # Keep only topk scoring predictions
-                num_topk = min(self.topk_candidates, int(topk_idxs.size(0)))
+                num_topk = min(self.topk_candidates, topk_idxs.size(0))
                 scores_per_level, idxs = scores_per_level.topk(num_topk)
                 topk_idxs = topk_idxs[idxs]
 
@@ -405,7 +403,7 @@ class RetinaNet(DetectionBaseNet):
                 boxes_per_level = self.box_coder.decode_single(
                     box_regression_per_level[anchor_idxs], anchors_per_level[anchor_idxs]
                 )
-                boxes_per_level = box_ops.clip_boxes_to_image(boxes_per_level, image_shape)
+                boxes_per_level = clip_boxes_to_image(boxes_per_level, image_shape)
 
                 image_boxes_list.append(boxes_per_level)
                 image_scores_list.append(scores_per_level)
@@ -415,24 +413,40 @@ class RetinaNet(DetectionBaseNet):
             image_scores = torch.concat(image_scores_list, dim=0)
             image_labels = torch.concat(image_labels_list, dim=0)
 
-            # Non-maximum suppression
-            if self.soft_nms is not None:
-                soft_scores, keep = self.soft_nms(image_boxes, image_scores, image_labels, score_threshold=0.001)
-                image_scores[keep] = soft_scores
+            if self.export_mode is False:
+                # Non-maximum suppression
+                if self.soft_nms is not None:
+                    soft_scores, keep = self.soft_nms(image_boxes, image_scores, image_labels, score_threshold=0.001)
+                    image_scores[keep] = soft_scores
+                else:
+                    keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh)
+
+                keep = keep[: self.detections_per_img]
+
+                detections.append(
+                    {
+                        "boxes": image_boxes[keep],
+                        "scores": image_scores[keep],
+                        "labels": image_labels[keep],
+                    }
+                )
             else:
-                keep = box_ops.batched_nms(image_boxes, image_scores, image_labels, self.nms_thresh)
+                detections.append(
+                    {
+                        "boxes": image_boxes,
+                        "scores": image_scores,
+                        "labels": image_labels,
+                    }
+                )
 
-            keep = keep[: self.detections_per_img]
+        return detections
 
-            detections.append(
-                {
-                    "boxes": image_boxes[keep],
-                    "scores": image_scores[keep],
-                    "labels": image_labels[keep],
-                }
-            )
+    def forward_net(self, x: torch.Tensor) -> tuple[list[torch.Tensor], dict[str, torch.Tensor]]:
+        features: dict[str, torch.Tensor] = self.backbone_with_fpn(x)
+        feature_list = list(features.values())
+        head_outputs = self.head(feature_list)
 
-        return detections
+        return (feature_list, head_outputs)
 
     # pylint: disable=invalid-name
     def forward(
@@ -440,14 +454,12 @@ class RetinaNet(DetectionBaseNet):
         x: torch.Tensor,
         targets: Optional[list[dict[str, torch.Tensor]]] = None,
         masks: Optional[torch.Tensor] = None,
-        image_sizes: Optional[list[list[int]]] = None,
+        image_sizes: Optional[list[tuple[int, int]]] = None,
     ) -> tuple[list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]:
         self._input_check(targets)
         images = self._to_img_list(x, image_sizes)
 
-        features: dict[str, torch.Tensor] = self.backbone_with_fpn(x)
-        feature_list = list(features.values())
-        head_outputs = self.head(feature_list)
+        feature_list, head_outputs = self.forward_net(x)
         anchors = self.anchor_generator(images, feature_list)
 
         losses: dict[str, torch.Tensor] = {}

birder/net/detection/rt_detr_v1.py

@@ -47,9 +47,6 @@ def get_contrastive_denoising_training_group(  # pylint: disable=too-many-locals
     label_noise_ratio: float,
     box_noise_scale: float,
 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[dict[str, Any]]]:
-    if num_denoising_queries <= 0:
-        return (None, None, None, None)
-
     num_ground_truths = [len(t["labels"]) for t in targets]
     device = targets[0]["labels"].device
 
@@ -596,18 +593,18 @@ class RT_DETRDecoder(nn.Module):
 
         # Gather reference points
         reference_points_unact = enc_outputs_coord_unact.gather(
-            dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_coord_unact.shape[-1])
+            dim=1, index=topk_ind.unsqueeze(-1).expand(-1, -1, enc_outputs_coord_unact.shape[-1])
         )
 
         enc_topk_bboxes = reference_points_unact.sigmoid()
 
         # Gather encoder logits for loss computation
         enc_topk_logits = enc_outputs_class.gather(
-            dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, enc_outputs_class.shape[-1])
+            dim=1, index=topk_ind.unsqueeze(-1).expand(-1, -1, enc_outputs_class.shape[-1])
         )
 
         # Extract region features
-        target = output_memory.gather(dim=1, index=topk_ind.unsqueeze(-1).repeat(1, 1, output_memory.shape[-1]))
+        target = output_memory.gather(dim=1, index=topk_ind.unsqueeze(-1).expand(-1, -1, output_memory.shape[-1]))
         target = target.detach()
 
         return (target, reference_points_unact.detach(), enc_topk_bboxes, enc_topk_logits)
@@ -621,6 +618,7 @@ class RT_DETRDecoder(nn.Module):
         denoising_bbox_unact: Optional[torch.Tensor] = None,
         attn_mask: Optional[torch.Tensor] = None,
         padding_mask: Optional[list[torch.Tensor]] = None,
+        return_intermediates: bool = True,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
         memory = []
         mask_flatten = []
@@ -648,12 +646,12 @@ class RT_DETRDecoder(nn.Module):
         level_start_index_tensor = torch.tensor(level_start_index, dtype=torch.long, device=memory.device)
 
         # Decoder forward
-        out_bboxes = []
-        out_logits = []
+        bboxes_list: list[torch.Tensor] = []
+        logits_list: list[torch.Tensor] = []
         reference_points = init_ref_points_unact.sigmoid()
         for decoder_layer, bbox_head, class_head in zip(self.layers, self.bbox_embed, self.class_embed):
             query_pos = self.query_pos_head(reference_points)
-            reference_points_input = reference_points.unsqueeze(2).repeat(1, 1, len(spatial_shapes), 1)
+            reference_points_input = reference_points.unsqueeze(2).expand(-1, -1, len(spatial_shapes), -1)
             target = decoder_layer(
                 target,
                 query_pos,
@@ -663,6 +661,7 @@ class RT_DETRDecoder(nn.Module):
                 level_start_index_tensor,
                 memory_padding_mask,
                 attn_mask,
+                src_shapes=spatial_shapes,
             )
 
             bbox_delta = bbox_head(target)
@@ -672,14 +671,19 @@ class RT_DETRDecoder(nn.Module):
             # Classification
             class_logits = class_head(target)
 
-            out_bboxes.append(new_reference_points)
-            out_logits.append(class_logits)
+            if return_intermediates is True:
+                bboxes_list.append(new_reference_points)
+                logits_list.append(class_logits)
 
             # Update reference points for next layer
             reference_points = new_reference_points.detach()
 
-        out_bboxes = torch.stack(out_bboxes)
-        out_logits = torch.stack(out_logits)
+        if return_intermediates is True:
+            out_bboxes = torch.stack(bboxes_list)
+            out_logits = torch.stack(logits_list)
+        else:
+            out_bboxes = new_reference_points
+            out_logits = class_logits
 
         return (out_bboxes, out_logits, enc_topk_bboxes, enc_topk_logits)
 
@@ -743,7 +747,7 @@ class RT_DETR_v1(DetectionBaseNet):
         self.decoder = RT_DETRDecoder(
             hidden_dim=hidden_dim,
             num_classes=self.num_classes,
-            num_queries=num_queries,
+            num_queries=self.num_queries,
             num_decoder_layers=num_decoder_layers,
             num_levels=self.num_levels,
             num_heads=num_heads,
@@ -810,7 +814,8 @@ class RT_DETR_v1(DetectionBaseNet):
                 for param in self.denoising_class_embed.parameters():
                     param.requires_grad_(True)
 
-    def _get_src_permutation_idx(self, indices: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
+    @staticmethod
+    def _get_src_permutation_idx(indices: list[tuple[torch.Tensor, torch.Tensor]]) -> tuple[torch.Tensor, torch.Tensor]:
         batch_idx = torch.concat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
         src_idx = torch.concat([src for (src, _) in indices])
         return (batch_idx, src_idx)
@@ -820,7 +825,7 @@ class RT_DETR_v1(DetectionBaseNet):
         cls_logits: torch.Tensor,
         box_output: torch.Tensor,
         targets: list[dict[str, torch.Tensor]],
-        indices: list[torch.Tensor],
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
         num_boxes: float,
     ) -> torch.Tensor:
         idx = self._get_src_permutation_idx(indices)
@@ -859,7 +864,7 @@ class RT_DETR_v1(DetectionBaseNet):
         self,
         box_output: torch.Tensor,
         targets: list[dict[str, torch.Tensor]],
-        indices: list[torch.Tensor],
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
         num_boxes: float,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         idx = self._get_src_permutation_idx(indices)
@@ -927,8 +932,6 @@ class RT_DETR_v1(DetectionBaseNet):
 
         return (loss_ce_dn, loss_bbox_dn, loss_giou_dn)
 
-    @torch.jit.unused  # type: ignore[untyped-decorator]
-    @torch.compiler.disable()  # type: ignore[untyped-decorator]
     def _compute_loss_from_outputs(  # pylint: disable=too-many-locals
         self,
         targets: list[dict[str, torch.Tensor]],
@@ -946,7 +949,7 @@ class RT_DETR_v1(DetectionBaseNet):
         if training_utils.is_dist_available_and_initialized() is True:
             torch.distributed.all_reduce(num_boxes)
 
-        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1).item()
+        num_boxes = torch.clamp(num_boxes / training_utils.get_world_size(), min=1)
 
         loss_ce_list = []
         loss_bbox_list = []
@@ -1001,11 +1004,11 @@ class RT_DETR_v1(DetectionBaseNet):
         images: Any,
         masks: Optional[list[torch.Tensor]] = None,
     ) -> dict[str, torch.Tensor]:
-        device = encoder_features[0].device
         for idx, target in enumerate(targets):
             boxes = target["boxes"]
             boxes = box_ops.box_convert(boxes, in_fmt="xyxy", out_fmt="cxcywh")
-            boxes = boxes / torch.tensor(images.image_sizes[idx][::-1] * 2, dtype=torch.float32, device=device)
+            scale = images.image_sizes[idx].flip(0).repeat(2).float()  # flip to [W, H], repeat to [W, H, W, H]
+            boxes = boxes / scale
             targets[idx]["boxes"] = boxes
             targets[idx]["labels"] = target["labels"] - 1  # No background
 
@@ -1038,7 +1041,7 @@ class RT_DETR_v1(DetectionBaseNet):
         return losses
 
     def postprocess_detections(
-        self, class_logits: torch.Tensor, box_regression: torch.Tensor, image_shapes: list[tuple[int, int]]
+        self, class_logits: torch.Tensor, box_regression: torch.Tensor, image_sizes: torch.Tensor
     ) -> list[dict[str, torch.Tensor]]:
         prob = class_logits.sigmoid()
         topk_values, topk_indexes = torch.topk(prob.view(class_logits.shape[0], -1), k=self.decoder.num_queries, dim=1)
@@ -1047,14 +1050,12 @@ class RT_DETR_v1(DetectionBaseNet):
         labels = topk_indexes % class_logits.shape[2]
         labels += 1  # Background offset
 
-        target_sizes = torch.tensor(image_shapes, device=class_logits.device)
-
         # Convert to [x0, y0, x1, y1] format
         boxes = box_ops.box_convert(box_regression, in_fmt="cxcywh", out_fmt="xyxy")
-        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4))
+        boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).expand(-1, -1, 4))
 
         # Convert from relative [0, 1] to absolute [0, height] coordinates
-        img_h, img_w = target_sizes.unbind(1)
+        img_h, img_w = image_sizes.unbind(1)
         scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
         boxes = boxes * scale_fct[:, None, :]
 
@@ -1090,32 +1091,34 @@ class RT_DETR_v1(DetectionBaseNet):
 
         return (None, None, None, None)
 
+    def forward_net(
+        self, x: torch.Tensor, masks: Optional[torch.Tensor]
+    ) -> tuple[list[torch.Tensor], Optional[list[torch.Tensor]]]:
+        features: dict[str, torch.Tensor] = self.backbone.detection_features(x)
+        feature_list = list(features.values())
+
+        mask_list: Optional[list[torch.Tensor]] = None
+        if masks is not None:
+            mask_list = []
+            for feat in feature_list:
+                m = F.interpolate(masks[None].float(), size=feat.shape[-2:], mode="nearest").to(torch.bool)[0]
+                mask_list.append(m)
+
+        encoder_features = self.encoder(feature_list, masks=mask_list)
+
+        return (encoder_features, mask_list)
+
     def forward(
         self,
         x: torch.Tensor,
         targets: Optional[list[dict[str, torch.Tensor]]] = None,
         masks: Optional[torch.Tensor] = None,
-        image_sizes: Optional[list[list[int]]] = None,
+        image_sizes: Optional[list[tuple[int, int]]] = None,
     ) -> tuple[list[dict[str, torch.Tensor]], dict[str, torch.Tensor]]:
         self._input_check(targets)
         images = self._to_img_list(x, image_sizes)
 
-        # Backbone features
-        features: dict[str, torch.Tensor] = self.backbone.detection_features(x)
-        feature_list = list(features.values())
-
-        # Hybrid encoder
-        mask_list: list[torch.Tensor] = []
-        for feat in feature_list:
-            if masks is not None:
-                mask_size = feat.shape[-2:]
-                m = F.interpolate(masks[None].float(), size=mask_size, mode="nearest").to(torch.bool)[0]
-            else:
-                B, _, H, W = feat.size()
-                m = torch.zeros(B, H, W, dtype=torch.bool, device=x.device)
-            mask_list.append(m)
-
-        encoder_features = self.encoder(feature_list, masks=mask_list)
+        encoder_features, mask_list = self.forward_net(x, masks)
 
         # Prepare spatial shapes and level start index
         spatial_shapes: list[list[int]] = []
@@ -1136,9 +1139,9 @@ class RT_DETR_v1(DetectionBaseNet):
         else:
             # Inference path - no CDN
             out_bboxes, out_logits, _, _ = self.decoder(
-                encoder_features, spatial_shapes, level_start_index, padding_mask=mask_list
+                encoder_features, spatial_shapes, level_start_index, padding_mask=mask_list, return_intermediates=False
             )
-            detections = self.postprocess_detections(out_logits[-1], out_bboxes[-1], images.image_sizes)
+            detections = self.postprocess_detections(out_logits, out_bboxes, images.image_sizes)
 
         return (detections, losses)