diff --git a/=1.11.0 b/=1.11.0
new file mode 100644
index 0000000..e69de29
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/detectors.py b/detectors.py
index 81e86d7..959d5b3 100644
--- a/detectors.py
+++ b/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,
+ )
diff --git a/feature_extractors.py b/feature_extractors.py
new file mode 100644
index 0000000..697dc21
--- /dev/null
+++ b/feature_extractors.py
@@ -0,0 +1,297 @@
+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:
+ 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.
+ 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,
+ model_name: str = None,
+ 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
+
+ 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)
+
+ # 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([
+ 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
+ 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.
+
+ 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:
+ # 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})
+ )
+ )
+
+ z = self.model(x)
+
+ for h in hooks:
+ h.remove()
+
+ 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:
+ feat = self.norms[idx](feat)
+ # Remove the [CLS] token
+ feat = feat[:, 1:, :]
+ B, P, D = feat.shape
+ 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:
+ 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:
+ 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].
+ """
+
+ def __init__(
+ self,
+ model_name: str = None,
+ finetuning: bool = False,
+ layers: list = None,
+ ):
+ super().__init__()
+ if layers is None:
+ layers = [0, 1, 2, 3]
+ self.finetuning = finetuning
+ self.layers = layers
+
+ 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()
+
+ 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:
+ # 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})
+ )
+ )
+
+ 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:
+ 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,
+ model_name: str = None,
+ out_channels: int = 256,
+ finetuning: bool = False,
+ ):
+ super().__init__()
+ 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)
+ 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` (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,
+ model_name: str = None,
+ out_channels: int = 256,
+ finetuning: bool = False,
+ ):
+ super().__init__()
+ 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,
+ out_channels=out_channels,
+ extra_blocks=LastLevelMaxPool(),
+ )
+ self.out_channels = out_channels
+
+ 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(stage_idx), f)
+ for stage_idx, f in zip(self.body.layers, stage_feats)
+ )
+ return self.fpn(feat_dict)
+
diff --git a/hungarian_head.py b/hungarian_head.py
new file mode 100644
index 0000000..efc37e5
--- /dev/null
+++ b/hungarian_head.py
@@ -0,0 +1,729 @@
+"""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:`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).
+* :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)
+
+
+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
+# ---------------------------------------------------------------------------
+
+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 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:
+
+ * 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.
+ 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.
+ 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,
+ num_feature_levels: int = 4,
+ 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
+
+ # 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,
+ )
+
+ # 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 ``'classification'``,
+ ``'bbox'``, and ``'giou'``.
+ * **Inference** – a list of dicts each containing ``'boxes'``
+ (absolute *xyxy*), ``'labels'`` (1-indexed), and ``'scores'``.
+ """
+ decoded = self.neck(features) # (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 ``'classification'``, ``'bbox'``,
+ ``'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 {
+ "classification": self.loss_weight_class * loss_cls,
+ "bbox": self.loss_weight_bbox * loss_bbox,
+ "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/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..b04985c
--- /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 `data.py`, `detectors.py`, and `inference.py`.
diff --git a/mkdocs/inference.md b/mkdocs/inference.md
new file mode 100644
index 0000000..2a3da27
--- /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 `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..04c6fcf
--- /dev/null
+++ b/mkdocs/training.md
@@ -0,0 +1,79 @@
+# Training a detector
+
+All detector wrappers live in `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.
diff --git a/requirements.txt b/requirements.txt
index f370a15..d04e0ec 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,3 +2,5 @@ torch==2.12.1
torchvision==0.27.1
pandas==3.0.3
Pillow==12.2.0
+transformers>=4.38.0
+scipy>=1.11.0