Add DINOv3 OWEED ViTDet training pipeline

detectron2/modeling/backbone/__init__.py

@@ -13,6 +13,7 @@
     BottleneckBlock,
 )
 from .vit import ViT, SimpleFeaturePyramid, get_vit_lr_decay_rate
+from .hf_dinov3 import HFDINOv3ViT
 from .mvit import MViT
 from .swin import SwinTransformer
 

detectron2/modeling/backbone/hf_dinov3.py

@@ -0,0 +1,108 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from typing import Any, Dict, Optional
+
+import torch
+
+from detectron2.layers import ShapeSpec
+
+from .backbone import Backbone
+
+__all__ = ["HFDINOv3ViT"]
+
+
+class HFDINOv3ViT(Backbone):
+    """
+    Hugging Face DINOv3 ViT backbone adapter for ViTDet-style models.
+
+    The adapter exposes the DINOv3 patch tokens as a single spatial feature map
+    with stride equal to the patch size, matching the contract expected by
+    :class:`SimpleFeaturePyramid`.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "facebook/dinov3-vitb16-pretrain-lvd1689m",
+        *,
+        out_feature: str = "last_feat",
+        pretrained: bool = True,
+        freeze: bool = True,
+        config_kwargs: Optional[Dict[str, Any]] = None,
+        model_kwargs: Optional[Dict[str, Any]] = None,
+    ):
+        super().__init__()
+        config_kwargs = {} if config_kwargs is None else dict(config_kwargs)
+        model_kwargs = {} if model_kwargs is None else dict(model_kwargs)
+
+        try:
+            from transformers import DINOv3ViTConfig, DINOv3ViTModel
+        except ImportError as exc:
+            raise ImportError(
+                "HFDINOv3ViT requires Hugging Face Transformers. "
+                "Install it with `pip install transformers`."
+            ) from exc
+
+        if pretrained:
+            self.model = DINOv3ViTModel.from_pretrained(model_name, **model_kwargs)
+        else:
+            self.model = DINOv3ViTModel(DINOv3ViTConfig(**config_kwargs))
+
+        self.out_feature = out_feature
+        self.freeze_encoder = freeze
+
+        cfg = self.model.config
+        self.patch_size = int(cfg.patch_size)
+        self.num_prefix_tokens = 1 + int(getattr(cfg, "num_register_tokens", 0))
+        self._out_features = [out_feature]
+        self._out_feature_channels = {out_feature: int(cfg.hidden_size)}
+        self._out_feature_strides = {out_feature: self.patch_size}
+
+        if self.freeze_encoder:
+            for parameter in self.model.parameters():
+                parameter.requires_grad = False
+            self.model.eval()
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        if self.freeze_encoder:
+            self.model.eval()
+        return self
+
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
+        height, width = x.shape[-2:]
+        if height % self.patch_size != 0 or width % self.patch_size != 0:
+            raise ValueError(
+                f"Input height/width must be divisible by patch_size={self.patch_size}; "
+                f"got {(height, width)}."
+            )
+
+        outputs = self.model(pixel_values=x, return_dict=True)
+        hidden = outputs.last_hidden_state
+
+        if hidden.dim() == 4:
+            features = hidden
+        else:
+            grid_h, grid_w = height // self.patch_size, width // self.patch_size
+            expected_tokens = grid_h * grid_w
+            if hidden.shape[1] == expected_tokens:
+                patch_tokens = hidden
+            elif hidden.shape[1] >= self.num_prefix_tokens + expected_tokens:
+                patch_tokens = hidden[
+                    :, self.num_prefix_tokens : self.num_prefix_tokens + expected_tokens, :
+                ]
+            else:
+                raise ValueError(
+                    "DINOv3 output token count cannot be reshaped to a spatial feature map: "
+                    f"got {hidden.shape[1]} tokens, expected {expected_tokens} patch tokens "
+                    f"after {self.num_prefix_tokens} prefix tokens."
+                )
+            features = patch_tokens.reshape(x.shape[0], grid_h, grid_w, -1).permute(0, 3, 1, 2)
+
+        return {self.out_feature: features}
+
+    def output_shape(self):
+        return {
+            self.out_feature: ShapeSpec(
+                channels=self._out_feature_channels[self.out_feature],
+                stride=self._out_feature_strides[self.out_feature],
+            )
+        }

detectron2/modeling/meta_arch/__init__.py

@@ -7,6 +7,7 @@
 
 # import all the meta_arch, so they will be registered
 from .rcnn import GeneralizedRCNN, ProposalNetwork
+from .bbox_only_eval_rcnn import GeneralizedRCNNWithBBoxOnlyEval
 from .dense_detector import DenseDetector
 from .retinanet import RetinaNet
 from .fcos import FCOS

detectron2/modeling/meta_arch/bbox_only_eval_rcnn.py

@@ -0,0 +1,37 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+from typing import Dict, List, Optional
+
+import torch
+
+from detectron2.structures import Instances
+
+from .rcnn import GeneralizedRCNN
+
+__all__ = ["GeneralizedRCNNWithBBoxOnlyEval"]
+
+
+class GeneralizedRCNNWithBBoxOnlyEval(GeneralizedRCNN):
+    """
+    GeneralizedRCNN variant that trains all configured heads, but skips mask
+    prediction during inference/evaluation.
+
+    This is useful for bbox-only COCO evaluation of mask models on very large
+    images, where postprocessing predicted masks would materialize full-size
+    bitmasks for every detection.
+    """
+
+    def inference(
+        self,
+        batched_inputs: List[Dict[str, torch.Tensor]],
+        detected_instances: Optional[List[Instances]] = None,
+        do_postprocess: bool = True,
+    ):
+        old_mask_on = getattr(self.roi_heads, "mask_on", None)
+        if old_mask_on is None:
+            return super().inference(batched_inputs, detected_instances, do_postprocess)
+
+        self.roi_heads.mask_on = False
+        try:
+            return super().inference(batched_inputs, detected_instances, do_postprocess)
+        finally:
+            self.roi_heads.mask_on = old_mask_on

projects/ViTDet/configs/COCO/mask_rcnn_dinov3_vitb16_100ep.py

@@ -0,0 +1,21 @@
+from detectron2.config import LazyCall as L
+from detectron2.modeling.backbone import HFDINOv3ViT
+
+from .mask_rcnn_vitdet_b_100ep import dataloader, lr_multiplier, model, optimizer, train
+
+
+# Keep the ViTDet Mask R-CNN architecture unchanged, but replace the MAE ViT-B
+# encoder with a frozen Hugging Face DINOv3 ViT-B/16 encoder.
+model.backbone.net = L(HFDINOv3ViT)(
+    model_name="facebook/dinov3-vitb16-pretrain-lvd1689m",
+    out_feature="last_feat",
+    pretrained=True,
+    freeze=True,
+)
+
+# Initialize trainable detector components from the COCO ViTDet-B detector
+# checkpoint. The DINOv3 encoder loads its own Hugging Face weights above.
+train.init_checkpoint = (
+    "https://dl.fbaipublicfiles.com/detectron2/ViTDet/COCO/"
+    "mask_rcnn_vitdet_b/f325346929/model_final_61ccd1.pkl"
+)

projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen.py

@@ -0,0 +1,64 @@
+from detectron2.data.datasets import register_coco_instances
+from detectron2.modeling.meta_arch import GeneralizedRCNNWithBBoxOnlyEval
+
+from ..COCO.mask_rcnn_dinov3_vitb16_100ep import (
+    dataloader,
+    lr_multiplier,
+    model,
+    optimizer,
+    train,
+)
+
+
+_DATA_ROOT = "/home/nikhileswara/Datasets/oweed_structure_v0_100_images_coco"
+_TRAIN_NAME = "oweed_structure_v0_train"
+_VAL_NAME = "oweed_structure_v0_val"
+
+register_coco_instances(
+    _TRAIN_NAME,
+    {},
+    f"{_DATA_ROOT}/annotations/instances_train2017.json",
+    f"{_DATA_ROOT}/train2017",
+)
+register_coco_instances(
+    _VAL_NAME,
+    {},
+    f"{_DATA_ROOT}/annotations/instances_val2017.json",
+    f"{_DATA_ROOT}/val2017",
+)
+
+dataloader.train.dataset.names = _TRAIN_NAME
+# 2 RTX 4090s: total_batch_size is global across DDP workers, so this gives
+# two 1024x1024 LSJ images per GPU when launched with --num-gpus 2.
+dataloader.train.total_batch_size = 4
+dataloader.train.num_workers = 8
+dataloader.train.pin_memory = True
+dataloader.train.persistent_workers = True
+dataloader.train.prefetch_factor = 4
+dataloader.test.dataset.names = _VAL_NAME
+dataloader.test.num_workers = 8
+
+dataloader.evaluator.dataset_name = _VAL_NAME
+dataloader.evaluator.tasks = ("bbox",)
+
+model._target_ = GeneralizedRCNNWithBBoxOnlyEval
+model.roi_heads.num_classes = 44
+model.roi_heads.box_predictor.test_score_thresh = 0.02
+model.roi_heads.box_predictor.test_topk_per_image = 500
+
+# The dataset has up to ~460 annotated instances per image, so keep more
+# proposals and final detections than the COCO defaults.
+model.proposal_generator.pre_nms_topk = (4000, 4000)
+model.proposal_generator.post_nms_topk = (2000, 2000)
+
+optimizer.lr = 1e-4
+
+train.output_dir = "./output/oweed_dinov3_vitb16_frozen_mask_rcnn"
+train.max_iter = 5000
+train.eval_period = 1000
+train.checkpointer.period = 500
+train.checkpointer.max_to_keep = 10
+
+lr_multiplier.scheduler.milestones = [3500, 4500]
+lr_multiplier.scheduler.num_updates = train.max_iter
+lr_multiplier.warmup_length = 100 / train.max_iter

projects/ViTDet/configs/OWEED/mask_rcnn_dinov3_vitb16_frozen_tiled_keepfrag.py

@@ -0,0 +1,99 @@
+from fvcore.common.param_scheduler import MultiStepParamScheduler
+
+import detectron2.data.transforms as T
+from detectron2.config import LazyCall as L
+from detectron2.data.datasets import register_coco_instances
+from detectron2.modeling.meta_arch import GeneralizedRCNNWithBBoxOnlyEval
+from detectron2.solver import WarmupParamScheduler
+
+from ..COCO.mask_rcnn_dinov3_vitb16_100ep import (
+    dataloader,
+    lr_multiplier,
+    model,
+    optimizer,
+    train,
+)
+
+
+_DATA_ROOT = "/home/nikhileswara/Datasets/oweed_structure_v0_100_images_coco_tiled_1024_s512_keepfrag"
+_TRAIN_NAME = "oweed_structure_v0_tiled_keepfrag_train"
+_VAL_NAME = "oweed_structure_v0_tiled_keepfrag_val"
+
+register_coco_instances(
+    _TRAIN_NAME,
+    {},
+    f"{_DATA_ROOT}/annotations/instances_train2017.json",
+    f"{_DATA_ROOT}/train2017",
+)
+register_coco_instances(
+    _VAL_NAME,
+    {},
+    f"{_DATA_ROOT}/annotations/instances_val2017.json",
+    f"{_DATA_ROOT}/val2017",
+)
+
+dataloader.train.dataset.names = _TRAIN_NAME
+dataloader.train.total_batch_size = 8
+dataloader.train.num_workers = 8
+dataloader.train.pin_memory = True
+dataloader.train.persistent_workers = True
+dataloader.train.prefetch_factor = 4
+dataloader.train.mapper.augmentations = [
+    L(T.RandomFlip)(horizontal=True),
+    L(T.ResizeScale)(
+        min_scale=0.8,
+        max_scale=1.25,
+        target_height=1024,
+        target_width=1024,
+    ),
+    L(T.FixedSizeCrop)(crop_size=(1024, 1024), pad=True),
+]
+dataloader.train.mapper.image_format = "RGB"
+dataloader.train.mapper.recompute_boxes = True
+
+dataloader.test.dataset.names = _VAL_NAME
+dataloader.test.num_workers = 8
+dataloader.test.mapper.augmentations = [
+    L(T.ResizeShortestEdge)(short_edge_length=1024, max_size=1024),
+]
+
+dataloader.evaluator.dataset_name = _VAL_NAME
+dataloader.evaluator.tasks = ("bbox",)
+dataloader.evaluator.max_dets_per_image = 500
+
+model._target_ = GeneralizedRCNNWithBBoxOnlyEval
+model.roi_heads.num_classes = 44
+model.roi_heads.batch_size_per_image = 1024
+model.roi_heads.positive_fraction = 0.5
+model.roi_heads.box_predictor.test_score_thresh = 0.02
+model.roi_heads.box_predictor.test_nms_thresh = 0.7
+model.roi_heads.box_predictor.test_topk_per_image = 500
+
+model.proposal_generator.anchor_generator.sizes = [[8], [16], [32], [64], [128]]
+model.proposal_generator.batch_size_per_image = 512
+model.proposal_generator.positive_fraction = 0.5
+model.proposal_generator.pre_nms_topk = (4000, 4000)
+model.proposal_generator.post_nms_topk = (2000, 2000)
+
+optimizer.lr = 1e-4
+
+train.output_dir = "./output/oweed_tiled_keepfrag_dinov3_vitb16_frozen_mask_rcnn"
+train.max_iter = 50000
+train.eval_period = 2000
+train.checkpointer.period = 2000
+train.checkpointer.max_to_keep = 10
+train.best_checkpointer = dict(
+    val_metric="bbox/APs",
+    mode="max",
+    file_prefix="model_best_bbox_aps",
+)
+
+lr_multiplier = L(WarmupParamScheduler)(
+    scheduler=L(MultiStepParamScheduler)(
+        values=[1.0, 0.1, 0.01],
+        milestones=[35000, 45000],
+        num_updates=train.max_iter,
+    ),
+    warmup_length=500 / train.max_iter,
+    warmup_factor=0.001,
+)