From 3d4f4c6958505a297d5a096a3738e739abcadde3 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 18:32:02 +0000 Subject: [PATCH 01/10] Initial plan From 5800205ec8529f332bf5efd800dc27334b6eeba1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 18:41:01 +0000 Subject: [PATCH 02/10] Add DINOv2ViT and DINOv2ConvNext feature extractors with RetinaNet detectors --- detectors.py | 107 +++++++++++++++++++ feature_extractors.py | 236 ++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 3 files changed, 344 insertions(+) create mode 100644 feature_extractors.py diff --git a/detectors.py b/detectors.py index 81e86d7..b772f0c 100644 --- a/detectors.py +++ b/detectors.py @@ -7,6 +7,7 @@ from datetime import datetime import logger import inference +from feature_extractors import DINOv2ViTBackbone, DINOv2ConvNextBackbone class AbstractDetector(): @@ -567,3 +568,109 @@ def replace_head(self, model, num_classes: int): num_anchors=num_anchors, num_classes=num_classes+1, ) + + +class DINOv2ViTRetinaNetDetector(AbstractDetector): + """RetinaNet detector with a DINOv2 ViT feature extractor backbone. + + The DINOv2 ViT last-layer features are projected and downsampled into a + four-level feature pyramid, which is consumed by a RetinaNet head. + + Args: + num_classes (int | None): Number of foreground classes. When + provided the detection head is replaced to match. + resume (str | None): Path to a checkpoint to resume from. + device (str): ``"cpu"`` or a CUDA device string. + root_dir (str | None): Root directory used for logging. + finetuning (bool): If ``True`` the DINOv2 backbone weights are + updated during training. Defaults to ``False`` (frozen). + """ + + def __init__( + self, + num_classes: int = None, + resume: str = None, + device: str = "cpu", + root_dir: str = None, + finetuning: bool = False, + ): + self.finetuning = finetuning + super().__init__( + name="dinov2_vit_retinanet", + num_classes=num_classes, + resume=resume, + device=device, + root_dir=root_dir, + ) + + def get_loss_names(self) -> list[str]: + return ["classification", "bbox_regression"] + + def load_pretrained_model(self): + backbone = DINOv2ViTBackbone(out_channels=256, finetuning=self.finetuning) + return torchvision.models.detection.RetinaNet( + backbone=backbone, + num_classes=91, + ) + + def replace_head(self, model, num_classes: int): + in_channels = model.head.classification_head.cls_logits.in_channels + num_anchors = model.head.classification_head.num_anchors + model.head = torchvision.models.detection.retinanet.RetinaNetHead( + in_channels=in_channels, + num_anchors=num_anchors, + num_classes=num_classes + 1, + ) + + +class DINOv2ConvNextRetinaNetDetector(AbstractDetector): + """RetinaNet detector with a ConvNext feature extractor backbone. + + All four ConvNext stages are passed through an FPN to produce a + five-level (4 + max-pool) feature pyramid that feeds the RetinaNet head. + + Args: + num_classes (int | None): Number of foreground classes. When + provided the detection head is replaced to match. + resume (str | None): Path to a checkpoint to resume from. + device (str): ``"cpu"`` or a CUDA device string. + root_dir (str | None): Root directory used for logging. + finetuning (bool): If ``True`` the ConvNext backbone weights are + updated during training. Defaults to ``False`` (frozen). + """ + + def __init__( + self, + num_classes: int = None, + resume: str = None, + device: str = "cpu", + root_dir: str = None, + finetuning: bool = False, + ): + self.finetuning = finetuning + super().__init__( + name="dinov2_convnext_retinanet", + num_classes=num_classes, + resume=resume, + device=device, + root_dir=root_dir, + ) + + def get_loss_names(self) -> list[str]: + return ["classification", "bbox_regression"] + + def load_pretrained_model(self): + backbone = DINOv2ConvNextBackbone(out_channels=256, finetuning=self.finetuning) + return torchvision.models.detection.RetinaNet( + backbone=backbone, + num_classes=91, + ) + + def replace_head(self, model, num_classes: int): + in_channels = model.head.classification_head.cls_logits.in_channels + num_anchors = model.head.classification_head.num_anchors + model.head = torchvision.models.detection.retinanet.RetinaNetHead( + in_channels=in_channels, + num_anchors=num_anchors, + num_classes=num_classes + 1, + ) diff --git a/feature_extractors.py b/feature_extractors.py new file mode 100644 index 0000000..82fa813 --- /dev/null +++ b/feature_extractors.py @@ -0,0 +1,236 @@ +import torch +import torch.nn as nn +from collections import OrderedDict +from transformers import Dinov2Config, Dinov2Model, ConvNextConfig, ConvNextModel +from torchvision.ops import FeaturePyramidNetwork +from torchvision.ops.feature_pyramid_network import LastLevelMaxPool + + +class DINOv2ViT(nn.Module): + """DINOv2 Vision Transformer feature extractor. + + Extracts features from intermediate transformer layers and reshapes them + from (B, N, D) to spatial feature maps (B, D, H, W). + + Args: + finetuning (bool): If True the backbone weights are trainable. + output_patches (bool): If True include the raw patch embeddings as an + additional output level. + layers (list[int]): Indices of transformer encoder layers whose output + to capture. Default: [2, 5, 8, 11]. + layer_norm (bool): Apply a learned LayerNorm to each captured feature + before reshaping. + """ + + def __init__( + self, + finetuning: bool = False, + output_patches: bool = False, + layers: list = None, + layer_norm: bool = True, + ): + super().__init__() + if layers is None: + layers = [2, 5, 8, 11] + self.finetuning = finetuning + self.layers = layers + self.output_patches = output_patches + self.config = Dinov2Config() + self.model = Dinov2Model(self.config) + self.layer_norm = layer_norm and (len(layers) + int(output_patches)) > 0 + if self.layer_norm: + self.norms = nn.ModuleList([ + nn.LayerNorm(self.config.hidden_size, eps=1e-5, elementwise_affine=True) + for _ in range(len(layers) + int(output_patches)) + ]) + if not self.finetuning: + self._freeze() + + def _freeze(self): + for param in self.model.parameters(): + param.requires_grad = False + + def forward(self, x: torch.Tensor): + """Forward image through DINOv2 ViT. + + Args: + x (torch.Tensor): Image tensor of shape (B, C, H, W). + + Returns: + torch.Tensor or tuple[torch.Tensor]: Spatial feature map(s) of + shape (B, D, H/patch, W/patch). A tuple is returned when + more than one layer is selected. + """ + captured = OrderedDict() + hooks = [] + + if self.output_patches: + hooks.append( + self.model.embeddings.register_forward_hook( + lambda m, i, o, key="embeddings": captured.update({key: o}) + ) + ) + for layer_idx in self.layers: + hooks.append( + self.model.encoder.layer[layer_idx].register_forward_hook( + lambda m, i, o, key=layer_idx: captured.update({key: o}) + ) + ) + + z = self.model(x) + + for h in hooks: + h.remove() + + if not captured: + captured[len(self.model.encoder.layer) - 1] = z.last_hidden_state + + feature_maps = OrderedDict() + for idx, (k, feat) in enumerate(captured.items()): + if self.layer_norm: + feat = self.norms[idx](feat) + # Remove the [CLS] token + feat = feat[:, 1:, :] + B, P, D = feat.shape + h = w = int(P ** 0.5) + feat = feat.permute(0, 2, 1).reshape(B, D, h, w) + feature_maps[k] = feat + + if len(feature_maps) > 1: + return tuple(feature_maps[k] for k in feature_maps) + return feature_maps[next(iter(feature_maps))] + + +class DINOv2ConvNext(nn.Module): + """ConvNext feature extractor for use in detection pipelines. + + Extracts hierarchical feature maps from a ConvNext backbone. Each + selected stage outputs a spatial map at a different resolution, making + this extractor naturally suitable for FPN-based detectors. + + Args: + finetuning (bool): If True the backbone weights are trainable. + layers (list[int]): Indices of ConvNext stages to capture. + Default: [0, 1, 2, 3]. + """ + + def __init__( + self, + finetuning: bool = False, + layers: list = None, + ): + super().__init__() + if layers is None: + layers = [0, 1, 2, 3] + self.finetuning = finetuning + self.layers = layers + self.config = ConvNextConfig() + self.model = ConvNextModel(self.config) + if not self.finetuning: + self._freeze() + + def _freeze(self): + for param in self.model.parameters(): + param.requires_grad = False + + def forward(self, x: torch.Tensor): + """Forward image through ConvNext. + + Args: + x (torch.Tensor): Image tensor of shape (B, C, H, W). + + Returns: + torch.Tensor or tuple[torch.Tensor]: Feature map(s) from the + selected stages. A tuple is returned when more than one + stage is selected. + """ + captured = OrderedDict() + hooks = [] + + for i in self.layers: + hooks.append( + self.model.encoder.stages[i].register_forward_hook( + lambda m, inp, out, key=i: captured.update({key: out}) + ) + ) + + z = self.model(x) + + for h in hooks: + h.remove() + + if not captured: + captured[0] = z.last_hidden_state + + if len(captured) > 1: + return tuple(captured[k] for k in captured) + return captured[next(iter(captured))] + + +class DINOv2ViTBackbone(nn.Module): + """DINOv2 ViT backbone with multi-scale projection for detection. + + Wraps :class:`DINOv2ViT` (single-scale ViT output) into a five-level + feature pyramid compatible with torchvision detection heads such as + :class:`~torchvision.models.detection.RetinaNet`. + + The ViT last-layer features are first projected to ``out_channels`` via a + 1×1 convolution, then downsampled four times with 3×3 strided convolutions + to produce five scales. + + Args: + out_channels (int): Number of channels in every output feature level. + finetuning (bool): Passed to :class:`DINOv2ViT`. + """ + + def __init__(self, out_channels: int = 256, finetuning: bool = False): + super().__init__() + self.body = DINOv2ViT(finetuning=finetuning, layers=[], layer_norm=False) + hidden = self.body.config.hidden_size + + self.proj = nn.Conv2d(hidden, out_channels, kernel_size=1) + self.down1 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1) + self.down2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1) + self.down3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1) + self.down4 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1) + self.out_channels = out_channels + + def forward(self, x: torch.Tensor) -> OrderedDict: + feat = self.body(x) + p0 = self.proj(feat) + p1 = self.down1(p0) + p2 = self.down2(p1) + p3 = self.down3(p2) + p4 = self.down4(p3) + return OrderedDict([("0", p0), ("1", p1), ("2", p2), ("3", p3), ("4", p4)]) + + +class DINOv2ConvNextBackbone(nn.Module): + """ConvNext backbone with FPN for detection. + + Wraps :class:`DINOv2ConvNext` and applies a + :class:`~torchvision.ops.FeaturePyramidNetwork` across all four ConvNext + stages to produce a uniform ``out_channels``-wide feature pyramid. + + Args: + out_channels (int): Number of channels in every FPN output level. + finetuning (bool): Passed to :class:`DINOv2ConvNext`. + """ + + def __init__(self, out_channels: int = 256, finetuning: bool = False): + super().__init__() + self.body = DINOv2ConvNext(finetuning=finetuning, layers=[0, 1, 2, 3]) + in_channels_list = list(self.body.config.hidden_sizes) + self.fpn = FeaturePyramidNetwork( + in_channels_list=in_channels_list, + out_channels=out_channels, + extra_blocks=LastLevelMaxPool(), + ) + self.out_channels = out_channels + + def forward(self, x: torch.Tensor) -> OrderedDict: + stage_feats = self.body(x) + feat_dict = OrderedDict( + (str(i), f) for i, f in enumerate(stage_feats) + ) + return self.fpn(feat_dict) diff --git a/requirements.txt b/requirements.txt index f370a15..cae4967 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ torch==2.12.1 torchvision==0.27.1 pandas==3.0.3 Pillow==12.2.0 +transformers>=4.38.0 From 5d17c11db728878db2c1fe6e6a2bf0d7121dca59 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 18:44:25 +0000 Subject: [PATCH 03/10] Address review feedback: add model_name param for pretrained weights, fix FPN key mapping --- detectors.py | 27 ++++++++++++++--- feature_extractors.py | 70 +++++++++++++++++++++++++++++++++++-------- 2 files changed, 81 insertions(+), 16 deletions(-) diff --git a/detectors.py b/detectors.py index b772f0c..91ae629 100644 --- a/detectors.py +++ b/detectors.py @@ -574,7 +574,7 @@ class DINOv2ViTRetinaNetDetector(AbstractDetector): """RetinaNet detector with a DINOv2 ViT feature extractor backbone. The DINOv2 ViT last-layer features are projected and downsampled into a - four-level feature pyramid, which is consumed by a RetinaNet head. + five-level feature pyramid, which is consumed by a RetinaNet head. Args: num_classes (int | None): Number of foreground classes. When @@ -584,6 +584,9 @@ class DINOv2ViTRetinaNetDetector(AbstractDetector): root_dir (str | None): Root directory used for logging. finetuning (bool): If ``True`` the DINOv2 backbone weights are updated during training. Defaults to ``False`` (frozen). + model_name (str | None): HuggingFace model identifier for pretrained + DINOv2 weights (e.g. ``'facebook/dinov2-base'``). When ``None`` + the backbone is randomly initialised. """ def __init__( @@ -593,8 +596,10 @@ def __init__( device: str = "cpu", root_dir: str = None, finetuning: bool = False, + model_name: str = None, ): self.finetuning = finetuning + self.model_name = model_name super().__init__( name="dinov2_vit_retinanet", num_classes=num_classes, @@ -607,7 +612,11 @@ def get_loss_names(self) -> list[str]: return ["classification", "bbox_regression"] def load_pretrained_model(self): - backbone = DINOv2ViTBackbone(out_channels=256, finetuning=self.finetuning) + backbone = DINOv2ViTBackbone( + model_name=self.model_name, + out_channels=256, + finetuning=self.finetuning, + ) return torchvision.models.detection.RetinaNet( backbone=backbone, num_classes=91, @@ -627,7 +636,8 @@ class DINOv2ConvNextRetinaNetDetector(AbstractDetector): """RetinaNet detector with a ConvNext feature extractor backbone. All four ConvNext stages are passed through an FPN to produce a - five-level (4 + max-pool) feature pyramid that feeds the RetinaNet head. + five-level (4 FPN + max-pool) feature pyramid that feeds the RetinaNet + head. Args: num_classes (int | None): Number of foreground classes. When @@ -637,6 +647,9 @@ class DINOv2ConvNextRetinaNetDetector(AbstractDetector): root_dir (str | None): Root directory used for logging. finetuning (bool): If ``True`` the ConvNext backbone weights are updated during training. Defaults to ``False`` (frozen). + model_name (str | None): HuggingFace model identifier for pretrained + ConvNext weights (e.g. ``'facebook/convnext-base-224'``). When + ``None`` the backbone is randomly initialised. """ def __init__( @@ -646,8 +659,10 @@ def __init__( device: str = "cpu", root_dir: str = None, finetuning: bool = False, + model_name: str = None, ): self.finetuning = finetuning + self.model_name = model_name super().__init__( name="dinov2_convnext_retinanet", num_classes=num_classes, @@ -660,7 +675,11 @@ def get_loss_names(self) -> list[str]: return ["classification", "bbox_regression"] def load_pretrained_model(self): - backbone = DINOv2ConvNextBackbone(out_channels=256, finetuning=self.finetuning) + backbone = DINOv2ConvNextBackbone( + model_name=self.model_name, + out_channels=256, + finetuning=self.finetuning, + ) return torchvision.models.detection.RetinaNet( backbone=backbone, num_classes=91, diff --git a/feature_extractors.py b/feature_extractors.py index 82fa813..49c1290 100644 --- a/feature_extractors.py +++ b/feature_extractors.py @@ -13,6 +13,10 @@ class DINOv2ViT(nn.Module): from (B, N, D) to spatial feature maps (B, D, H, W). Args: + model_name (str | None): HuggingFace model identifier to load + pretrained weights (e.g. ``'facebook/dinov2-base'``). When + ``None`` a randomly-initialised model using the default + :class:`~transformers.Dinov2Config` is created. finetuning (bool): If True the backbone weights are trainable. output_patches (bool): If True include the raw patch embeddings as an additional output level. @@ -24,6 +28,7 @@ class DINOv2ViT(nn.Module): def __init__( self, + model_name: str = None, finetuning: bool = False, output_patches: bool = False, layers: list = None, @@ -35,8 +40,14 @@ def __init__( self.finetuning = finetuning self.layers = layers self.output_patches = output_patches - self.config = Dinov2Config() - self.model = Dinov2Model(self.config) + + if model_name is not None: + self.model = Dinov2Model.from_pretrained(model_name) + self.config = self.model.config + else: + self.config = Dinov2Config() + self.model = Dinov2Model(self.config) + self.layer_norm = layer_norm and (len(layers) + int(output_patches)) > 0 if self.layer_norm: self.norms = nn.ModuleList([ @@ -65,12 +76,14 @@ def forward(self, x: torch.Tensor): hooks = [] if self.output_patches: + # key= default-argument captures current string at definition time hooks.append( self.model.embeddings.register_forward_hook( lambda m, i, o, key="embeddings": captured.update({key: o}) ) ) for layer_idx in self.layers: + # key=layer_idx default-argument captures the current int at definition time hooks.append( self.model.encoder.layer[layer_idx].register_forward_hook( lambda m, i, o, key=layer_idx: captured.update({key: o}) @@ -109,6 +122,10 @@ class DINOv2ConvNext(nn.Module): this extractor naturally suitable for FPN-based detectors. Args: + model_name (str | None): HuggingFace model identifier to load + pretrained weights (e.g. ``'facebook/convnext-base-224'``). When + ``None`` a randomly-initialised model using the default + :class:`~transformers.ConvNextConfig` is created. finetuning (bool): If True the backbone weights are trainable. layers (list[int]): Indices of ConvNext stages to capture. Default: [0, 1, 2, 3]. @@ -116,6 +133,7 @@ class DINOv2ConvNext(nn.Module): def __init__( self, + model_name: str = None, finetuning: bool = False, layers: list = None, ): @@ -124,8 +142,14 @@ def __init__( layers = [0, 1, 2, 3] self.finetuning = finetuning self.layers = layers - self.config = ConvNextConfig() - self.model = ConvNextModel(self.config) + + if model_name is not None: + self.model = ConvNextModel.from_pretrained(model_name) + self.config = self.model.config + else: + self.config = ConvNextConfig() + self.model = ConvNextModel(self.config) + if not self.finetuning: self._freeze() @@ -148,6 +172,7 @@ def forward(self, x: torch.Tensor): hooks = [] for i in self.layers: + # key=i default-argument captures the current int at definition time hooks.append( self.model.encoder.stages[i].register_forward_hook( lambda m, inp, out, key=i: captured.update({key: out}) @@ -179,13 +204,21 @@ class DINOv2ViTBackbone(nn.Module): to produce five scales. Args: + model_name (str | None): Passed to :class:`DINOv2ViT`. Provide a + HuggingFace identifier (e.g. ``'facebook/dinov2-base'``) to load + pretrained weights. out_channels (int): Number of channels in every output feature level. finetuning (bool): Passed to :class:`DINOv2ViT`. """ - def __init__(self, out_channels: int = 256, finetuning: bool = False): + def __init__( + self, + model_name: str = None, + out_channels: int = 256, + finetuning: bool = False, + ): super().__init__() - self.body = DINOv2ViT(finetuning=finetuning, layers=[], layer_norm=False) + self.body = DINOv2ViT(model_name=model_name, finetuning=finetuning, layers=[], layer_norm=False) hidden = self.body.config.hidden_size self.proj = nn.Conv2d(hidden, out_channels, kernel_size=1) @@ -208,18 +241,27 @@ def forward(self, x: torch.Tensor) -> OrderedDict: class DINOv2ConvNextBackbone(nn.Module): """ConvNext backbone with FPN for detection. - Wraps :class:`DINOv2ConvNext` and applies a - :class:`~torchvision.ops.FeaturePyramidNetwork` across all four ConvNext - stages to produce a uniform ``out_channels``-wide feature pyramid. + Wraps :class:`DINOv2ConvNext` (all four stages) and applies a + :class:`~torchvision.ops.FeaturePyramidNetwork` to produce a uniform + ``out_channels``-wide five-level feature pyramid (four FPN levels plus one + max-pool level). Args: + model_name (str | None): Passed to :class:`DINOv2ConvNext`. Provide a + HuggingFace identifier (e.g. ``'facebook/convnext-base-224'``) to + load pretrained weights. out_channels (int): Number of channels in every FPN output level. finetuning (bool): Passed to :class:`DINOv2ConvNext`. """ - def __init__(self, out_channels: int = 256, finetuning: bool = False): + def __init__( + self, + model_name: str = None, + out_channels: int = 256, + finetuning: bool = False, + ): super().__init__() - self.body = DINOv2ConvNext(finetuning=finetuning, layers=[0, 1, 2, 3]) + self.body = DINOv2ConvNext(model_name=model_name, finetuning=finetuning, layers=[0, 1, 2, 3]) in_channels_list = list(self.body.config.hidden_sizes) self.fpn = FeaturePyramidNetwork( in_channels_list=in_channels_list, @@ -230,7 +272,11 @@ def __init__(self, out_channels: int = 256, finetuning: bool = False): def forward(self, x: torch.Tensor) -> OrderedDict: stage_feats = self.body(x) + # Map each feature to its actual stage index so the FPN in_channels + # list aligns correctly with the captured feature maps. feat_dict = OrderedDict( - (str(i), f) for i, f in enumerate(stage_feats) + (str(stage_idx), f) + for stage_idx, f in zip(self.body.layers, stage_feats) ) return self.fpn(feat_dict) + From debe067752400c9ec19d8602ec438e795839ebeb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 18:48:37 +0000 Subject: [PATCH 04/10] Fix patch reshape to support rectangular images and improve validation message --- feature_extractors.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/feature_extractors.py b/feature_extractors.py index 49c1290..ed3bf17 100644 --- a/feature_extractors.py +++ b/feature_extractors.py @@ -48,6 +48,8 @@ def __init__( self.config = Dinov2Config() self.model = Dinov2Model(self.config) + # Only allocate LayerNorm modules when there is at least one feature + # level to normalise (either a captured encoder layer or patch embeddings). self.layer_norm = layer_norm and (len(layers) + int(output_patches)) > 0 if self.layer_norm: self.norms = nn.ModuleList([ @@ -98,6 +100,12 @@ def forward(self, x: torch.Tensor): if not captured: captured[len(self.model.encoder.layer) - 1] = z.last_hidden_state + # Spatial grid dimensions derived from the input, not from patch count, + # so that rectangular images are handled correctly. + patch_size = self.config.patch_size + feat_h = x.shape[2] // patch_size + feat_w = x.shape[3] // patch_size + feature_maps = OrderedDict() for idx, (k, feat) in enumerate(captured.items()): if self.layer_norm: @@ -105,8 +113,13 @@ def forward(self, x: torch.Tensor): # Remove the [CLS] token feat = feat[:, 1:, :] B, P, D = feat.shape - h = w = int(P ** 0.5) - feat = feat.permute(0, 2, 1).reshape(B, D, h, w) + if feat_h * feat_w != P: + raise ValueError( + f"Derived spatial grid ({feat_h}×{feat_w}={feat_h * feat_w}) does not " + f"match number of patch tokens ({P}). Ensure the input image dimensions " + f"are divisible by the model's patch size ({patch_size})." + ) + feat = feat.permute(0, 2, 1).reshape(B, D, feat_h, feat_w) feature_maps[k] = feat if len(feature_maps) > 1: From 846d982956548c3affcdbffd855c753a85ac473f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 19:08:24 +0000 Subject: [PATCH 05/10] Add DINOv2HungarianDetector with Hungarian matching detection head --- =1.11.0 | 0 detectors.py | 101 +++++++- hungarian_head.py | 592 ++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 4 files changed, 693 insertions(+), 1 deletion(-) create mode 100644 =1.11.0 create mode 100644 hungarian_head.py diff --git a/=1.11.0 b/=1.11.0 new file mode 100644 index 0000000..e69de29 diff --git a/detectors.py b/detectors.py index 91ae629..259803a 100644 --- a/detectors.py +++ b/detectors.py @@ -7,7 +7,8 @@ from datetime import datetime import logger import inference -from feature_extractors import DINOv2ViTBackbone, DINOv2ConvNextBackbone +from feature_extractors import DINOv2ViT, DINOv2ViTBackbone, DINOv2ConvNextBackbone +from hungarian_head import HungarianDetectionHead, DINOv2HungarianDetectionModel class AbstractDetector(): @@ -693,3 +694,101 @@ def replace_head(self, model, num_classes: int): num_anchors=num_anchors, num_classes=num_classes + 1, ) + + +class DINOv2HungarianDetector(AbstractDetector): + """DINOv2 detector with a DETR-style Hungarian-matching detection head. + + Combines a frozen (or optionally fine-tuned) :class:`DINOv2ViT` backbone + with a :class:`~hungarian_head.HungarianDetectionHead` that uses the + Hungarian algorithm to assign predicted queries to ground-truth objects + during training. + + The model is trained end-to-end with a set-based loss consisting of three + terms: cross-entropy classification loss, L1 bounding-box regression loss, + and GIoU bounding-box regression loss. + + Args: + num_classes (int | None): Number of foreground object classes. + resume (str | None): Path to a checkpoint file to load weights from. + device (str): ``'cpu'`` or a CUDA device string (e.g. ``'cuda:0'``). + root_dir (str | None): Root directory for logging. + finetuning (bool): If ``True`` the DINOv2 backbone weights are updated + during training. Defaults to ``False`` (backbone frozen). + model_name (str | None): HuggingFace model identifier for pretrained + DINOv2 weights (e.g. ``'facebook/dinov2-base'``). When ``None`` + the backbone is randomly initialised. + hidden_dim (int): Transformer / embedding dimensionality. Default: 256. + num_queries (int): Number of learnable object queries. Default: 100. + nhead (int): Number of attention heads per transformer decoder layer. + Default: 8. + num_decoder_layers (int): Depth of the transformer decoder. + Default: 6. + score_threshold (float): Minimum foreground score to keep during + inference. Default: 0.5. + """ + + def __init__( + self, + num_classes: int = None, + resume: str = None, + device: str = "cpu", + root_dir: str = None, + finetuning: bool = False, + model_name: str = None, + hidden_dim: int = 256, + num_queries: int = 100, + nhead: int = 8, + num_decoder_layers: int = 6, + score_threshold: float = 0.5, + ): + # Set attributes before calling super().__init__ because + # AbstractDetector.__init__ immediately calls build_model → load_pretrained_model. + self.finetuning = finetuning + self.model_name = model_name + self.hidden_dim = hidden_dim + self.num_queries = num_queries + self.nhead = nhead + self.num_decoder_layers = num_decoder_layers + self.score_threshold = score_threshold + super().__init__( + name="dinov2_hungarian", + num_classes=num_classes, + resume=resume, + device=device, + root_dir=root_dir, + ) + + def get_loss_names(self) -> list[str]: + return ["loss_classification", "loss_bbox", "loss_giou"] + + def load_pretrained_model(self): + backbone = DINOv2ViT( + model_name=self.model_name, + finetuning=self.finetuning, + layers=[], + layer_norm=False, + ) + in_channels = backbone.config.hidden_size + head = HungarianDetectionHead( + in_channels=in_channels, + num_classes=91, # COCO default; replaced by replace_head when num_classes is set + hidden_dim=self.hidden_dim, + num_queries=self.num_queries, + nhead=self.nhead, + num_decoder_layers=self.num_decoder_layers, + score_threshold=self.score_threshold, + ) + return DINOv2HungarianDetectionModel(backbone=backbone, head=head) + + def replace_head(self, model, num_classes: int): + in_channels = model.backbone.config.hidden_size + model.head = HungarianDetectionHead( + in_channels=in_channels, + num_classes=num_classes, + hidden_dim=self.hidden_dim, + num_queries=self.num_queries, + nhead=self.nhead, + num_decoder_layers=self.num_decoder_layers, + score_threshold=self.score_threshold, + ) diff --git a/hungarian_head.py b/hungarian_head.py new file mode 100644 index 0000000..b8fc38b --- /dev/null +++ b/hungarian_head.py @@ -0,0 +1,592 @@ +"""DETR-style detection head with Hungarian-matching set-based loss. + +This module provides: + +* :class:`SinusoidalPositionEncoding2D` – 2-D sinusoidal position encoding for + spatial feature maps. +* :class:`HungarianMatcher` – optimal bipartite matching between predicted and + ground-truth objects using the Hungarian algorithm + (:func:`scipy.optimize.linear_sum_assignment`). +* :class:`HungarianDetectionHead` – transformer-decoder detection head that + combines the matcher with a set-based training loss (classification, + L1-box, and GIoU). +* :class:`DINOv2HungarianDetectionModel` – complete ``nn.Module`` that pairs a + :class:`~feature_extractors.DINOv2ViT` backbone with a + :class:`HungarianDetectionHead`. +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F +from scipy.optimize import linear_sum_assignment +from torchvision.ops import box_convert, generalized_box_iou + +from feature_extractors import DINOv2ViT + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _build_mlp(input_dim: int, hidden_dim: int, output_dim: int, num_layers: int) -> nn.Sequential: + """Build a fully-connected MLP with ReLU activations between layers. + + Args: + input_dim (int): Dimensionality of the input features. + hidden_dim (int): Width of every hidden layer. + output_dim (int): Dimensionality of the final output. + num_layers (int): Total number of linear layers (including the output + layer). Must be ≥ 1. + + Returns: + nn.Sequential: The constructed MLP. + """ + layers: list[nn.Module] = [] + for i in range(num_layers): + in_d = input_dim if i == 0 else hidden_dim + out_d = output_dim if i == num_layers - 1 else hidden_dim + layers.append(nn.Linear(in_d, out_d)) + if i < num_layers - 1: + layers.append(nn.ReLU(inplace=True)) + return nn.Sequential(*layers) + + +# --------------------------------------------------------------------------- +# Positional encoding +# --------------------------------------------------------------------------- + +class SinusoidalPositionEncoding2D(nn.Module): + """2-D sinusoidal positional encoding for spatial feature maps. + + Generates a fixed (non-learned) positional encoding following the DETR + convention: the spatial axes are encoded independently with sine/cosine + functions, then concatenated along the channel dimension. + + Args: + hidden_dim (int): Total number of encoding channels (must be even). + temperature (float): Denominator base used in the sinusoidal + frequency computation. Default: 10 000. + """ + + def __init__(self, hidden_dim: int, temperature: float = 10000.0): + super().__init__() + if hidden_dim % 2 != 0: + raise ValueError("hidden_dim must be even for 2-D sinusoidal encoding.") + self.hidden_dim = hidden_dim + self.temperature = temperature + + def forward(self, mask: torch.Tensor) -> torch.Tensor: + """Compute the positional encoding for a spatial mask. + + Args: + mask (torch.Tensor): Boolean tensor of shape ``(B, H, W)`` where + ``True`` marks padding positions. + + Returns: + torch.Tensor: Positional encoding of shape ``(B, H*W, hidden_dim)``. + """ + B, H, W = mask.shape + device = mask.device + + not_mask = ~mask + # Cumulative sum gives a monotonically increasing coordinate per row/col + y_embed = not_mask.cumsum(1, dtype=torch.float32) # (B, H, W) + x_embed = not_mask.cumsum(2, dtype=torch.float32) # (B, H, W) + # Normalise to [0, 2π] + y_embed = y_embed / (y_embed[:, -1:, :] + 1e-6) * 2.0 * math.pi + x_embed = x_embed / (x_embed[:, :, -1:] + 1e-6) * 2.0 * math.pi + + half = self.hidden_dim // 2 + dim_t = torch.arange(half, device=device, dtype=torch.float32) + dim_t = self.temperature ** (2.0 * (dim_t // 2) / half) + + pos_x = x_embed[:, :, :, None] / dim_t # (B, H, W, half) + pos_y = y_embed[:, :, :, None] / dim_t + + # Interleave sin / cos + pos_x = torch.stack( + [pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()], dim=4 + ).flatten(3) # (B, H, W, half) + pos_y = torch.stack( + [pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()], dim=4 + ).flatten(3) # (B, H, W, half) + + pos = torch.cat([pos_y, pos_x], dim=3) # (B, H, W, hidden_dim) + return pos.flatten(1, 2) # (B, H*W, hidden_dim) + + +# --------------------------------------------------------------------------- +# Hungarian matcher +# --------------------------------------------------------------------------- + +class HungarianMatcher(nn.Module): + """Optimal bipartite matching of predictions to ground-truth objects. + + The matching cost is a weighted sum of three terms: + + * **Classification cost** – negative softmax probability assigned to the + target class. + * **L1 bounding-box cost** – ℓ¹ distance between predicted and target boxes + in normalised *cxcywh* format. + * **GIoU cost** – negative Generalised IoU between the matched boxes. + + Args: + weight_class (float): Weight for the classification cost. Default: 1. + weight_bbox (float): Weight for the L1 bounding-box cost. Default: 5. + weight_giou (float): Weight for the GIoU cost. Default: 2. + """ + + def __init__( + self, + weight_class: float = 1.0, + weight_bbox: float = 5.0, + weight_giou: float = 2.0, + ): + super().__init__() + self.weight_class = weight_class + self.weight_bbox = weight_bbox + self.weight_giou = weight_giou + + @torch.no_grad() + def forward( + self, + pred_logits: torch.Tensor, + pred_boxes: torch.Tensor, + targets: list[dict], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + """Compute the optimal matching for a batch of images. + + Args: + pred_logits (torch.Tensor): Raw class logits of shape + ``(B, num_queries, num_classes + 1)``. + pred_boxes (torch.Tensor): Predicted boxes in normalised *cxcywh* + format of shape ``(B, num_queries, 4)``. + targets (list[dict]): One dict per image. Each dict must contain: + + * ``'boxes'`` – ``(num_gt, 4)`` absolute *xyxy* coordinates. + * ``'labels'`` – ``(num_gt,)`` integer class labels + (1-indexed; 0 is reserved for background). + * ``'image_size'`` – ``(H, W)`` tuple of the input image + spatial dimensions. + + Returns: + list[tuple[Tensor, Tensor]]: For each image a pair + ``(src_idx, tgt_idx)`` of 1-D :class:`torch.long` tensors giving + the matched query indices and the matched ground-truth indices + respectively. + """ + B = pred_logits.shape[0] + indices = [] + + for b in range(B): + gt_boxes = targets[b]["boxes"] # (num_gt, 4) absolute xyxy + gt_labels = targets[b]["labels"] # (num_gt,) 1-indexed + img_h, img_w = targets[b]["image_size"] + num_gt = len(gt_labels) + + if num_gt == 0: + indices.append(( + torch.zeros(0, dtype=torch.long), + torch.zeros(0, dtype=torch.long), + )) + continue + + # ---- Classification cost ---------------------------------------- + prob = pred_logits[b].softmax(-1) # (Q, C+1) + cost_class = -prob[:, gt_labels] # (Q, num_gt) + + # ---- Normalise GT boxes to [0, 1] cxcywh ----------------------- + gt_boxes_norm = gt_boxes.clone().float() + gt_boxes_norm[:, [0, 2]] = gt_boxes_norm[:, [0, 2]] / img_w + gt_boxes_norm[:, [1, 3]] = gt_boxes_norm[:, [1, 3]] / img_h + gt_cxcywh = box_convert(gt_boxes_norm, in_fmt="xyxy", out_fmt="cxcywh") + + # ---- L1 bounding-box cost ---------------------------------------- + cost_bbox = torch.cdist(pred_boxes[b], gt_cxcywh, p=1) # (Q, num_gt) + + # ---- GIoU cost -------------------------------------------------- + pred_xyxy = box_convert(pred_boxes[b], in_fmt="cxcywh", out_fmt="xyxy") + gt_xyxy = box_convert(gt_cxcywh, in_fmt="cxcywh", out_fmt="xyxy") + cost_giou = -generalized_box_iou(pred_xyxy, gt_xyxy) # (Q, num_gt) + + # ---- Solve assignment problem ----------------------------------- + C = ( + self.weight_class * cost_class + + self.weight_bbox * cost_bbox + + self.weight_giou * cost_giou + ).cpu().numpy() + + row_idx, col_idx = linear_sum_assignment(C) + indices.append(( + torch.as_tensor(row_idx, dtype=torch.long), + torch.as_tensor(col_idx, dtype=torch.long), + )) + + return indices + + +# --------------------------------------------------------------------------- +# Detection head +# --------------------------------------------------------------------------- + +class HungarianDetectionHead(nn.Module): + """DETR-style detection head that uses Hungarian matching for supervision. + + The head accepts a single spatial feature map from a DINOv2 backbone, + projects it to ``hidden_dim`` channels, adds 2-D sinusoidal positional + encodings, and feeds the result as the *memory* of a standard + :class:`torch.nn.TransformerDecoder`. A set of ``num_queries`` learnable + object queries act as the *target* sequence. The decoded queries are + independently passed through a classification head and a bounding-box + regression MLP. + + **Training** – ground-truth objects are assigned to predicted queries via + the :class:`HungarianMatcher`. Three losses are computed: + + * Cross-entropy classification loss over all queries (matched queries carry + their GT class; unmatched queries carry the background class 0). + * L1 loss on the normalised *cxcywh* boxes of the matched pairs. + * GIoU loss on the matched pairs. + + **Inference** – scores are derived from the softmax over foreground classes + (background excluded). Only predictions whose score exceeds + ``score_threshold`` are returned. + + Args: + in_channels (int): Number of channels in the input feature map produced + by the backbone. + num_classes (int): Number of foreground classes. The head produces + ``num_classes + 1`` class outputs (index 0 = background). + hidden_dim (int): Transformer / embedding dimensionality. Default: 256. + num_queries (int): Number of learnable object queries. Default: 100. + nhead (int): Number of attention heads per transformer layer. + Default: 8. + num_decoder_layers (int): Depth of the transformer decoder. + Default: 6. + dim_feedforward (int): Feedforward dimension inside each transformer + layer. Default: 2048. + dropout (float): Dropout probability in the transformer. Default: 0.1. + weight_class (float): Matcher cost weight for classification. + Default: 1. + weight_bbox (float): Matcher cost weight for L1 boxes. Default: 5. + weight_giou (float): Matcher cost weight for GIoU. Default: 2. + loss_weight_class (float): Coefficient for the classification loss term. + Default: 1. + loss_weight_bbox (float): Coefficient for the L1 box loss term. + Default: 5. + loss_weight_giou (float): Coefficient for the GIoU loss term. + Default: 2. + score_threshold (float): Minimum foreground score to keep a prediction + during inference. Default: 0.5. + """ + + def __init__( + self, + in_channels: int, + num_classes: int, + hidden_dim: int = 256, + num_queries: int = 100, + nhead: int = 8, + num_decoder_layers: int = 6, + dim_feedforward: int = 2048, + dropout: float = 0.1, + weight_class: float = 1.0, + weight_bbox: float = 5.0, + weight_giou: float = 2.0, + loss_weight_class: float = 1.0, + loss_weight_bbox: float = 5.0, + loss_weight_giou: float = 2.0, + score_threshold: float = 0.5, + ): + super().__init__() + self.num_classes = num_classes + self.num_queries = num_queries + self.score_threshold = score_threshold + self.loss_weight_class = loss_weight_class + self.loss_weight_bbox = loss_weight_bbox + self.loss_weight_giou = loss_weight_giou + + # Project backbone features to hidden_dim + self.input_proj = nn.Conv2d(in_channels, hidden_dim, kernel_size=1) + + # Positional encoding for memory tokens + self.pos_encoding = SinusoidalPositionEncoding2D(hidden_dim) + + # Learnable object queries + self.query_embed = nn.Embedding(num_queries, hidden_dim) + + # Transformer decoder + decoder_layer = nn.TransformerDecoderLayer( + d_model=hidden_dim, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + batch_first=True, + ) + self.transformer_decoder = nn.TransformerDecoder( + decoder_layer=decoder_layer, + num_layers=num_decoder_layers, + ) + + # Prediction heads + self.class_head = nn.Linear(hidden_dim, num_classes + 1) + self.box_head = _build_mlp(hidden_dim, hidden_dim, 4, num_layers=3) + + # Matcher (used only during training) + self.matcher = HungarianMatcher( + weight_class=weight_class, + weight_bbox=weight_bbox, + weight_giou=weight_giou, + ) + + def forward( + self, + features: torch.Tensor, + targets: list[dict] | None = None, + image_sizes: list[tuple[int, int]] | None = None, + ): + """Run the detection head on a batch of feature maps. + + Args: + features (torch.Tensor): Backbone feature map of shape + ``(B, in_channels, H, W)``. + targets (list[dict] | None): Ground-truth dicts (only needed during + training). Each dict must have ``'boxes'``, ``'labels'``, and + ``'image_size'`` keys – see :class:`HungarianMatcher` for the + expected formats. + image_sizes (list[tuple[int, int]] | None): ``(H, W)`` of each + original input image. Required during inference to convert + normalised predictions back to absolute pixel coordinates. + + Returns: + dict | list[dict]: + * **Training** – a dict with keys ``'loss_classification'``, + ``'loss_bbox'``, and ``'loss_giou'``. + * **Inference** – a list of dicts each containing ``'boxes'`` + (absolute *xyxy*), ``'labels'`` (1-indexed), and ``'scores'``. + """ + B, _, H, W = features.shape + + # Project features and add positional encoding + feat = self.input_proj(features) # (B, D, H, W) + mask = torch.zeros(B, H, W, dtype=torch.bool, device=features.device) + pos = self.pos_encoding(mask) # (B, H*W, D) + memory = feat.flatten(2).permute(0, 2, 1) + pos # (B, H*W, D) + + # Expand learnable queries for the whole batch + queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1) # (B, Q, D) + + # Decode + decoded = self.transformer_decoder(tgt=queries, memory=memory) # (B, Q, D) + + # Predictions + pred_logits = self.class_head(decoded) # (B, Q, C+1) + pred_boxes = self.box_head(decoded).sigmoid() # (B, Q, 4) normalised cxcywh + + if self.training: + return self._compute_loss(pred_logits, pred_boxes, targets) + return self._post_process(pred_logits, pred_boxes, image_sizes) + + # ------------------------------------------------------------------ + # Training helpers + # ------------------------------------------------------------------ + + def _compute_loss( + self, + pred_logits: torch.Tensor, + pred_boxes: torch.Tensor, + targets: list[dict], + ) -> dict[str, torch.Tensor]: + """Compute the set-based detection loss for a batch. + + Args: + pred_logits: ``(B, Q, C+1)`` raw class logits. + pred_boxes: ``(B, Q, 4)`` normalised *cxcywh* predictions. + targets: list of ground-truth dicts (see :meth:`forward`). + + Returns: + dict with keys ``'loss_classification'``, ``'loss_bbox'``, + ``'loss_giou'``. + """ + B, Q, _ = pred_logits.shape + device = pred_logits.device + + # Compute optimal matching + indices = self.matcher(pred_logits, pred_boxes, targets) + + # ---- Classification loss ------------------------------------------- + # Default: all queries predict background (class 0) + target_labels = torch.zeros(B, Q, dtype=torch.long, device=device) + for b, (src_idx, tgt_idx) in enumerate(indices): + if len(src_idx) > 0: + target_labels[b, src_idx] = targets[b]["labels"][tgt_idx].to(device) + + # Down-weight background to balance against foreground queries + num_gt_total = max(1, sum(len(t["labels"]) for t in targets)) + bg_weight = num_gt_total / (B * Q) + class_weights = torch.ones(self.num_classes + 1, device=device) + class_weights[0] = bg_weight + + loss_cls = F.cross_entropy( + pred_logits.reshape(B * Q, -1), + target_labels.reshape(B * Q), + weight=class_weights, + ) + + # ---- Box losses (matched pairs only) -------------------------------- + src_boxes_list: list[torch.Tensor] = [] + tgt_boxes_list: list[torch.Tensor] = [] + + for b, (src_idx, tgt_idx) in enumerate(indices): + if len(src_idx) == 0: + continue + src_boxes_list.append(pred_boxes[b][src_idx]) + + img_h, img_w = targets[b]["image_size"] + gt_boxes = targets[b]["boxes"][tgt_idx].float().to(device) + gt_boxes_norm = gt_boxes.clone() + gt_boxes_norm[:, [0, 2]] = gt_boxes_norm[:, [0, 2]] / img_w + gt_boxes_norm[:, [1, 3]] = gt_boxes_norm[:, [1, 3]] / img_h + tgt_boxes_list.append( + box_convert(gt_boxes_norm, in_fmt="xyxy", out_fmt="cxcywh") + ) + + if src_boxes_list: + src_boxes = torch.cat(src_boxes_list, dim=0) # (M, 4) + tgt_boxes = torch.cat(tgt_boxes_list, dim=0) # (M, 4) + num_matched = max(1, src_boxes.shape[0]) + + loss_bbox = F.l1_loss(src_boxes, tgt_boxes, reduction="sum") / num_matched + + src_xyxy = box_convert(src_boxes, in_fmt="cxcywh", out_fmt="xyxy") + tgt_xyxy = box_convert(tgt_boxes, in_fmt="cxcywh", out_fmt="xyxy") + giou = generalized_box_iou(src_xyxy, tgt_xyxy) + loss_giou = (1.0 - giou.diag()).sum() / num_matched + else: + # No matched pairs in this batch – return zero-gradient tensors + loss_bbox = pred_boxes.sum() * 0.0 + loss_giou = pred_boxes.sum() * 0.0 + + return { + "loss_classification": self.loss_weight_class * loss_cls, + "loss_bbox": self.loss_weight_bbox * loss_bbox, + "loss_giou": self.loss_weight_giou * loss_giou, + } + + # ------------------------------------------------------------------ + # Inference helpers + # ------------------------------------------------------------------ + + def _post_process( + self, + pred_logits: torch.Tensor, + pred_boxes: torch.Tensor, + image_sizes: list[tuple[int, int]], + ) -> list[dict[str, torch.Tensor]]: + """Convert raw model output to detection results. + + Args: + pred_logits: ``(B, Q, C+1)`` raw class logits. + pred_boxes: ``(B, Q, 4)`` normalised *cxcywh* predictions. + image_sizes: ``(H, W)`` for each image in the batch. + + Returns: + list of dicts each with keys ``'boxes'`` (absolute *xyxy*), + ``'labels'`` (1-indexed), and ``'scores'``. + """ + probs = pred_logits.softmax(-1) # (B, Q, C+1) + # Exclude background class (index 0) when computing scores + fg_probs = probs[:, :, 1:] # (B, Q, C) + scores, labels = fg_probs.max(-1) # (B, Q) + labels = labels + 1 # restore 1-indexed class labels + + results = [] + for b in range(pred_logits.shape[0]): + img_h, img_w = image_sizes[b] + + # Denormalise boxes from [0, 1] cxcywh → absolute xyxy + boxes_abs = pred_boxes[b].clone() + boxes_abs[:, [0, 2]] = boxes_abs[:, [0, 2]] * img_w + boxes_abs[:, [1, 3]] = boxes_abs[:, [1, 3]] * img_h + boxes_xyxy = box_convert(boxes_abs, in_fmt="cxcywh", out_fmt="xyxy") + + keep = scores[b] >= self.score_threshold + results.append({ + "boxes": boxes_xyxy[keep], + "labels": labels[b][keep], + "scores": scores[b][keep], + }) + + return results + + +# --------------------------------------------------------------------------- +# Complete model +# --------------------------------------------------------------------------- + +class DINOv2HungarianDetectionModel(nn.Module): + """Complete DINOv2 + :class:`HungarianDetectionHead` detection model. + + This ``nn.Module`` is the object stored in + :attr:`DINOv2HungarianDetector.model`. It wires together a + :class:`~feature_extractors.DINOv2ViT` backbone (frozen by default) and + a :class:`HungarianDetectionHead`. + + The model's ``forward`` method mirrors the interface expected by + :class:`~detectors.AbstractDetector`: + + * **Training mode** – call with ``(images, targets)``; returns a loss dict. + * **Eval mode** – call with ``(images,)``; returns a list of detection + dicts (``'boxes'``, ``'labels'``, ``'scores'``). + + .. note:: + All images in a batch must have the *same* spatial dimensions because + the DINOv2 ViT backbone maps each image to a fixed-size patch grid. + Use a resize transform to ensure consistent dimensions during training + and inference. + + Args: + backbone (DINOv2ViT): Feature extractor. + head (HungarianDetectionHead): Detection head. + """ + + def __init__(self, backbone: DINOv2ViT, head: HungarianDetectionHead): + super().__init__() + self.backbone = backbone + self.head = head + + def forward( + self, + images: list[torch.Tensor] | torch.Tensor, + targets: list[dict] | None = None, + ): + """Run a full forward pass. + + Args: + images: Either a list of ``(C, H, W)`` tensors (all the same size) + or a pre-stacked ``(B, C, H, W)`` tensor. + targets (list[dict] | None): Ground-truth dicts required during + training. Each dict should contain ``'boxes'`` and + ``'labels'``; ``'image_size'`` is filled in automatically. + + Returns: + dict | list[dict]: Loss dict during training; list of detection + dicts during inference. + """ + if isinstance(images, (list, tuple)): + image_sizes = [(img.shape[-2], img.shape[-1]) for img in images] + x = torch.stack(images, dim=0) + else: + image_sizes = [(images.shape[-2], images.shape[-1])] * images.shape[0] + x = images + + # Attach image sizes to targets so the matcher can normalise GT boxes + if targets is not None: + for i, t in enumerate(targets): + if "image_size" not in t: + t["image_size"] = image_sizes[i] + + features = self.backbone(x) # (B, D, H_p, W_p) + return self.head(features, targets=targets, image_sizes=image_sizes) diff --git a/requirements.txt b/requirements.txt index cae4967..d04e0ec 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,3 +3,4 @@ torchvision==0.27.1 pandas==3.0.3 Pillow==12.2.0 transformers>=4.38.0 +scipy>=1.11.0 From f15be25345a49c24791844637775ba166e119bb8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 19:09:53 +0000 Subject: [PATCH 06/10] Fix loss key naming: use no-prefix convention consistent with existing detectors --- detectors.py | 2 +- hungarian_head.py | 14 +++++++------- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/detectors.py b/detectors.py index 259803a..0dd181d 100644 --- a/detectors.py +++ b/detectors.py @@ -760,7 +760,7 @@ def __init__( ) def get_loss_names(self) -> list[str]: - return ["loss_classification", "loss_bbox", "loss_giou"] + return ["classification", "bbox", "giou"] def load_pretrained_model(self): backbone = DINOv2ViT( diff --git a/hungarian_head.py b/hungarian_head.py index b8fc38b..14818a4 100644 --- a/hungarian_head.py +++ b/hungarian_head.py @@ -362,8 +362,8 @@ def forward( Returns: dict | list[dict]: - * **Training** – a dict with keys ``'loss_classification'``, - ``'loss_bbox'``, and ``'loss_giou'``. + * **Training** – a dict with keys ``'classification'``, + ``'bbox'``, and ``'giou'``. * **Inference** – a list of dicts each containing ``'boxes'`` (absolute *xyxy*), ``'labels'`` (1-indexed), and ``'scores'``. """ @@ -407,8 +407,8 @@ def _compute_loss( targets: list of ground-truth dicts (see :meth:`forward`). Returns: - dict with keys ``'loss_classification'``, ``'loss_bbox'``, - ``'loss_giou'``. + dict with keys ``'classification'``, ``'bbox'``, + ``'giou'``. """ B, Q, _ = pred_logits.shape device = pred_logits.device @@ -470,9 +470,9 @@ def _compute_loss( loss_giou = pred_boxes.sum() * 0.0 return { - "loss_classification": self.loss_weight_class * loss_cls, - "loss_bbox": self.loss_weight_bbox * loss_bbox, - "loss_giou": self.loss_weight_giou * loss_giou, + "classification": self.loss_weight_class * loss_cls, + "bbox": self.loss_weight_bbox * loss_bbox, + "giou": self.loss_weight_giou * loss_giou, } # ------------------------------------------------------------------ From adb3dcede8ba5cf2e9a1724bd83a65bc2f121c8c Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 19:49:16 +0000 Subject: [PATCH 07/10] Add mkdocs usage documentation --- README.md | 8 +++++ mkdocs/dataset.md | 56 ++++++++++++++++++++++++++++++++ mkdocs/index.md | 9 ++++++ mkdocs/inference.md | 69 +++++++++++++++++++++++++++++++++++++++ mkdocs/training.md | 79 +++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 221 insertions(+) create mode 100644 mkdocs/dataset.md create mode 100644 mkdocs/index.md create mode 100644 mkdocs/inference.md create mode 100644 mkdocs/training.md diff --git a/README.md b/README.md index 399711a..321cf20 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,14 @@ There are six models and their subtypes that you can use with this library: For more details about these models see the official PyTorch Documentation. +## Additional documentation + +Detailed usage notes are available in the `mkdocs/` directory: + +- `mkdocs/dataset.md` +- `mkdocs/training.md` +- `mkdocs/inference.md` + ## Training models First you need to import the necessary libraries: diff --git a/mkdocs/dataset.md b/mkdocs/dataset.md new file mode 100644 index 0000000..36cefd5 --- /dev/null +++ b/mkdocs/dataset.md @@ -0,0 +1,56 @@ +# Dataset organization + +The training code expects a COCO-style dataset: + +- image files live in a directory that you pass as `train_data_dir` or `val_data_dir` +- annotations live in a COCO JSON file that you pass as `train_annotation_file` or `val_annotation_file` +- each image entry in the annotation file must match a file inside the corresponding image directory + +A practical layout looks like this: + +```text +project-root/ +├── datasets/ +│ └── my_dataset/ +│ ├── train/ +│ │ ├── image_0001.jpg +│ │ ├── image_0002.jpg +│ │ └── ... +│ ├── val/ +│ │ ├── image_0101.jpg +│ │ ├── image_0102.jpg +│ │ └── ... +│ └── annotations/ +│ ├── train.json +│ └── val.json +└── train.py +``` + +## What the loader reads + +`build_coco_dataset(...)` uses `torchvision.datasets.CocoDetection`, so the JSON file must follow the COCO detection schema. + +During training, `extract_images_targets(...)` converts each annotation into PyTorch detection targets: + +- `bbox` is expected in COCO `xywh` format and is converted to `xyxy` +- `category_id` is converted to the training label with `category_id + 1` +- label `0` is therefore reserved for the background class + +## Class indexing rule + +Set `num_classes` to the number of foreground classes in your dataset. + +Because the code shifts every `category_id` by `+1`, the safest setup is to keep annotation class ids dense and zero-based: + +- first class -> `category_id: 0` +- second class -> `category_id: 1` +- ... +- last class -> `category_id: num_classes - 1` + +## Split handling + +The current training loop always requires a training split and can optionally use a validation split: + +- `train_data_dir` + `train_annotation_file` are required to train +- `val_data_dir` + `val_annotation_file` enable validation mAP logging +- test-related arguments exist on `train(...)`, but the current implementation does not consume them yet diff --git a/mkdocs/index.md b/mkdocs/index.md new file mode 100644 index 0000000..385ce74 --- /dev/null +++ b/mkdocs/index.md @@ -0,0 +1,9 @@ +# PyTorch Detector Documentation + +This directory collects the project-specific usage notes that are only briefly covered in the root README. + +- [Dataset organization](dataset.md) explains the folder layout and annotation assumptions expected by `build_coco_dataset(...)` and `extract_images_targets(...)`. +- [Training a detector](training.md) shows how the detector wrappers are instantiated and how `AbstractDetector.train(...)` uses the dataset inputs, logs, checkpoints, and validation split. +- [Inference with a detector](inference.md) explains what `detector.inference()` returns and how batched prediction and visualization work. + +The documentation in this folder reflects the current code in `/home/runner/work/pytorch-detector/pytorch-detector/data.py`, `/home/runner/work/pytorch-detector/pytorch-detector/detectors.py`, and `/home/runner/work/pytorch-detector/pytorch-detector/inference.py`. diff --git a/mkdocs/inference.md b/mkdocs/inference.md new file mode 100644 index 0000000..7f8ed82 --- /dev/null +++ b/mkdocs/inference.md @@ -0,0 +1,69 @@ +# Inference with a detector + +Every detector wrapper exposes an `inference()` helper: + +```python +detector = FasterRCNNV2Detector(num_classes=num_classes, resume=checkpoint, device=device) +inference = detector.inference() +``` + +`detector.inference()` returns an `Inference` object from `/home/runner/work/pytorch-detector/pytorch-detector/inference.py`. + +## Default behavior + +The inference helper: + +- keeps a reference to the detector instance +- uses a default batch size of `4` +- converts images to float tensors with torchvision v2 transforms +- switches the model to evaluation mode and runs under `torch.no_grad()` + +## Running predictions + +There are two main prediction flows. + +### `get_results(images)` + +Use this when you already have a list of loaded PIL images. + +It returns the raw model output for each image, including: + +- `boxes` +- `labels` +- `scores` + +### `get_results_df(image_paths, resize=224)` + +Use this when your input is a list of file paths. + +For each batch, the helper: + +1. loads the images from disk +2. resizes every image to a square `resize x resize` +3. runs the detector +4. flattens all detections into a Pandas DataFrame + +The resulting DataFrame contains: + +- `filename` +- `width` +- `height` +- `class` +- `class_index` +- `xmin`, `ymin`, `xmax`, `ymax` +- `score` + +## Drawing boxes + +`draw_results(...)` overlays predictions on PIL images. + +Useful options: + +- `score_threshold`: hides predictions below the threshold +- `color`: bounding-box and label color +- `show_labels`: toggles numeric class labels +- `make_copy`: keeps the original images unchanged when `True` + +## Important inference detail + +The helper reports numeric class ids, not class names. If you need human-readable labels, keep your own mapping from dataset class id to display name next to your inference script. diff --git a/mkdocs/training.md b/mkdocs/training.md new file mode 100644 index 0000000..7c08707 --- /dev/null +++ b/mkdocs/training.md @@ -0,0 +1,79 @@ +# Training a detector + +All detector wrappers live in `/home/runner/work/pytorch-detector/pytorch-detector/detectors.py`. They expose the same high-level `train(...)` entry point through `AbstractDetector`. + +## 1. Pick a detector class + +Examples include: + +- `FasterRCNNDetector` +- `FasterRCNNV2Detector` +- `RetinaNetResNet50FPNDetector` +- `FCOSResNet50FPNDetector` +- `SSD300VGG16Detector` +- `DINOv2HungarianDetector` + +Create the detector with the number of foreground classes and the device you want to use. + +## 2. Prepare transforms and paths + +The `train(...)` method expects: + +- a training image directory +- a training COCO annotation file +- optional validation image and annotation paths +- torchvision v2 transforms for the training and validation sets +- a `log_dir` where epoch logs and checkpoints should be written + +## 3. Call `train(...)` + +The training loop performs these steps: + +1. rebuilds the model head for `num_classes` +2. constructs a COCO training dataloader +3. optionally constructs a validation dataloader +4. trains for `n_epochs` +5. saves checkpoints every `save_every` epochs when enabled +6. logs metrics to `epochs.yaml` inside `log_dir` +7. computes train/validation mAP values when a validation split is provided + +## Key training arguments + +- `n_epochs`: total number of epochs to run +- `lr`: SGD learning rate +- `batch_size`: batch size for the training dataloader +- `resume`: checkpoint path to load model weights from +- `save_every`: checkpoint frequency +- `lr_step_every`: epoch interval used before stepping the learning-rate scheduler +- `log_dir`: output directory for logs and checkpoints +- `n_batches_validation`: number of batches evaluated when computing mAP metrics + +## Logging and checkpoints + +Training creates the log directory if it does not exist. + +Inside that directory, the current implementation writes: + +- `epochs.yaml` with one record per epoch +- checkpoint files named like `_epoch0001.pth` + +Each checkpoint stores: + +- `model_state` +- `optim_state` +- `scheduler_state` +- the detector name and class count +- the metrics collected for that epoch + +## Validation behavior + +Validation only runs when `val_data_dir` is provided. + +When validation is enabled, the trainer computes: + +- `train_map50` +- `train_mAP50_95` +- `val_map50` +- `val_mAP50_95` + +The training loop also builds a small non-shuffled training dataloader for the train-side metric snapshot. From 1730b55750201c0f7edffc29a7df2fbeb294418f Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 19 Jun 2026 19:50:00 +0000 Subject: [PATCH 08/10] Fix mkdocs path references --- mkdocs/index.md | 2 +- mkdocs/inference.md | 2 +- mkdocs/training.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/mkdocs/index.md b/mkdocs/index.md index 385ce74..b04985c 100644 --- a/mkdocs/index.md +++ b/mkdocs/index.md @@ -6,4 +6,4 @@ This directory collects the project-specific usage notes that are only briefly c - [Training a detector](training.md) shows how the detector wrappers are instantiated and how `AbstractDetector.train(...)` uses the dataset inputs, logs, checkpoints, and validation split. - [Inference with a detector](inference.md) explains what `detector.inference()` returns and how batched prediction and visualization work. -The documentation in this folder reflects the current code in `/home/runner/work/pytorch-detector/pytorch-detector/data.py`, `/home/runner/work/pytorch-detector/pytorch-detector/detectors.py`, and `/home/runner/work/pytorch-detector/pytorch-detector/inference.py`. +The documentation in this folder reflects the current code in `data.py`, `detectors.py`, and `inference.py`. diff --git a/mkdocs/inference.md b/mkdocs/inference.md index 7f8ed82..2a3da27 100644 --- a/mkdocs/inference.md +++ b/mkdocs/inference.md @@ -7,7 +7,7 @@ detector = FasterRCNNV2Detector(num_classes=num_classes, resume=checkpoint, devi inference = detector.inference() ``` -`detector.inference()` returns an `Inference` object from `/home/runner/work/pytorch-detector/pytorch-detector/inference.py`. +`detector.inference()` returns an `Inference` object from `inference.py`. ## Default behavior diff --git a/mkdocs/training.md b/mkdocs/training.md index 7c08707..04c6fcf 100644 --- a/mkdocs/training.md +++ b/mkdocs/training.md @@ -1,6 +1,6 @@ # Training a detector -All detector wrappers live in `/home/runner/work/pytorch-detector/pytorch-detector/detectors.py`. They expose the same high-level `train(...)` entry point through `AbstractDetector`. +All detector wrappers live in `detectors.py`. They expose the same high-level `train(...)` entry point through `AbstractDetector`. ## 1. Pick a detector class From 1649f8106c70d4b902445cbb04ce6c0413429393 Mon Sep 17 00:00:00 2001 From: Tilo Flasche Date: Fri, 19 Jun 2026 22:10:59 +0200 Subject: [PATCH 09/10] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- feature_extractors.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/feature_extractors.py b/feature_extractors.py index ed3bf17..697dc21 100644 --- a/feature_extractors.py +++ b/feature_extractors.py @@ -62,7 +62,9 @@ def __init__( def _freeze(self): for param in self.model.parameters(): param.requires_grad = False - + if getattr(self, "layer_norm", False): + for param in self.norms.parameters(): + param.requires_grad = False def forward(self, x: torch.Tensor): """Forward image through DINOv2 ViT. From 09202d24800df7e38006e7f5e86e1cc631b49445 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 20 Jun 2026 09:22:04 +0000 Subject: [PATCH 10/10] Add decoder-only deformable neck for Hungarian detector --- detectors.py | 6 ++ hungarian_head.py | 211 ++++++++++++++++++++++++++++++++++++++-------- 2 files changed, 180 insertions(+), 37 deletions(-) diff --git a/detectors.py b/detectors.py index 0dd181d..959d5b3 100644 --- a/detectors.py +++ b/detectors.py @@ -724,6 +724,8 @@ class DINOv2HungarianDetector(AbstractDetector): Default: 8. num_decoder_layers (int): Depth of the transformer decoder. Default: 6. + num_feature_levels (int): Number of multi-scale feature levels used by + the deformable decoder neck. Default: 4. score_threshold (float): Minimum foreground score to keep during inference. Default: 0.5. """ @@ -740,6 +742,7 @@ def __init__( num_queries: int = 100, nhead: int = 8, num_decoder_layers: int = 6, + num_feature_levels: int = 4, score_threshold: float = 0.5, ): # Set attributes before calling super().__init__ because @@ -750,6 +753,7 @@ def __init__( self.num_queries = num_queries self.nhead = nhead self.num_decoder_layers = num_decoder_layers + self.num_feature_levels = num_feature_levels self.score_threshold = score_threshold super().__init__( name="dinov2_hungarian", @@ -777,6 +781,7 @@ def load_pretrained_model(self): num_queries=self.num_queries, nhead=self.nhead, num_decoder_layers=self.num_decoder_layers, + num_feature_levels=self.num_feature_levels, score_threshold=self.score_threshold, ) return DINOv2HungarianDetectionModel(backbone=backbone, head=head) @@ -790,5 +795,6 @@ def replace_head(self, model, num_classes: int): num_queries=self.num_queries, nhead=self.nhead, num_decoder_layers=self.num_decoder_layers, + num_feature_levels=self.num_feature_levels, score_threshold=self.score_threshold, ) diff --git a/hungarian_head.py b/hungarian_head.py index 14818a4..efc37e5 100644 --- a/hungarian_head.py +++ b/hungarian_head.py @@ -7,6 +7,8 @@ * :class:`HungarianMatcher` – optimal bipartite matching between predicted and ground-truth objects using the Hungarian algorithm (:func:`scipy.optimize.linear_sum_assignment`). +* :class:`DeformableDecoderNeck` – decoder-only deformable-transformer-inspired + neck that refines backbone features before detection. * :class:`HungarianDetectionHead` – transformer-decoder detection head that combines the matcher with a set-based training loss (classification, L1-box, and GIoU). @@ -117,6 +119,162 @@ def forward(self, mask: torch.Tensor) -> torch.Tensor: return pos.flatten(1, 2) # (B, H*W, hidden_dim) +class DeformableDecoderLayer(nn.Module): + """Single decoder block used by :class:`DeformableDecoderNeck`.""" + + def __init__( + self, + hidden_dim: int, + nhead: int, + dim_feedforward: int, + dropout: float, + ): + super().__init__() + self.self_attn = nn.MultiheadAttention( + embed_dim=hidden_dim, + num_heads=nhead, + dropout=dropout, + batch_first=True, + ) + self.cross_attn = nn.MultiheadAttention( + embed_dim=hidden_dim, + num_heads=nhead, + dropout=dropout, + batch_first=True, + ) + self.linear1 = nn.Linear(hidden_dim, dim_feedforward) + self.dropout = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, hidden_dim) + + self.norm1 = nn.LayerNorm(hidden_dim) + self.norm2 = nn.LayerNorm(hidden_dim) + self.norm3 = nn.LayerNorm(hidden_dim) + self.dropout1 = nn.Dropout(dropout) + self.dropout2 = nn.Dropout(dropout) + self.dropout3 = nn.Dropout(dropout) + self.activation = nn.ReLU(inplace=True) + + def forward( + self, + tgt: torch.Tensor, + query_pos: torch.Tensor, + memory: torch.Tensor, + ) -> torch.Tensor: + q = k = tgt + query_pos + tgt2 = self.self_attn(q, k, value=tgt)[0] + tgt = self.norm1(tgt + self.dropout1(tgt2)) + + tgt2 = self.cross_attn(query=tgt + query_pos, key=memory, value=memory)[0] + tgt = self.norm2(tgt + self.dropout2(tgt2)) + + tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) + tgt = self.norm3(tgt + self.dropout3(tgt2)) + return tgt + + +class DeformableDecoderNeck(nn.Module): + """Decoder-only deformable-transformer-inspired neck. + + The neck creates a small multi-scale feature set from the backbone output, + flattens levels into a shared memory sequence, and runs a stack of decoder + layers over learnable object queries. No encoder stack is used. + """ + + def __init__( + self, + in_channels: int, + hidden_dim: int = 256, + num_queries: int = 100, + nhead: int = 8, + num_decoder_layers: int = 6, + dim_feedforward: int = 2048, + dropout: float = 0.1, + num_feature_levels: int = 4, + ): + super().__init__() + if num_feature_levels < 1: + raise ValueError("num_feature_levels must be at least 1.") + + self.hidden_dim = hidden_dim + self.num_feature_levels = num_feature_levels + self.input_proj = nn.Conv2d(in_channels, hidden_dim, kernel_size=1) + self.downsample_blocks = nn.ModuleList( + [ + nn.Conv2d(hidden_dim, hidden_dim, kernel_size=3, stride=2, padding=1) + for _ in range(num_feature_levels - 1) + ] + ) + self.level_embed = nn.Embedding(num_feature_levels, hidden_dim) + self.pos_encoding = SinusoidalPositionEncoding2D(hidden_dim) + + self.query_embed = nn.Embedding(num_queries, hidden_dim) + self.reference_points = nn.Linear(hidden_dim, 2) + self.ref_point_proj = _build_mlp(hidden_dim, hidden_dim, hidden_dim, num_layers=2) + + self.decoder_layers = nn.ModuleList( + [ + DeformableDecoderLayer( + hidden_dim=hidden_dim, + nhead=nhead, + dim_feedforward=dim_feedforward, + dropout=dropout, + ) + for _ in range(num_decoder_layers) + ] + ) + + def _build_memory(self, features: torch.Tensor) -> torch.Tensor: + feat = self.input_proj(features) + levels = [feat] + for block in self.downsample_blocks: + levels.append(block(levels[-1])) + + memory_tokens = [] + for level_idx, level_feat in enumerate(levels): + B, _, H, W = level_feat.shape + mask = torch.zeros(B, H, W, dtype=torch.bool, device=level_feat.device) + pos = self.pos_encoding(mask) + tokens = level_feat.flatten(2).permute(0, 2, 1) + level_bias = self.level_embed.weight[level_idx].view(1, 1, -1) + memory_tokens.append(tokens + pos + level_bias) + return torch.cat(memory_tokens, dim=1) + + def _reference_positional_encoding(self, reference_points: torch.Tensor) -> torch.Tensor: + # reference_points: (B, Q, 2) in [0, 1] + scale = 2.0 * math.pi + dim_t = torch.arange( + self.hidden_dim // 2, + dtype=reference_points.dtype, + device=reference_points.device, + ) + dim_t = 10000 ** (2 * (dim_t // 2) / max(1, self.hidden_dim // 2)) + + ref = reference_points[:, :, :, None] * scale + ref = ref / dim_t + pos = torch.stack((ref[..., 0::2].sin(), ref[..., 1::2].cos()), dim=-1).flatten(-2) + pos = pos.flatten(-2) + if pos.shape[-1] != self.hidden_dim: + pos = pos[..., : self.hidden_dim] + return self.ref_point_proj(pos) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + B = features.shape[0] + memory = self._build_memory(features) + query = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1) + reference_points = self.reference_points(query).sigmoid() + query_pos = self._reference_positional_encoding(reference_points) + + refined_queries = query + for layer in self.decoder_layers: + refined_queries = layer( + tgt=refined_queries, + query_pos=query_pos, + memory=memory, + ) + + return refined_queries + + # --------------------------------------------------------------------------- # Hungarian matcher # --------------------------------------------------------------------------- @@ -234,13 +392,11 @@ def forward( class HungarianDetectionHead(nn.Module): """DETR-style detection head that uses Hungarian matching for supervision. - The head accepts a single spatial feature map from a DINOv2 backbone, - projects it to ``hidden_dim`` channels, adds 2-D sinusoidal positional - encodings, and feeds the result as the *memory* of a standard - :class:`torch.nn.TransformerDecoder`. A set of ``num_queries`` learnable - object queries act as the *target* sequence. The decoded queries are - independently passed through a classification head and a bounding-box - regression MLP. + The head accepts a single spatial feature map from a DINOv2 backbone and + first routes it through a decoder-only + :class:`~hungarian_head.DeformableDecoderNeck` (no transformer encoder). + The decoded query features are then passed through a classification head + and a bounding-box regression MLP. **Training** – ground-truth objects are assigned to predicted queries via the :class:`HungarianMatcher`. Three losses are computed: @@ -265,6 +421,8 @@ class HungarianDetectionHead(nn.Module): Default: 8. num_decoder_layers (int): Depth of the transformer decoder. Default: 6. + num_feature_levels (int): Number of memory feature levels built by the + deformable decoder neck. Default: 4. dim_feedforward (int): Feedforward dimension inside each transformer layer. Default: 2048. dropout (float): Dropout probability in the transformer. Default: 0.1. @@ -290,6 +448,7 @@ def __init__( num_queries: int = 100, nhead: int = 8, num_decoder_layers: int = 6, + num_feature_levels: int = 4, dim_feedforward: int = 2048, dropout: float = 0.1, weight_class: float = 1.0, @@ -308,26 +467,16 @@ def __init__( self.loss_weight_bbox = loss_weight_bbox self.loss_weight_giou = loss_weight_giou - # Project backbone features to hidden_dim - self.input_proj = nn.Conv2d(in_channels, hidden_dim, kernel_size=1) - - # Positional encoding for memory tokens - self.pos_encoding = SinusoidalPositionEncoding2D(hidden_dim) - - # Learnable object queries - self.query_embed = nn.Embedding(num_queries, hidden_dim) - - # Transformer decoder - decoder_layer = nn.TransformerDecoderLayer( - d_model=hidden_dim, + # Decoder-only deformable neck (no encoder stack) + self.neck = DeformableDecoderNeck( + in_channels=in_channels, + hidden_dim=hidden_dim, + num_queries=num_queries, nhead=nhead, + num_decoder_layers=num_decoder_layers, + num_feature_levels=num_feature_levels, dim_feedforward=dim_feedforward, dropout=dropout, - batch_first=True, - ) - self.transformer_decoder = nn.TransformerDecoder( - decoder_layer=decoder_layer, - num_layers=num_decoder_layers, ) # Prediction heads @@ -367,19 +516,7 @@ def forward( * **Inference** – a list of dicts each containing ``'boxes'`` (absolute *xyxy*), ``'labels'`` (1-indexed), and ``'scores'``. """ - B, _, H, W = features.shape - - # Project features and add positional encoding - feat = self.input_proj(features) # (B, D, H, W) - mask = torch.zeros(B, H, W, dtype=torch.bool, device=features.device) - pos = self.pos_encoding(mask) # (B, H*W, D) - memory = feat.flatten(2).permute(0, 2, 1) + pos # (B, H*W, D) - - # Expand learnable queries for the whole batch - queries = self.query_embed.weight.unsqueeze(0).expand(B, -1, -1) # (B, Q, D) - - # Decode - decoded = self.transformer_decoder(tgt=queries, memory=memory) # (B, Q, D) + decoded = self.neck(features) # (B, Q, D) # Predictions pred_logits = self.class_head(decoded) # (B, Q, C+1)