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

birder/net/detection/retinanet.py

@@ -120,7 +120,7 @@ class RetinaNetClassificationHead(nn.Module):
             cls_logits = self.cls_logits(cls_logits)
 
             # Permute classification output from (N, A * K, H, W) to (N, HWA, K).
-            (N, _, H, W) = cls_logits.shape
+            N, _, H, W = cls_logits.shape
             cls_logits = cls_logits.view(N, -1, self.num_classes, H, W)
             cls_logits = cls_logits.permute(0, 3, 4, 1, 2)
             cls_logits = cls_logits.reshape(N, -1, self.num_classes)  # Size=(N, HWA, K)
@@ -202,7 +202,7 @@ class RetinaNetRegressionHead(nn.Module):
             bbox_regression = self.bbox_reg(bbox_regression)
 
             # Permute bbox regression output from (N, 4 * A, H, W) to (N, HWA, 4).
-            (N, _, H, W) = bbox_regression.size()
+            N, _, H, W = bbox_regression.size()
             bbox_regression = bbox_regression.view(N, -1, 4, H, W)
             bbox_regression = bbox_regression.permute(0, 3, 4, 1, 2)
             bbox_regression = bbox_regression.reshape(N, -1, 4)  # Size=(N, HWA, 4)
@@ -395,7 +395,7 @@ class RetinaNet(DetectionBaseNet):
 
                 # Keep only topk scoring predictions
                 num_topk = min(self.topk_candidates, int(topk_idxs.size(0)))
-                (scores_per_level, idxs) = scores_per_level.topk(num_topk)
+                scores_per_level, idxs = scores_per_level.topk(num_topk)
                 topk_idxs = topk_idxs[idxs]
 
                 anchor_idxs = torch.div(topk_idxs, num_classes, rounding_mode="floor")
@@ -417,7 +417,7 @@ class RetinaNet(DetectionBaseNet):
 
             # 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)
+                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)

birder/net/detection/rt_detr_v1.py

@@ -1,5 +1,5 @@
 """
-RT-DETR (Real-Time DEtection TRansformer), adapted from
+RT-DETR v1 (Real-Time DEtection TRansformer), adapted from
 https://github.com/lyuwenyu/RT-DETR/tree/main/rtdetr_pytorch
 
 Paper "DETRs Beat YOLOs on Real-time Object Detection", https://arxiv.org/abs/2304.08069
@@ -114,7 +114,7 @@ def get_contrastive_denoising_training_group(  # pylint: disable=too-many-locals
     # Embed class labels
     input_query_class = class_embed(input_query_class)
 
-    # Create attention mask
+    # Create attention mask (True = masked)
     target_size = total_denoising_queries + num_queries
     attn_mask = torch.zeros([target_size, target_size], dtype=torch.bool, device=device)
     attn_mask[total_denoising_queries:, :total_denoising_queries] = True
@@ -212,10 +212,69 @@ class CSPRepLayer(nn.Module):
         return self.conv3(x1 + x2)
 
 
+class MultiheadAttention(nn.Module):
+    def __init__(self, d_model: int, num_heads: int, attn_drop: float = 0.0, proj_drop: float = 0.0) -> None:
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model should be divisible by num_heads"
+
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(d_model, d_model)
+        self.k_proj = nn.Linear(d_model, d_model)
+        self.v_proj = nn.Linear(d_model, d_model)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(d_model, d_model)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.xavier_uniform_(self.q_proj.weight)
+        nn.init.xavier_uniform_(self.k_proj.weight)
+        nn.init.xavier_uniform_(self.v_proj.weight)
+        nn.init.xavier_uniform_(self.proj.weight)
+        if self.q_proj.bias is not None:
+            nn.init.zeros_(self.q_proj.bias)
+            nn.init.zeros_(self.k_proj.bias)
+            nn.init.zeros_(self.v_proj.bias)
+            nn.init.zeros_(self.proj.bias)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, l_q, C = query.shape
+        q = self.q_proj(query).reshape(B, l_q, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(key).reshape(B, key.size(1), self.num_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(value).reshape(B, value.size(1), self.num_heads, self.head_dim).transpose(1, 2)
+
+        if key_padding_mask is not None:
+            # key_padding_mask is expected to be boolean (True = masked)
+            # SDPA expects True = attend, so we invert
+            attn_mask = ~key_padding_mask[:, None, None, :]
+        else:
+            attn_mask = None
+
+        attn = F.scaled_dot_product_attention(  # pylint: disable=not-callable
+            q, k, v, attn_mask=attn_mask, dropout_p=self.attn_drop.p if self.training else 0.0, scale=self.scale
+        )
+
+        attn = attn.transpose(1, 2).reshape(B, l_q, C)
+        x = self.proj(attn)
+        x = self.proj_drop(x)
+
+        return x
+
+
 class TransformerEncoderLayer(nn.Module):
     def __init__(self, d_model: int, num_heads: int, dim_feedforward: int, dropout: float) -> None:
         super().__init__()
-        self.self_attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
+        self.self_attn = MultiheadAttention(d_model, num_heads, attn_drop=dropout)
 
         self.linear1 = nn.Linear(d_model, dim_feedforward)
         self.dropout = nn.Dropout(dropout)
@@ -231,10 +290,8 @@ class TransformerEncoderLayer(nn.Module):
     def forward(
         self, src: torch.Tensor, pos: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None
     ) -> torch.Tensor:
-        q = src + pos
-        k = src + pos
-
-        (src2, _) = self.self_attn(q, k, value=src, key_padding_mask=key_padding_mask, need_weights=False)
+        qk = src + pos
+        src2 = self.self_attn(qk, qk, value=src, key_padding_mask=key_padding_mask)
         src = src + self.dropout1(src2)
         src = self.norm1(src)
 
@@ -268,7 +325,7 @@ class AIFI(nn.Module):
         self._pos_cache.clear()
 
     def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor:
-        (B, C, H, W) = x.size()
+        B, C, H, W = x.size()
         x = x.flatten(2).permute(0, 2, 1)
 
         use_cache = self.use_cache is True and torch.jit.is_tracing() is False and torch.jit.is_scripting() is False
@@ -522,7 +579,7 @@ class RT_DETRDecoder(nn.Module):
         spatial_shapes: list[list[int]],
         memory_padding_mask: Optional[torch.Tensor] = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        (anchors, valid_mask) = self._generate_anchors(spatial_shapes, device=memory.device, dtype=memory.dtype)
+        anchors, valid_mask = self._generate_anchors(spatial_shapes, device=memory.device, dtype=memory.dtype)
         if memory_padding_mask is not None:
             valid_mask = valid_mask & ~memory_padding_mask.unsqueeze(-1)
 
@@ -535,7 +592,7 @@ class RT_DETRDecoder(nn.Module):
         enc_outputs_coord_unact = self.enc_bbox_head(output_memory) + anchors
 
         # Select top-k queries based on classification confidence
-        (_, topk_ind) = torch.topk(enc_outputs_class.max(dim=-1).values, self.num_queries, dim=1)
+        _, topk_ind = torch.topk(enc_outputs_class.max(dim=-1).values, self.num_queries, dim=1)
 
         # Gather reference points
         reference_points_unact = enc_outputs_coord_unact.gather(
@@ -577,7 +634,7 @@ class RT_DETRDecoder(nn.Module):
         memory_padding_mask = torch.concat(mask_flatten, dim=1) if mask_flatten else None
 
         # Get decoder input (query selection)
-        (target, init_ref_points_unact, enc_topk_bboxes, enc_topk_logits) = self._get_decoder_input(
+        target, init_ref_points_unact, enc_topk_bboxes, enc_topk_logits = self._get_decoder_input(
             memory, spatial_shapes, memory_padding_mask
         )
 
@@ -858,7 +915,7 @@ class RT_DETR_v1(DetectionBaseNet):
             loss_ce = self._class_loss(
                 dn_out_logits[layer_idx], dn_out_bboxes[layer_idx], targets, indices, dn_num_boxes
             )
-            (loss_bbox, loss_giou) = self._box_loss(dn_out_bboxes[layer_idx], targets, indices, dn_num_boxes)
+            loss_bbox, loss_giou = self._box_loss(dn_out_bboxes[layer_idx], targets, indices, dn_num_boxes)
 
             loss_ce_list.append(loss_ce)
             loss_bbox_list.append(loss_bbox)
@@ -899,7 +956,7 @@ class RT_DETR_v1(DetectionBaseNet):
         for layer_idx in range(out_logits.shape[0]):
             indices = self.matcher(out_logits[layer_idx], out_bboxes[layer_idx], targets)
             loss_ce = self._class_loss(out_logits[layer_idx], out_bboxes[layer_idx], targets, indices, num_boxes)
-            (loss_bbox, loss_giou) = self._box_loss(out_bboxes[layer_idx], targets, indices, num_boxes)
+            loss_bbox, loss_giou = self._box_loss(out_bboxes[layer_idx], targets, indices, num_boxes)
             loss_ce_list.append(loss_ce)
             loss_bbox_list.append(loss_bbox)
             loss_giou_list.append(loss_giou)
@@ -907,7 +964,7 @@ class RT_DETR_v1(DetectionBaseNet):
         # Encoder auxiliary loss
         enc_indices = self.matcher(enc_topk_logits, enc_topk_bboxes, targets)
         loss_ce_enc = self._class_loss(enc_topk_logits, enc_topk_bboxes, targets, enc_indices, num_boxes)
-        (loss_bbox_enc, loss_giou_enc) = self._box_loss(enc_topk_bboxes, targets, enc_indices, num_boxes)
+        loss_bbox_enc, loss_giou_enc = self._box_loss(enc_topk_bboxes, targets, enc_indices, num_boxes)
         loss_ce_list.append(loss_ce_enc)
         loss_bbox_list.append(loss_bbox_enc)
         loss_giou_list.append(loss_giou_enc)
@@ -918,7 +975,7 @@ class RT_DETR_v1(DetectionBaseNet):
 
         # Add denoising loss if available
         if dn_out_bboxes is not None and dn_out_logits is not None and dn_meta is not None:
-            (loss_ce_dn, loss_bbox_dn, loss_giou_dn) = self._compute_denoising_loss(
+            loss_ce_dn, loss_bbox_dn, loss_giou_dn = self._compute_denoising_loss(
                 dn_out_bboxes, dn_out_logits, targets, dn_meta, num_boxes
             )
             loss_ce = loss_ce + loss_ce_dn
@@ -952,9 +1009,9 @@ class RT_DETR_v1(DetectionBaseNet):
             targets[idx]["boxes"] = boxes
             targets[idx]["labels"] = target["labels"] - 1  # No background
 
-        (denoising_class, denoising_bbox_unact, attn_mask, dn_meta) = self._prepare_cdn_queries(targets)
+        denoising_class, denoising_bbox_unact, attn_mask, dn_meta = self._prepare_cdn_queries(targets)
 
-        (out_bboxes, out_logits, enc_topk_bboxes, enc_topk_logits) = self.decoder(
+        out_bboxes, out_logits, enc_topk_bboxes, enc_topk_logits = self.decoder(
             encoder_features,
             spatial_shapes,
             level_start_index,
@@ -965,7 +1022,7 @@ class RT_DETR_v1(DetectionBaseNet):
         )
 
         if dn_meta is not None:
-            (dn_num_split, _num_queries) = dn_meta["dn_num_split"]
+            dn_num_split, _num_queries = dn_meta["dn_num_split"]
             dn_out_bboxes = out_bboxes[:, :, :dn_num_split]
             dn_out_logits = out_logits[:, :, :dn_num_split]
             out_bboxes = out_bboxes[:, :, dn_num_split:]
@@ -984,9 +1041,7 @@ class RT_DETR_v1(DetectionBaseNet):
         self, class_logits: torch.Tensor, box_regression: torch.Tensor, image_shapes: list[tuple[int, int]]
     ) -> 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
-        )
+        topk_values, topk_indexes = torch.topk(prob.view(class_logits.shape[0], -1), k=self.decoder.num_queries, dim=1)
         scores = topk_values
         topk_boxes = topk_indexes // class_logits.shape[2]
         labels = topk_indexes % class_logits.shape[2]
@@ -999,7 +1054,7 @@ class RT_DETR_v1(DetectionBaseNet):
         boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(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 = target_sizes.unbind(1)
         scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1)
         boxes = boxes * scale_fct[:, None, :]
 
@@ -1056,7 +1111,7 @@ class RT_DETR_v1(DetectionBaseNet):
                 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()
+                B, _, H, W = feat.size()
                 m = torch.zeros(B, H, W, dtype=torch.bool, device=x.device)
             mask_list.append(m)
 
@@ -1080,7 +1135,7 @@ class RT_DETR_v1(DetectionBaseNet):
             losses = self.compute_loss(encoder_features, spatial_shapes, level_start_index, targets, images, mask_list)
         else:
             # Inference path - no CDN
-            (out_bboxes, out_logits, _, _) = self.decoder(
+            out_bboxes, out_logits, _, _ = self.decoder(
                 encoder_features, spatial_shapes, level_start_index, padding_mask=mask_list
             )
             detections = self.postprocess_detections(out_logits[-1], out_bboxes[-1], images.image_sizes)
@@ -1100,6 +1155,7 @@ class RT_DETR_v1(DetectionBaseNet):
 
 
 registry.register_model_config(
-    "rt_detr_v1_s", RT_DETR_v1, config={"num_decoder_layers": 3, "expansion": 0.5, "depth_multiplier": 0.33}
+    "rt_detr_v1_t", RT_DETR_v1, config={"num_decoder_layers": 3, "expansion": 0.5, "depth_multiplier": 0.33}
 )
+registry.register_model_config("rt_detr_v1_s", RT_DETR_v1, config={"num_decoder_layers": 3, "expansion": 0.5})
 registry.register_model_config("rt_detr_v1", RT_DETR_v1, config={"num_decoder_layers": 6})