Add decoder-only deformable neck to Hungarian DETR detector

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 <a href="https://pytorch.org/vision/stable/models.html#object-detection" target="_blank">official PyTorch Documentation</a>.
 
+## 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:

detectors.py

@@ -7,6 +7,8 @@
 from datetime import datetime
 import logger
 import inference
+from feature_extractors import DINOv2ViT, DINOv2ViTBackbone, DINOv2ConvNextBackbone
+from hungarian_head import HungarianDetectionHead, DINOv2HungarianDetectionModel
 
 
 class AbstractDetector():
@@ -567,3 +569,232 @@ 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
+    five-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).
+        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__(
+        self,
+        num_classes: int = None,
+        resume: str = None,
+        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,
+            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(
+            model_name=self.model_name,
+            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 FPN + 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).
+        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__(
+        self,
+        num_classes: int = None,
+        resume: str = None,
+        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,
+            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(
+            model_name=self.model_name,
+            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 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.
+        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.
+    """
+
+    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,
+        num_feature_levels: int = 4,
+        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.num_feature_levels = num_feature_levels
+        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 ["classification", "bbox", "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,
+            num_feature_levels=self.num_feature_levels,
+            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,
+            num_feature_levels=self.num_feature_levels,
+            score_threshold=self.score_threshold,
+        )