Coverage for collectra/tasks/object_detection/rfdetr_input_adaptation.py: 79.52%
83 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-16 09:56 +0000
1"""Adapt RGB RF-DETR checkpoints to concatenated aligned views."""
3from pathlib import Path
5import torch
8def _patch_embeddings(detector):
9 try:
10 embeddings = detector.model.model.backbone[
11 0
12 ].encoder.encoder.embeddings.patch_embeddings
13 except (AttributeError, IndexError, TypeError) as exc:
14 raise RuntimeError(
15 "Could not locate RF-DETR's DINOv2 patch embeddings."
16 ) from exc
17 if not isinstance(embeddings.projection, torch.nn.Conv2d):
18 raise RuntimeError("Expected the patch projection to be torch.nn.Conv2d.")
19 return embeddings
22def _expand_model(detector, num_views: int) -> None:
23 if num_views < 1:
24 raise ValueError("num_views must be >= 1")
25 embeddings = _patch_embeddings(detector)
26 old_conv = embeddings.projection
27 if old_conv.in_channels != 3:
28 raise ValueError(
29 "Expected a 3-channel pretrained model, but the patch projection "
30 f"has {old_conv.in_channels} input channels."
31 )
32 if old_conv.groups != 1:
33 raise RuntimeError("Grouped input convolutions are not supported.")
35 channels = 3 * num_views
36 new_conv = torch.nn.Conv2d(
37 channels,
38 old_conv.out_channels,
39 kernel_size=old_conv.kernel_size,
40 stride=old_conv.stride,
41 padding=old_conv.padding,
42 dilation=old_conv.dilation,
43 groups=old_conv.groups,
44 bias=old_conv.bias is not None,
45 padding_mode=old_conv.padding_mode,
46 ).to(device=old_conv.weight.device, dtype=old_conv.weight.dtype)
47 with torch.no_grad():
48 new_conv.weight.copy_(old_conv.weight.repeat(1, num_views, 1, 1) / num_views)
49 if old_conv.bias is not None:
50 new_conv.bias.copy_(old_conv.bias)
51 new_conv.weight.requires_grad_(old_conv.weight.requires_grad)
52 if old_conv.bias is not None:
53 new_conv.bias.requires_grad_(old_conv.bias.requires_grad)
54 new_conv.train(old_conv.training)
55 embeddings.projection = new_conv
56 # DINOv2 checks this attribute before calling the convolution.
57 embeddings.num_channels = channels
58 detector.model.model.backbone[0].encoder.encoder.config.num_channels = channels
59 detector.model_config.num_channels = channels
60 detector.model.args.num_channels = channels
61 detector.means = list(detector.means[:3]) * num_views
62 detector.stds = list(detector.stds[:3]) * num_views
65def expand_rfdetr_input_channels(
66 input_model: Path,
67 output_model: Path,
68 num_views: int,
69 variant: str | None = None,
70) -> None:
71 """Save expanded model weights and reconstruction metadata (no optimizer)."""
72 if num_views < 1:
73 raise ValueError("num_views must be >= 1")
74 import rfdetr
76 if variant is None:
77 try:
78 detector = rfdetr.RFDETR.from_checkpoint(str(input_model), device="cpu")
79 except (KeyError, ValueError) as exc:
80 raise ValueError(
81 "Could not reconstruct the input checkpoint automatically. "
82 "Use --variant with its model size, e.g. --variant nano. "
83 f"RF-DETR reported: {exc}"
84 ) from exc
85 else:
86 variants = {
87 "nano": "RFDETRNano",
88 "small": "RFDETRSmall",
89 "medium": "RFDETRMedium",
90 "base": "RFDETRBase",
91 "large": "RFDETRLarge",
92 "xlarge": "RFDETRXLarge",
93 "2xlarge": "RFDETR2XLarge",
94 }
95 name = variants.get(variant.lower(), variant)
96 model_class = getattr(rfdetr, name, None)
97 if model_class is None or not name.startswith("RFDETR"):
98 raise ValueError(f"Unknown or unavailable RF-DETR variant: {variant}")
99 detector = model_class(pretrain_weights=str(input_model), device="cpu")
101 _expand_model(detector, num_views)
102 config = detector.model_config.model_dump(mode="json")
103 args = dict(config)
104 args["class_names"] = list(detector.model.class_names or [])
105 checkpoint = {
106 "model": detector.model.model.state_dict(),
107 "model_name": type(detector).__name__,
108 "model_config": config,
109 "args": args,
110 "multichannel": {
111 "num_views": num_views,
112 "num_channels": 3 * num_views,
113 "means": detector.means,
114 "stds": detector.stds,
115 },
116 }
117 output_model.parent.mkdir(parents=True, exist_ok=True)
118 torch.save(checkpoint, output_model)
121def load_adapted_model(output_model: Path, device: str = "cpu"):
122 """Reconstruct a checkpoint produced here, expanding before loading weights.
124 Returns the RF-DETR wrapper, including repeated RGB normalization statistics.
125 Use CHW tensors with 3 * num_views channels for prediction.
126 """
127 import rfdetr
129 checkpoint = torch.load(output_model, map_location="cpu", weights_only=True)
130 metadata = checkpoint["multichannel"]
131 num_views = metadata["num_views"]
132 if num_views < 1 or metadata["num_channels"] != 3 * num_views:
133 raise ValueError("Inconsistent multichannel checkpoint metadata.")
134 config = dict(checkpoint["model_config"])
135 # Avoid the upstream order: loading expanded weights into an RGB layer.
136 config.update(pretrain_weights=None, num_channels=3, device=device)
137 model_class = getattr(rfdetr, checkpoint["model_name"])
138 detector = model_class(**config)
139 _expand_model(detector, num_views)
140 detector.model.model.load_state_dict(checkpoint["model"], strict=True)
141 detector.model.class_names = checkpoint["args"].get("class_names", [])
142 detector.means = list(metadata["means"])
143 detector.stds = list(metadata["stds"])
144 return detector
147def validate_saved_model(output_model: Path, expected_channels: int) -> None:
148 reloaded = load_adapted_model(output_model)
149 embeddings = _patch_embeddings(reloaded)
150 conv = embeddings.projection
151 if conv.in_channels != expected_channels:
152 raise RuntimeError(
153 f"Reloaded model has {conv.in_channels} input channels; "
154 f"expected {expected_channels}."
155 )
156 # Exercise the channel guard as well as the convolution after strict reload.
157 with torch.no_grad():
158 embeddings(torch.zeros(1, expected_channels, *conv.kernel_size))