Coverage for collectra/tasks/object_detection/rfdetr.py: 79.42%

1from __future__ import annotations 

2 

3import json 

4import shutil 

5from pathlib import Path 

6from typing import Annotated, TYPE_CHECKING 

7 

8from cappa import Arg 

9 

10from collectra.cli import command 

11from .adaptation_commands import RFDETRAdaptationCommands 

12from collectra.types.images import ( 

13 Image, 

14 ImageCrop, 

15 image_channel_count, 

16 image_size, 

17 read_tiff_channels, 

18) 

19from collectra.utils import change_dir 

20 

21from ...logger import get_logger 

22from ..base import Task 

23from collectra.utils import threading_locked 

24from ..machine_learning.training import ( 

25 prepare_object_detection_inputs, 

26 print_distribution_table, 

27 run_training_command, 

28) 

29from .detr import DetectionTrainResult 

30 

31if TYPE_CHECKING: 

32 from rfdetr.detr import RFDETR 

33 

34logger = get_logger(__name__) 

35 

36__all__ = ["ObjectDetectionRFDETR"] 

37 

38_MODEL_VARIANTS = { 

39 "base": "RFDETRBase", 

40 "nano": "RFDETRNano", 

41 "small": "RFDETRSmall", 

42 "medium": "RFDETRMedium", 

43 "large": "RFDETRLarge", 

44 "xlarge": "RFDETRXLarge", 

45 "2xlarge": "RFDETR2XLarge", 

46} 

47 

48 

49def __getattr__(name: str): 

50 """Expose patchable backend classes without importing them during CLI help.""" 

51 if name == "RFDETR": 

52 from rfdetr.detr import RFDETR 

53 

54 return RFDETR 

55 if name in _MODEL_VARIANTS.values(): 

56 import rfdetr 

57 

58 return getattr(rfdetr, name) 

59 raise AttributeError(name) 

60 

61 

62def _backend_class(name: str): 

63 return globals().get(name) or __getattr__(name) 

64 

65 

66def load_rfdetr_model(model_path: str | Path) -> RFDETR: 

67 from .rfdetr_channels import multichannel_checkpoints, sync_channels 

68 

69 # Fusion checkpoints retain a three-channel projection but consume many 

70 # RGB views. Ordinary RF-DETR loading would silently discard fusion weights. 

71 import torch 

72 

73 try: 

74 checkpoint = torch.load(model_path, map_location="cpu", weights_only=True) 

75 except Exception: 

76 checkpoint = {} 

77 if "backbone.0.encoder.fusion_spec" in checkpoint.get("model", {}): 

78 from .rfdetr_feature_fusion import load_adapted_model 

79 

80 return load_adapted_model(model_path) 

81 del checkpoint 

82 

83 errors = [] 

84 try: 

85 # Resolves the model size from checkpoint metadata. The loader widens 

86 # the input layer to the stored channel count; num_channels=3 stops 

87 # RF-DETR from widening it a second time. 

88 with multichannel_checkpoints(): 

89 model = _backend_class("RFDETR").from_checkpoint( 

90 str(model_path), num_channels=3 

91 ) 

92 if isinstance(model, _backend_class("RFDETR")): 

93 sync_channels(model) 

94 return model 

95 except Exception as e: 

96 errors.append(("RFDETR.from_checkpoint", str(e))) 

97 

98 # Fallback for checkpoints without metadata. RF-DETR loads non-strictly, so 

99 # a smaller variant can "load" a larger checkpoint while dropping weights. 

100 for class_name in _MODEL_VARIANTS.values(): 

101 try: 

102 model_class = _backend_class(class_name) 

103 with multichannel_checkpoints(): 

104 model = model_class(pretrain_weights=str(model_path)) 

105 sync_channels(model) 

106 return model 

107 except Exception as e: 

108 errors.append((class_name, str(e))) 

109 

110 logger.error("Failed to load RF-DETR model from %s. Errors: %s", model_path, errors) 

111 

112 

113class ObjectDetectionRFDETR(RFDETRAdaptationCommands, Task): 

114 """RF-DETR task backed by the Roboflow rfdetr package.""" 

115 

116 model: str | Path | object | None 

117 original_model_path: str | Path | None = None 

118 

119 def __init__( 

120 self, 

121 name: str, 

122 model: str | Path | object | None = None, 

123 **kwargs, 

124 ): 

125 super().__init__(name, model=model, **kwargs) 

126 self._device = None 

127 self._categories: list[str] = [] 

128 

129 def _select_device(self) -> str: 

130 import platform 

131 

132 import torch 

133 

134 return ( 

135 "mps" 

136 if platform.system() == "Darwin" 

137 else "cuda" if torch.cuda.is_available() else "cpu" 

138 ) 

139 

140 def _init_model(self) -> None: 

141 self._load() 

142 

143 if not isinstance(self.model, _backend_class("RFDETR")): 

144 raise ValueError("Model must be an RFDETR instance") 

145 

146 def _reload(self) -> None: 

147 if not self.original_model_path: 

148 logger.warning("Original model path is not set. Skipping...") 

149 return 

150 

151 self.model = load_rfdetr_model(self.original_model_path) 

152 self._load_categories_from_json(Path(self.original_model_path)) 

153 

154 def _classes_sidecar_candidates(self, checkpoint_path: Path) -> list[Path]: 

155 weights_dir = checkpoint_path.parent 

156 return [ 

157 weights_dir / f"{checkpoint_path.stem}.classes.json", 

158 weights_dir / "classes.json", 

159 ] 

160 

161 def _load_categories_from_json(self, checkpoint_path: Path) -> None: 

162 for classes_file in self._classes_sidecar_candidates(checkpoint_path): 

163 if classes_file.exists(): 

164 with open(classes_file, "r") as f: 

165 metadata = json.load(f) 

166 self._categories = metadata.get("classes", self._categories) 

167 return 

168 

169 def _load(self) -> None: 

170 if isinstance(self.model, _backend_class("RFDETR")): 

171 return 

172 

173 self._device = self._select_device() 

174 

175 variant = "base" if self.model is None else None 

176 if isinstance(self.model, str): 

177 name = self.model.lower() 

178 if name == "default": 

179 variant = "base" 

180 elif name in _MODEL_VARIANTS: 

181 variant = name 

182 else: 

183 variant = next( 

184 ( 

185 key 

186 for key, value in _MODEL_VARIANTS.items() 

187 if value.lower() == name 

188 ), 

189 None, 

190 ) 

191 if variant is not None: 

192 self.model = _backend_class(_MODEL_VARIANTS[variant])() 

193 self.original_model_path = None 

194 from rfdetr.assets.coco_classes import COCO_CLASSES 

195 

196 self._categories = list(COCO_CLASSES) 

197 return 

198 

199 if not isinstance(self.model, (str, Path)): 

200 raise ValueError( 

201 "Model must be None, 'default', a model variant name, " 

202 "a checkpoint path, or an RFDETR instance" 

203 ) 

204 

205 checkpoint_path = Path(str(self.model)) 

206 if checkpoint_path.exists() and checkpoint_path.is_file(): 

207 self.original_model_path = checkpoint_path 

208 self.model = load_rfdetr_model(self.original_model_path) 

209 self._categories = [] 

210 self._load_categories_from_json(checkpoint_path) 

211 candidates = self._classes_sidecar_candidates(checkpoint_path) 

212 if not any(c.exists() for c in candidates): 

213 logger.warning( 

214 "No classes sidecar found alongside checkpoint %s " 

215 "(looked for %s); inference labels will fall back to " 

216 "numeric class indices because no class metadata is available.", 

217 checkpoint_path, 

218 [str(c) for c in candidates], 

219 ) 

220 return 

221 

222 raise ValueError( 

223 f"Model path does not exist: {checkpoint_path}. " 

224 "RF-DETR requires a valid checkpoint path, None/'default', " 

225 f"or a model variant: {', '.join(_MODEL_VARIANTS)}." 

226 ) 

227 

228 def _tiff_tensor(self, image: Image): 

229 """Read every channel of a TIFF image as a CHW float tensor in [0, 1].""" 

230 import numpy as np 

231 import torch 

232 

233 from .rfdetr_data import scale_tiff_pixels 

234 

235 pixels = read_tiff_channels(image.get_path()) 

236 if isinstance(image, ImageCrop): 

237 left, top, right, bottom = image.coordinates() 

238 if right <= left or bottom <= top: 

239 raise ValueError(f"ImageCrop {image.id!r} produces an empty pixel crop") 

240 pixels = pixels[top:bottom, left:right, :] 

241 pixels = np.rot90(pixels, k=image.orientation.to_degree() // 90) 

242 

243 expected = self.model.model_config.num_channels 

244 if pixels.shape[2] != expected: 

245 raise ValueError( 

246 f"{image.get_path()}: {pixels.shape[2]} image channels, " 

247 f"but the detection model expects {expected}" 

248 ) 

249 pixels = scale_tiff_pixels(pixels) 

250 return torch.from_numpy(np.ascontiguousarray(pixels.transpose(2, 0, 1))) 

251 

252 @threading_locked() 

253 def run(self, *args: Image) -> list[Image]: 

254 if len(args) != 1: 

255 raise ValueError("This task only supports a single Image input.") 

256 if not isinstance(args[0], Image): 

257 raise TypeError("Input must be an instance of Image.") 

258 

259 image = args[0] 

260 self._init_model() 

261 

262 threshold = float(getattr(self, "threshold", 0.5)) 

263 if image.get_path().suffix.lower() in {".tif", ".tiff"}: 

264 model_input = self._tiff_tensor(image) 

265 height, width = model_input.shape[1:] 

266 else: 

267 model_input = image.pil().convert("RGB") 

268 width, height = model_input.size 

269 

270 detections = self.model.predict( 

271 model_input, threshold=threshold, include_source_image=False 

272 ) 

273 

274 results: list[Image] = [] 

275 for i in range(len(detections.xyxy)): 

276 x1, y1, x2, y2 = detections.xyxy[i] 

277 x1 = max(0.0, min(float(width), float(x1))) 

278 x2 = max(0.0, min(float(width), float(x2))) 

279 y1 = max(0.0, min(float(height), float(y1))) 

280 y2 = max(0.0, min(float(height), float(y2))) 

281 

282 bbox_width = x2 - x1 

283 bbox_height = y2 - y1 

284 if bbox_width <= 0 or bbox_height <= 0: 

285 continue 

286 

287 x_center = (x1 + x2) / 2.0 / width 

288 y_center = (y1 + y2) / 2.0 / height 

289 width_relative = bbox_width / width 

290 height_relative = bbox_height / height 

291 

292 class_idx = int(detections.class_id[i]) 

293 class_name = ( 

294 self._categories[class_idx] 

295 if 0 <= class_idx < len(self._categories) 

296 else str(class_idx) 

297 ) 

298 

299 results.append( 

300 image.make_crop( 

301 x_center=x_center, 

302 y_center=y_center, 

303 width_relative=width_relative, 

304 height_relative=height_relative, 

305 orientation=image.orientation, 

306 name=class_name, 

307 ) 

308 ) 

309 

310 print(f"Found {len(results)} objects in the image.") 

311 return results 

312 

313 def _prepare_params(self, **kwargs) -> dict: 

314 params = { 

315 "epochs": int(kwargs.get("epochs", 1)), 

316 "batch_size": int(kwargs.get("batch", 4)), 

317 "grad_accum_steps": int(kwargs.get("grad_accum_steps", 4)), 

318 "lr": float(kwargs.get("lr", 1e-4)), 

319 "lr_encoder": float(kwargs.get("lr_encoder", 1.5e-4)), 

320 "lr_vit_layer_decay": float(kwargs.get("lr_vit_layer_decay", 0.8)), 

321 "lr_component_decay": float(kwargs.get("lr_component_decay", 0.7)), 

322 "warmup_epochs": float(kwargs.get("warmup_epochs", 0.0)), 

323 "weight_decay": float(kwargs.get("weight_decay", 1e-4)), 

324 "drop_path": float(kwargs.get("drop_path", 0.0)), 

325 "multi_scale": bool(kwargs.get("multi_scale", True)), 

326 "expanded_scales": bool(kwargs.get("expanded_scales", True)), 

327 "do_random_resize_via_padding": bool( 

328 kwargs.get("do_random_resize_via_padding", False) 

329 ), 

330 "augmentation_backend": kwargs.get("augmentation_backend", "cpu"), 

331 "use_ema": bool(kwargs.get("use_ema", True)), 

332 "ema_decay": float(kwargs.get("ema_decay", 0.993)), 

333 "ema_tau": int(kwargs.get("ema_tau", 100)), 

334 "ema_update_interval": int(kwargs.get("ema_update_interval", 1)), 

335 "output_dir": str(kwargs["output_dir"]), 

336 "wandb": bool(kwargs.get("wandb", False)), 

337 "project": kwargs.get("project", "runs/rfdetr"), 

338 "run": kwargs.get("log", "rfdetr-run"), 

339 "early_stopping": bool(kwargs.get("early_stopping", False)), 

340 "early_stopping_patience": int(kwargs.get("early_stopping_patience", 10)), 

341 "early_stopping_min_delta": float( 

342 kwargs.get("early_stopping_min_delta", 0.001) 

343 ), 

344 "early_stopping_use_ema": bool(kwargs.get("early_stopping_use_ema", False)), 

345 } 

346 

347 resolution = kwargs.get("resolution") 

348 if resolution is not None: 

349 params["resolution"] = int(resolution) 

350 

351 augmentation = kwargs.get("augmentation", "default").lower() 

352 if augmentation != "default": 

353 preset_names = { 

354 "none", 

355 "conservative", 

356 "aggressive", 

357 "aerial", 

358 "industrial", 

359 } 

360 if augmentation not in preset_names: 

361 raise ValueError( 

362 "augmentation must be one of: default, none, conservative, " 

363 "aggressive, aerial, industrial" 

364 ) 

365 from rfdetr.datasets.aug_configs import ( 

366 AUG_AERIAL, 

367 AUG_AGGRESSIVE, 

368 AUG_CONSERVATIVE, 

369 AUG_INDUSTRIAL, 

370 ) 

371 

372 presets = { 

373 "none": {}, 

374 "conservative": AUG_CONSERVATIVE, 

375 "aggressive": AUG_AGGRESSIVE, 

376 "aerial": AUG_AERIAL, 

377 "industrial": AUG_INDUSTRIAL, 

378 } 

379 params["aug_config"] = presets[augmentation] 

380 return params 

381 

382 def _check_distribution( 

383 self, 

384 classes: list[str], 

385 train_samples: list[dict], 

386 val_samples: list[dict], 

387 ) -> None: 

388 train_classes = {name: 0 for name in classes} 

389 val_classes = {name: 0 for name in classes} 

390 label_to_name = {index: name for index, name in enumerate(classes)} 

391 

392 for sample in train_samples: 

393 for label in sample["labels"]: 

394 class_name = label_to_name.get(label, "") 

395 if class_name: 

396 train_classes[class_name] += 1 

397 

398 for sample in val_samples: 

399 for label in sample["labels"]: 

400 class_name = label_to_name.get(label, "") 

401 if class_name: 

402 val_classes[class_name] += 1 

403 

404 print_distribution_table( 

405 "Class Distribution", classes, train_classes, val_classes 

406 ) 

407 

408 def _prepare_assets( 

409 self, 

410 classes: list[str], 

411 validation_flag: str, 

412 exclude_flag: str, 

413 *images: ImageCrop, 

414 ) -> tuple[list[dict], list[dict]]: 

415 class_to_idx = {name: idx for idx, name in enumerate(classes)} 

416 grouped: dict[Path, dict] = {} 

417 

418 for img in images: 

419 if exclude_flag and img.partition == exclude_flag: 

420 continue 

421 

422 base_img: Image | ImageCrop = img 

423 if isinstance(img, ImageCrop): 

424 if img.source_parent: 

425 base_img = img.source_parent 

426 if isinstance(img.source_parent, ImageCrop): 

427 img.set_rel_to_src_parent() 

428 

429 image_path = base_img.get_path() 

430 if image_path not in grouped: 

431 grouped[image_path] = { 

432 "image_path": image_path, 

433 "partition": img.partition, 

434 "boxes": [], 

435 "labels": [], 

436 } 

437 

438 if not isinstance(img, ImageCrop): 

439 continue 

440 if img.name not in class_to_idx: 

441 continue 

442 

443 width, height = base_img.width, base_img.height 

444 box_w = float(img.width_relative) * width 

445 box_h = float(img.height_relative) * height 

446 cx = float(img.x_center) * width 

447 cy = float(img.y_center) * height 

448 

449 x1 = max(0.0, cx - box_w / 2.0) 

450 y1 = max(0.0, cy - box_h / 2.0) 

451 x2 = min(float(width), cx + box_w / 2.0) 

452 y2 = min(float(height), cy + box_h / 2.0) 

453 if x2 <= x1 or y2 <= y1: 

454 continue 

455 

456 grouped[image_path]["boxes"].append([x1, y1, x2, y2]) 

457 grouped[image_path]["labels"].append(class_to_idx[img.name]) 

458 

459 train_samples: list[dict] = [] 

460 val_samples: list[dict] = [] 

461 for sample in grouped.values(): 

462 if validation_flag and sample["partition"] == validation_flag: 

463 val_samples.append(sample) 

464 elif not exclude_flag or sample["partition"] != exclude_flag: 

465 train_samples.append(sample) 

466 return train_samples, val_samples 

467 

468 def _write_coco_dataset( 

469 self, 

470 train_samples: list[dict], 

471 val_samples: list[dict], 

472 classes: list[str], 

473 dataset_dir: Path, 

474 ) -> Path: 

475 categories = [ 

476 {"id": idx, "name": name, "supercategory": "object"} 

477 for idx, name in enumerate(classes) 

478 ] 

479 

480 channel_counts: dict[str, int] = {} 

481 for split_name, samples in [("train", train_samples), ("valid", val_samples)]: 

482 split_dir = dataset_dir / split_name 

483 split_dir.mkdir(parents=True, exist_ok=True) 

484 

485 images_list: list[dict] = [] 

486 annotations_list: list[dict] = [] 

487 annotation_id = 0 

488 

489 for image_id, sample in enumerate(samples): 

490 src_path = Path(sample["image_path"]) 

491 dest_name = f"{image_id}_{src_path.name}" 

492 dest_path = split_dir / dest_name 

493 shutil.copy2(src_path, dest_path) 

494 

495 w, h = image_size(dest_path) 

496 channel_counts[str(src_path)] = image_channel_count(dest_path) 

497 

498 images_list.append( 

499 { 

500 "id": image_id, 

501 "file_name": dest_name, 

502 "width": w, 

503 "height": h, 

504 } 

505 ) 

506 

507 for box, label in zip(sample["boxes"], sample["labels"]): 

508 x1, y1, x2, y2 = box 

509 coco_bbox = [x1, y1, x2 - x1, y2 - y1] 

510 area = (x2 - x1) * (y2 - y1) 

511 

512 annotations_list.append( 

513 { 

514 "id": annotation_id, 

515 "image_id": image_id, 

516 "category_id": int(label), 

517 "bbox": coco_bbox, 

518 "area": float(area), 

519 "iscrowd": 0, 

520 } 

521 ) 

522 annotation_id += 1 

523 

524 coco_json = { 

525 "images": images_list, 

526 "annotations": annotations_list, 

527 "categories": categories, 

528 } 

529 

530 annotations_path = split_dir / "_annotations.coco.json" 

531 with open(annotations_path, "w") as f: 

532 json.dump(coco_json, f, indent=2) 

533 

534 if len(set(channel_counts.values())) > 1: 

535 details = ", ".join( 

536 f"{name}: {count} channels" for name, count in channel_counts.items() 

537 ) 

538 raise ValueError( 

539 f"Inconsistent input channel counts in prepared dataset: {details}" 

540 ) 

541 if channel_counts: 

542 logger.info( 

543 "RF-DETR dataset images have %d channels", 

544 next(iter(channel_counts.values())), 

545 ) 

546 return dataset_dir 

547 

548 def _train_fold( 

549 self, 

550 train_samples: list[dict], 

551 val_samples: list[dict], 

552 classes: list[str], 

553 log: Path, 

554 kwargs: dict, 

555 ) -> DetectionTrainResult: 

556 

557 weights_dir = log / "weights" 

558 weights_dir.mkdir(parents=True, exist_ok=True) 

559 

560 dataset_dir = log / "dataset" 

561 dataset_dir.mkdir(parents=True, exist_ok=True) 

562 self._write_coco_dataset(train_samples, val_samples, classes, dataset_dir) 

563 

564 params_kwargs = dict(kwargs) 

565 params_kwargs["log"] = f"{log.name}" 

566 params_kwargs["output_dir"] = str(weights_dir) 

567 params = self._prepare_params(**params_kwargs) 

568 has_tiffs = any( 

569 Path(sample["image_path"]).suffix.lower() in {".tif", ".tiff"} 

570 for sample in [*train_samples, *val_samples] 

571 ) 

572 if has_tiffs: 

573 from .rfdetr_channels import require_multichannel_rfdetr 

574 

575 require_multichannel_rfdetr() 

576 if has_tiffs and params["augmentation_backend"] != "torchvision": 

577 # Albumentations round-trips through PIL; the GPU path uses 

578 # three-channel normalization outside our dataset patch. "auto" 

579 # can select either path, so TIFFs must explicitly use torchvision. 

580 logger.info("TIFF inputs: using the torchvision augmentation backend") 

581 params["augmentation_backend"] = "torchvision" 

582 

583 from .rfdetr_channels import multichannel_checkpoints 

584 from .rfdetr_data import multichannel_datasets 

585 

586 model = self.model 

587 with multichannel_datasets(), multichannel_checkpoints(): 

588 model.train( 

589 dataset_dir=str(dataset_dir), 

590 **params, 

591 ) 

592 

593 self._categories = classes 

594 self.model = model 

595 

596 # Persist class names alongside checkpoint for later reloading 

597 classes_metadata = {"classes": classes} 

598 with open(weights_dir / "classes.json", "w") as f: 

599 json.dump(classes_metadata, f, indent=2) 

600 

601 # RF-DETR's total checkpoint selects the validation winner across 

602 # regular and EMA weights. Prefer it so best.pt keeps that winner. 

603 best_pt = weights_dir / "best.pt" 

604 for candidate in [ 

605 weights_dir / "checkpoint_best_total.pth", 

606 weights_dir / "checkpoint_best_regular.pth", 

607 ]: 

608 if candidate.exists(): 

609 shutil.copy2(candidate, best_pt) 

610 break 

611 

612 if best_pt.exists(): 

613 self.original_model_path = best_pt 

614 self._reload() 

615 

616 best_metrics = {"status": "completed", "classes": classes} 

617 

618 logger.info("RF-DETR training completed. Weights saved to: %s", weights_dir) 

619 return DetectionTrainResult(save_dir=log, results_dict=best_metrics) 

620 

621 @command 

622 def train( 

623 self, 

624 inputs: Annotated[ 

625 list[str], 

626 Arg( 

627 help="Pipeline data files or directories containing training annotations." 

628 ), 

629 ], 

630 output: Annotated[ 

631 Path | None, 

632 Arg(help="Directory for training logs, dataset files, and checkpoints."), 

633 ] = None, 

634 keep_log: Annotated[ 

635 bool, 

636 Arg( 

637 help="Keep the training directory after saving the model to the pipeline." 

638 ), 

639 ] = True, 

640 validation: Annotated[ 

641 str, 

642 Arg(help="Partition value identifying validation examples."), 

643 ] = "", 

644 exclude: Annotated[ 

645 str, 

646 Arg(help="Partition value identifying examples to exclude from training."), 

647 ] = "", 

648 epochs: Annotated[ 

649 int, 

650 Arg(help="Number of training epochs."), 

651 ] = 1, 

652 batch: Annotated[ 

653 int, 

654 Arg(help="Number of images per device in each training batch."), 

655 ] = 4, 

656 grad_accum_steps: Annotated[ 

657 int, 

658 Arg( 

659 help="Accumulate gradients over this many batches before an optimizer step." 

660 ), 

661 ] = 4, 

662 learning_rate: Annotated[ 

663 float, 

664 Arg(help="Learning rate for training."), 

665 ] = 1e-4, 

666 wandb: Annotated[ 

667 bool, 

668 Arg(help="Log training metrics to Weights & Biases."), 

669 ] = False, 

670 early_stopping: Annotated[ 

671 bool, 

672 Arg(help="Stop training when validation performance stops improving."), 

673 ] = False, 

674 weight_decay: Annotated[ 

675 float, 

676 Arg(help="Weight decay used for regularization."), 

677 ] = 1e-4, 

678 drop_path: Annotated[ 

679 float, 

680 Arg(help="Stochastic depth drop probability used for regularization."), 

681 ] = 0.0, 

682 augmentation: Annotated[ 

683 str, 

684 Arg( 

685 help="Augmentation preset: default, none, conservative, aggressive, aerial, or industrial." 

686 ), 

687 ] = "default", 

688 augmentation_backend: Annotated[ 

689 str, 

690 Arg(help="Augmentation execution backend: cpu, auto, or gpu."), 

691 ] = "cpu", 

692 multi_scale: Annotated[ 

693 bool, 

694 Arg(help="Train with multiple image resolutions."), 

695 ] = True, 

696 expanded_scales: Annotated[ 

697 bool, 

698 Arg(help="Use an expanded range of resolutions for multi-scale training."), 

699 ] = True, 

700 do_random_resize_via_padding: Annotated[ 

701 bool, 

702 Arg(help="Use padding for random image resizing."), 

703 ] = False, 

704 use_ema: Annotated[ 

705 bool, 

706 Arg(help="Maintain an exponential moving average (EMA) of model weights."), 

707 ] = True, 

708 ema_decay: Annotated[ 

709 float, 

710 Arg( 

711 help="Decay factor for the exponential moving average of model weights." 

712 ), 

713 ] = 0.993, 

714 ema_tau: Annotated[ 

715 int, 

716 Arg(help="Time constant controlling the warmup of EMA decay."), 

717 ] = 100, 

718 ema_update_interval: Annotated[ 

719 int, 

720 Arg(help="Number of optimizer steps between EMA updates."), 

721 ] = 1, 

722 early_stopping_patience: Annotated[ 

723 int, 

724 Arg( 

725 help="Number of validation checks without sufficient improvement before stopping." 

726 ), 

727 ] = 10, 

728 early_stopping_min_delta: Annotated[ 

729 float, 

730 Arg(help="Minimum improvement required to reset early-stopping patience."), 

731 ] = 0.001, 

732 early_stopping_use_ema: Annotated[ 

733 bool, 

734 Arg(help="Use EMA model metrics for early stopping."), 

735 ] = False, 

736 lr_encoder: Annotated[ 

737 float, 

738 Arg(help="Learning rate for the image encoder."), 

739 ] = 1.5e-4, 

740 lr_vit_layer_decay: Annotated[ 

741 float, 

742 Arg( 

743 help="Layer-wise learning-rate decay factor for the vision transformer." 

744 ), 

745 ] = 0.8, 

746 lr_component_decay: Annotated[ 

747 float, 

748 Arg(help="Learning-rate decay factor for encoder components."), 

749 ] = 0.7, 

750 warmup_epochs: Annotated[ 

751 float, 

752 Arg( 

753 help="Number of epochs for learning-rate warmup; fractional values are allowed." 

754 ), 

755 ] = 0.0, 

756 model: Annotated[ 

757 str, 

758 Arg( 

759 help="Checkpoint path or pretrained variant: nano, small, base, medium, large, xlarge, or 2xlarge (full RFDETR class names also accepted). XL variants require rfdetr[plus]. 'default' selects Base; omit to keep the configured model." 

760 ), 

761 ] = "", 

762 resolution: Annotated[ 

763 float | None, 

764 Arg(help="Change default resolution. Valid options will depend on the model type."), 

765 ] = None, 

766 max_items: Annotated[ 

767 int, 

768 Arg( 

769 help="Use at most this many images in each of the training and validation sets (0 uses all); for quick tests." 

770 ), 

771 ] = 0, 

772 ): 

773 """Train this RF-DETR object detector.""" 

774 return run_training_command( 

775 self, 

776 inputs, 

777 self._train, 

778 keep_log=keep_log, 

779 output=output, 

780 validation=validation, 

781 exclude=exclude, 

782 prepare_inputs=prepare_object_detection_inputs, 

783 include_unlabelled=True, 

784 epochs=epochs, 

785 batch=batch, 

786 grad_accum_steps=grad_accum_steps, 

787 lr=learning_rate, 

788 wandb=wandb, 

789 early_stopping=early_stopping, 

790 weight_decay=weight_decay, 

791 drop_path=drop_path, 

792 augmentation=augmentation, 

793 augmentation_backend=augmentation_backend,