# ------------------------------------------------------------------------ # RF-DETR # Copyright (c) 2025 Roboflow. All Rights Reserved. # Licensed under the Apache License, Version 2.0 [see LICENSE for details] # ------------------------------------------------------------------------ """LightningDataModule for RF-DETR dataset construction and loaders.""" from __future__ import annotations from pathlib import Path from typing import TYPE_CHECKING, Any, Literal import torch import torch.utils.data from pytorch_lightning import LightningDataModule from torch.utils.data import DataLoader from rfdetr._namespace import _namespace_from_configs from rfdetr.config import AugmentationBackend, ModelConfig, TrainConfig from rfdetr.datasets import build_dataset from rfdetr.datasets.aug_configs import AUG_CONFIG from rfdetr.datasets.yolo import YoloSplitUnavailableError from rfdetr.utilities.box_ops import box_xyxy_to_cxcywh from rfdetr.utilities.logger import get_logger from rfdetr.utilities.tensors import make_collate_fn logger = get_logger() _MIN_TRAIN_BATCHES = 5 if TYPE_CHECKING: from matplotlib.figure import Figure def _worker_init_fn(worker_id: int) -> None: """Seed NumPy and the ``random`` module per DataLoader worker. PyTorch seeds ``torch``'s RNG per worker automatically but leaves NumPy and the stdlib ``random`` module unseeded, so without this hook every worker would draw identical NumPy/``random`` sequences (a well-known augmentation duplication footgun). Deriving the seed from ``torch.initial_seed()`` keeps augmentation reproducible while still giving each worker a distinct stream. Args: worker_id: Index of the DataLoader worker (unused; seed is derived from the per-worker torch seed). """ import random import numpy as np seed = torch.initial_seed() % (2**32) np.random.seed(seed) random.seed(seed) def _has_cuda_device() -> bool: """Return ``True`` when the runtime has a CUDA accelerator available. Uses the fork-safe global ``DEVICE`` constant instead of direct ``torch.cuda.is_available()`` calls to avoid creating a CUDA context in fork-based notebook/DDP workflows. """ from rfdetr.config import DEVICE return str(DEVICE).startswith("cuda") class GradAccumAlignedDataset(torch.utils.data.Dataset[Any]): """Dataset wrapper that pads length to a multiple of ``effective_batch_size * world_size``. Workaround for https://github.com/Lightning-AI/pytorch-lightning/issues/19987: PTL fires the optimizer on partial accumulation windows at the tail of the dataset, causing the last optimizer step to be under-scaled. Padding the dataset to a multiple of ``effective_batch_size * world_size`` ensures that ``drop_last=True`` on the DataLoader becomes a true no-op — every accumulation window is always complete. Padding indices are drawn randomly from the original dataset. Because RF-DETR uses online augmentation, each padded sample receives a fresh random augmentation at ``__getitem__`` time, so it behaves like a new training example rather than a true duplicate. This wrapper can be removed once the upstream PTL issue is resolved. Args: dataset: The underlying dataset to wrap. effective_batch_size: ``batch_size * grad_accum_steps``. world_size: Number of DDP processes (default 1 for single-GPU/CPU). The alignment unit is ``effective_batch_size * world_size`` so that after PTL's ``DistributedSampler`` splits samples across ranks each rank still receives an exact multiple of ``effective_batch_size``. """ def __init__( self, dataset: torch.utils.data.Dataset[Any], effective_batch_size: int, world_size: int = 1, ) -> None: if effective_batch_size < 1: raise ValueError(f"effective_batch_size must be >= 1, got {effective_batch_size}") if world_size < 1: raise ValueError(f"world_size must be >= 1, got {world_size}") self._dataset = dataset self._dataset_length = len(dataset) # type: ignore[arg-type] pad_unit = effective_batch_size * world_size remainder = self._dataset_length % pad_unit pad_count = (pad_unit - remainder) % pad_unit pad_index_generator = torch.Generator() pad_index_generator.manual_seed(0) self._pad_indices: list[int] = ( torch.randint( 0, self._dataset_length, (pad_count,), generator=pad_index_generator, ).tolist() if pad_count > 0 else [] ) self._length = self._dataset_length + pad_count def __len__(self) -> int: """Return the padded dataset length (always a multiple of the alignment unit).""" return self._length def __getitem__(self, idx: int) -> Any: """Return the item at the (possibly remapped) index.""" # pad_indices are fixed at __init__ time; same indices reused every epoch # (different augmentations per epoch due to online augmentation) dataset_idx = idx if idx < self._dataset_length else self._pad_indices[idx - self._dataset_length] return self._dataset[dataset_idx] class RFDETRDataModule(LightningDataModule): """LightningDataModule wrapping RF-DETR dataset construction and data loading. Args: model_config: Architecture configuration (used for resolution, patch_size, etc.). train_config: Training hyperparameter configuration (used for dataset params). """ def __init__(self, model_config: ModelConfig, train_config: TrainConfig) -> None: super().__init__() self.model_config = model_config self.train_config = train_config # Backbone divisibility requirement: inputs with windowed attention must # have H and W divisible by patch_size * num_windows. The collate_fn # below rounds batch-max H/W up to this value so the mask accurately # marks every pad pixel. block_size = model_config.patch_size * model_config.num_windows if block_size <= 0: raise ValueError( "Computed collate block_size must be > 0, got " f"{block_size} from patch_size={model_config.patch_size} " f"and num_windows={model_config.num_windows}." ) self._collate_fn = make_collate_fn( block_size=block_size, ) self._dataset_train: torch.utils.data.Dataset[Any] | None = None self._dataset_val: torch.utils.data.Dataset[Any] | None = None self._dataset_test: torch.utils.data.Dataset[Any] | None = None # GPU augmentation pipeline (Kornia); built lazily in setup("fit"). self._kornia_pipeline: Any | None = None self._kornia_normalize: Any | None = None # Resolved backend for the pending/most recent _setup_kornia_pipeline() call; set by # setup("fit") just before that call so _setup_kornia_pipeline can stay a zero-arg method. self._resolved_augmentation_backend: AugmentationBackend | None = None # Sentinel: True once _setup_kornia_pipeline has run (even on fallback paths # where _kornia_pipeline stays None), preventing redundant re-runs on repeated # setup("fit") calls (e.g. during validation loops in some PTL strategies). self._kornia_setup_done: bool = False self._num_workers: int = self.train_config.num_workers # Use the fork-safe DEVICE constant instead of torch.cuda.is_available(), # which creates a CUDA driver context that breaks fork-based DDP. from rfdetr.config import DEVICE accelerator = str(self.train_config.accelerator).lower() uses_cuda_accelerator = accelerator in {"auto", "gpu", "cuda"} self._pin_memory: bool = ( (DEVICE == "cuda" and uses_cuda_accelerator) if self.train_config.pin_memory is None else bool(self.train_config.pin_memory) ) self._persistent_workers: bool = ( self._num_workers > 0 if self.train_config.persistent_workers is None else bool(self.train_config.persistent_workers) ) self._prefetch_factor: int | None if self._num_workers > 0: self._prefetch_factor = ( self.train_config.prefetch_factor if self.train_config.prefetch_factor is not None else 2 ) else: self._prefetch_factor = None # ------------------------------------------------------------------ # PTL lifecycle hooks # ------------------------------------------------------------------ def setup(self, stage: str) -> None: """Build datasets for the requested stage. PTL calls this on every process before the corresponding dataloader method. Datasets are built lazily — a dataset is only constructed once even if ``setup`` is called multiple times. Args: stage: PTL stage identifier — one of ``"fit"``, ``"validate"``, ``"test"``, or ``"predict"``. """ resolution = self.model_config.resolution ns = _namespace_from_configs(self.model_config, self.train_config) if stage == "fit": requested_backend = self.train_config.augmentation_backend # Keypoint transforms are incompatible with the Kornia GPU pipeline outright, so an # explicit 'kornia'/'gpu' request is rejected here -- before the CUDA readiness check # below -- so a keypoint model without CUDA still sees the keypoint error, not an # unrelated "no CUDA" one. if self.model_config.use_grouppose_keypoints and requested_backend in ( AugmentationBackend.KORNIA, "kornia", "gpu", ): raise ValueError( f"augmentation_backend={requested_backend!r} does not support keypoint transforms. " "Set augmentation_backend='cpu' or 'albumentations' when use_grouppose_keypoints=True." ) from rfdetr.datasets.kornia_transforms import ( is_gpu_postprocess, require_gpu_backend_ready, resolve_augmentation_backend, ) require_gpu_backend_ready(requested_backend, has_cuda=_has_cuda_device()) # Resolve 'auto' to an actual backend before building datasets so that # gpu_postprocess in dataset builders always matches what the DataModule # will actually do in on_after_batch_transfer. Without this, 'auto' on # a machine without CUDA/kornia would strip CPU Normalize from datasets # while _kornia_pipeline stays None, leaving training inputs unnormalized. resolved = resolve_augmentation_backend(requested_backend, has_cuda=_has_cuda_device()) if resolved != requested_backend: # 'auto'/'cpu' (and legacy aliases) were resolved above only to decide the # DataModule-side Kornia pipeline and AUG_CONFIG injection below. Dataset # builders do their own environment-aware CPU sub-backend pick (Albumentations # vs. torchvision), so forward the simplified sentinel -- not e.g. the specific # TV/ALBU pick -- to stay consistent with the gpu_postprocess flag they compute. ns.augmentation_backend = "kornia" if is_gpu_postprocess(resolved) else "cpu" # ALBU forces Albumentations even when aug_config is None if resolved == AugmentationBackend.ALBU and ns.aug_config is None: ns.aug_config = AUG_CONFIG if self.model_config.use_grouppose_keypoints and is_gpu_postprocess(resolved): raise ValueError( f"augmentation_backend='{resolved}' does not support keypoint transforms. " "Set augmentation_backend='cpu' or 'albumentations' when use_grouppose_keypoints=True." ) if self._dataset_train is None: self._dataset_train = build_dataset("train", ns, resolution) if self._dataset_val is None: self._dataset_val = build_dataset("val", ns, resolution) # Build Kornia pipeline (once); use _kornia_setup_done so fallback paths # (pipeline stays None) do not re-run on repeated setup("fit") calls. if not self._kornia_setup_done: self._resolved_augmentation_backend = resolved self._setup_kornia_pipeline() self._kornia_setup_done = True elif stage == "validate": if self._dataset_val is None: self._dataset_val = build_dataset("val", ns, resolution) elif stage == "test": if self._dataset_test is None: self._dataset_test = self._build_test_dataset(ns, resolution) elif stage == "predict": if self._dataset_val is None: self._dataset_val = build_dataset("val", ns, resolution) def _build_test_dataset(self, ns: Any, resolution: int) -> torch.utils.data.Dataset[Any]: """Build the dataset backing ``setup("test")``, preferring a real test split. ``roboflow`` and ``yolo`` datasets carry a labelled ``test`` split, so it is used directly. A YOLO-format dataset is not required to declare one: a ``roboflow`` dataset whose detected format is YOLO-style routes through the same builder as ``dataset_file="yolo"`` (:func:`rfdetr.datasets.build_roboflow` -> ``build_roboflow_from_yolo``), and Roboflow's export UI does not require a test split any more than a hand-built YOLO dataset does. When that YOLO-style split cannot be resolved, the ``val`` split is used instead and the substitution is logged, because falling back silently reports validation numbers under the name "test". A ``roboflow`` export in COCO format still has to ship ``test/_annotations.coco.json``: its absence raises a plain ``FileNotFoundError`` from the COCO builder, which this fallback does not catch. ``coco`` falls back to ``val`` because its ``test`` split is unlabelled COCO test-dev and cannot be scored locally. ``o365`` also falls back because its dataset builder exposes only ``train`` and ``val`` splits. Args: ns: Merged model/train config namespace forwarded to :func:`build_dataset`. resolution: Target square resolution in pixels. Returns: The dataset to evaluate during the ``test`` stage. """ dataset_file = self.train_config.dataset_file if dataset_file in ("roboflow", "yolo"): try: return build_dataset("test", ns, resolution) except YoloSplitUnavailableError as exc: dataset_label = {"roboflow": "Roboflow", "yolo": "YOLO"}[dataset_file] logger.warning( "No resolvable 'test' split for this %s dataset (%s); evaluating the 'val' split instead.", dataset_label, str(exc), ) return build_dataset("val", ns, resolution) def _resolve_batch_size(self) -> int: """Return the concrete training batch size. Raises: RuntimeError: If ``train_config.batch_size == "auto"`` was never resolved. Auto-batch sizing probes the live model and is only performed by ``RFDETR.train()`` (:func:`rfdetr.training.auto_batch.resolve_auto_batch_config`); the PTL CLI path does not resolve it automatically. """ batch_size = self.train_config.batch_size if batch_size == "auto": raise RuntimeError( "train_config.batch_size='auto' was not resolved before dataloader construction. " "Auto-batch sizing requires probing the live model and is only performed by " "RFDETR.train(); set batch_size to a concrete integer when using the PTL CLI path directly." ) return batch_size @staticmethod def _require_dataset(dataset: torch.utils.data.Dataset[Any] | None, split: str) -> torch.utils.data.Dataset[Any]: """Return *dataset*, raising if ``setup()`` has not built it yet.""" if dataset is None: raise RuntimeError(f"{split} dataset was not built; call setup({split!r}) before requesting a dataloader.") return dataset def train_dataloader(self) -> DataLoader[Any]: """Return the training DataLoader. Uses a replacement sampler when the dataset is too small to fill ``_MIN_TRAIN_BATCHES`` effective batches (matching legacy behaviour in ``main.py``). Otherwise wraps the dataset with :class:`GradAccumAlignedDataset` to ensure its length is an exact multiple of ``effective_batch_size * world_size`` (workaround for https://github.com/Lightning-AI/pytorch-lightning/issues/19987) and then uses ``shuffle=True, drop_last=True`` so that PTL can auto-inject ``DistributedSampler`` in DDP mode. Returns: DataLoader for the training dataset. """ dataset: torch.utils.data.Dataset[Any] = self._require_dataset(self._dataset_train, "fit") batch_size = self._resolve_batch_size() effective_batch_size = batch_size * self.train_config.grad_accum_steps num_workers = self._num_workers dataset_length = len(dataset) # type: ignore[arg-type] if dataset_length < effective_batch_size * _MIN_TRAIN_BATCHES: logger.info( "Training with uniform sampler because dataset is too small: %d < %d", dataset_length, effective_batch_size * _MIN_TRAIN_BATCHES, ) sampler = torch.utils.data.RandomSampler( dataset, # type: ignore[arg-type] replacement=True, num_samples=effective_batch_size * _MIN_TRAIN_BATCHES, ) return DataLoader( dataset, batch_size=batch_size, sampler=sampler, collate_fn=self._collate_fn, num_workers=num_workers, pin_memory=self._pin_memory, persistent_workers=self._persistent_workers, prefetch_factor=self._prefetch_factor, worker_init_fn=_worker_init_fn, ) # Pad the dataset to a multiple of effective_batch_size * world_size so # that drop_last=True below becomes a true no-op and PTL never fires the # optimizer on a partial accumulation window. # See https://github.com/Lightning-AI/pytorch-lightning/issues/19987 world_size: int = getattr(self.trainer, "world_size", 1) if self.trainer else 1 aligned_dataset = GradAccumAlignedDataset(dataset, effective_batch_size, world_size) return DataLoader( aligned_dataset, batch_size=batch_size, shuffle=True, drop_last=True, # no-op after alignment, but keeps intent explicit collate_fn=self._collate_fn, num_workers=num_workers, pin_memory=self._pin_memory, persistent_workers=self._persistent_workers, prefetch_factor=self._prefetch_factor, worker_init_fn=_worker_init_fn, ) def val_dataloader(self) -> DataLoader[Any]: """Return the validation DataLoader. Returns: DataLoader for the validation dataset with sequential sampling. """ dataset = self._require_dataset(self._dataset_val, "validate") return DataLoader( dataset, batch_size=self._resolve_batch_size(), sampler=torch.utils.data.SequentialSampler(dataset), # type: ignore[arg-type] drop_last=False, collate_fn=self._collate_fn, num_workers=self._num_workers, pin_memory=self._pin_memory, persistent_workers=self._persistent_workers, prefetch_factor=self._prefetch_factor, worker_init_fn=_worker_init_fn, ) def test_dataloader(self) -> DataLoader[Any]: """Return the test DataLoader. Returns: DataLoader for the test dataset with sequential sampling. """ dataset = self._require_dataset(self._dataset_test, "test") return DataLoader( dataset, batch_size=self._resolve_batch_size(), sampler=torch.utils.data.SequentialSampler(dataset), # type: ignore[arg-type] drop_last=False, collate_fn=self._collate_fn, num_workers=self._num_workers, pin_memory=self._pin_memory, persistent_workers=self._persistent_workers, prefetch_factor=self._prefetch_factor, worker_init_fn=_worker_init_fn, ) def predict_dataloader(self) -> DataLoader[Any]: """Return the predict DataLoader (reuses the validation dataset, no augmentation). Returns: DataLoader for the validation dataset with sequential sampling. """ dataset = self._require_dataset(self._dataset_val, "predict") return DataLoader( dataset, batch_size=self._resolve_batch_size(), sampler=torch.utils.data.SequentialSampler(dataset), # type: ignore[arg-type] drop_last=False, collate_fn=self._collate_fn, num_workers=self._num_workers, pin_memory=self._pin_memory, persistent_workers=self._persistent_workers, prefetch_factor=self._prefetch_factor, worker_init_fn=_worker_init_fn, ) def _show_samples( self, count: int, split: Literal["train", "val", "test"] = "train", *, columns: int = 3, figure_size: tuple[float, float] | None = None, ) -> "Figure": """Build a private diagnostic figure for transformed dataset samples. Samples the dataset after RF-DETR dataset transforms, so boxes and keypoints match the model input tensors rather than raw annotation JSON. Args: count: Maximum number of samples to render. split: Dataset split to visualize. columns: Number of subplot columns. figure_size: Optional Matplotlib figure size ``(width, height)`` in inches. When omitted, the size is derived from the grid shape. Returns: Matplotlib figure containing the annotated sample grid. When the dataset includes instance masks, they are rendered as coloured overlays before bounding boxes and labels. Raises: ValueError: If ``count`` or ``columns`` is not positive. Example: >>> # dm = RFDETRDataModule(model_config, train_config) >>> # figure = dm._show_samples(3, split="train") """ try: import matplotlib.pyplot as plt import numpy as np import supervision as sv import torchvision.transforms as T # noqa: N812 except ImportError as err: raise ImportError( "RFDETRDataModule._show_samples() requires visualization dependencies. " "Install them with `pip install 'rfdetr[visual]'`." ) from err from rfdetr.utilities.box_ops import box_cxcywh_to_xyxy if count <= 0: raise ValueError(f"count must be positive, got {count}.") if columns <= 0: raise ValueError(f"columns must be positive, got {columns}.") if figure_size is not None: if len(figure_size) != 2: raise ValueError(f"figure_size must contain two values, got {figure_size}.") if figure_size[0] <= 0 or figure_size[1] <= 0: raise ValueError(f"figure_size values must be positive, got {figure_size}.") dataset = self._get_dataset_for_visualization(split) if dataset is None: raise RuntimeError(f"Could not build dataset split {split!r} for visualization.") inv_normalize = T.Normalize( mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225], std=[1 / 0.229, 1 / 0.224, 1 / 0.225], ) dataset_length = len(dataset) # type: ignore[arg-type] rows = max(1, (min(count, dataset_length) + columns - 1) // columns) figure, axes = plt.subplots(rows, columns, figsize=figure_size or (5 * columns, 5 * rows)) axes_array = np.asarray(axes, dtype=object).reshape(-1) for axis in axes_array: axis.axis("off") class_names = self.class_names for axis, sample_index in zip(axes_array, range(min(count, dataset_length)), strict=False): image_tensor, target = dataset[sample_index] image_path = self._source_image_path(dataset, sample_index) image = inv_normalize(image_tensor) image_array = image.detach().cpu().numpy() scene = np.ascontiguousarray((np.clip(image_array.transpose(1, 2, 0), 0.0, 1.0) * 255).astype(np.uint8)) size = target.get("size") if isinstance(size, torch.Tensor): height, width = int(size[0]), int(size[1]) else: height, width = int(image_tensor.shape[-2]), int(image_tensor.shape[-1]) boxes = target.get("boxes", torch.zeros((0, 4), dtype=torch.float32)) labels = target.get("labels", torch.zeros((0,), dtype=torch.int64)) if boxes.numel() > 0: scale = torch.tensor([width, height, width, height], dtype=torch.float32) xyxy = box_cxcywh_to_xyxy(boxes.detach().cpu()) * scale class_ids = labels.detach().cpu().numpy().astype(int) masks_tensor = target.get("masks") mask_array = ( masks_tensor.detach().cpu().numpy() if masks_tensor is not None and masks_tensor.numel() > 0 else None ) detections = sv.Detections(xyxy=xyxy.numpy().astype(np.float32), class_id=class_ids, mask=mask_array) labels_text = [ class_names[class_id] if class_names is not None and class_id < len(class_names) else str(class_id) for class_id in class_ids ] if mask_array is not None: scene = sv.MaskAnnotator().annotate(scene=scene, detections=detections) scene = sv.BoxAnnotator(thickness=1).annotate(scene=scene, detections=detections) scene = sv.LabelAnnotator(text_scale=0.4, text_padding=2).annotate( scene=scene, detections=detections, labels=labels_text, ) keypoints = target.get("keypoints") if keypoints is not None and keypoints.numel() > 0: keypoints_array = keypoints.detach().cpu().numpy().astype(np.float32) keypoint_xy = keypoints_array[..., :2] * np.asarray([width, height], dtype=np.float32) keypoint_visibility = keypoints_array[..., 2] > 0 key_points = sv.KeyPoints( xy=keypoint_xy, keypoint_confidence=keypoint_visibility.astype(np.float32), class_id=labels.detach().cpu().numpy().astype(int), visible=keypoint_visibility, data={"visible": keypoint_visibility}, ) scene = sv.VertexAnnotator(radius=3).annotate(scene=scene, key_points=key_points) axis.imshow(scene) title = image_path.name if image_path is not None else f"{split}[{sample_index}]" axis.set_title(self._ellipsize_sample_title(title), fontsize=10) axis.axis("off") figure.tight_layout() return figure @staticmethod def _ellipsize_sample_title(title: str, max_length: int = 48) -> str: """Shorten long sample titles so subplot grids do not overflow. Args: title: Raw title text, usually an image file name. max_length: Maximum returned character count including the ellipsis. Returns: Original title when it already fits, otherwise a middle-ellipsized string that preserves the start and file suffix. """ if len(title) <= max_length: return title if max_length <= 3: return "." * max_length keep_left = max(1, (max_length - 3) // 2) keep_right = max_length - 3 - keep_left return f"{title[:keep_left]}...{title[-keep_right:]}" def _get_dataset_for_visualization( self, split: Literal["train", "val", "test"], ) -> torch.utils.data.Dataset[Any] | None: """Return a built dataset split for private visualization.""" if split == "train": self.setup("fit") return self._dataset_train if split == "val": self.setup("validate") return self._dataset_val if split == "test": self.setup("test") return self._dataset_test raise ValueError(f"Unsupported split {split!r}.") @staticmethod def _source_image_path(dataset: torch.utils.data.Dataset[Any], sample_index: int) -> Path | None: """Return a source image path for common COCO-style datasets.""" image_folder = getattr(dataset, "root", None) image_ids = getattr(dataset, "ids", None) coco = getattr(dataset, "coco", None) if image_folder is None or image_ids is None or coco is None: return None image_id = image_ids[sample_index] image_info = coco.loadImgs(image_id)[0] return Path(image_folder) / image_info["file_name"] # type: ignore[no-any-return] def _setup_kornia_pipeline(self) -> None: """Build the Kornia pipeline for ``self._resolved_augmentation_backend``. ``TV`` and ``ALBU`` are no-ops. ``KORNIA`` validates that kornia is installed then builds the pipeline on whatever device the batch arrives on. ``self._resolved_augmentation_backend`` (set by :meth:`setup` just before this call) is the concrete backend returned by :func:`rfdetr.datasets.kornia_transforms.resolve_augmentation_backend`. """ from rfdetr.datasets.kornia_transforms import is_gpu_postprocess resolved = self._resolved_augmentation_backend if resolved is None or not is_gpu_postprocess(resolved): return if not AugmentationBackend._is_kornia_available(): raise ImportError("Kornia augmentation requires kornia. Install with: pip install 'rfdetr[augment]'") from rfdetr.datasets.kornia_transforms import build_kornia_pipeline, build_normalize self._kornia_pipeline = build_kornia_pipeline( self.train_config.aug_config if self.train_config.aug_config is not None else AUG_CONFIG, self.model_config.resolution, with_masks=self.model_config.segmentation_head, ) self._kornia_normalize = build_normalize() logger.info("Kornia augmentation pipeline built (resolved=%s)", resolved) def on_after_batch_transfer(self, batch: tuple[Any, Any], dataloader_idx: int) -> tuple[Any, Any]: """Apply Kornia GPU augmentation after the batch is transferred to device. When ``_kornia_pipeline`` is set and the trainer is in training mode, augmentation and normalization are applied on the GPU. Validation and test batches pass through unchanged. Segmentation models use a mask-aware pipeline (``with_masks=True``) so images, boxes, and per-instance masks are augmented in sync. Args: batch: Tuple of ``(NestedTensor, list[dict])`` already on device. dataloader_idx: Index of the current dataloader. Returns: The (possibly augmented) batch. """ kornia_pipeline = self._kornia_pipeline kornia_normalize = self._kornia_normalize if self.trainer is None or not self.trainer.training or kornia_pipeline is None or kornia_normalize is None: return batch from rfdetr.datasets.kornia_transforms import collate_boxes, collate_masks, unpack_boxes from rfdetr.utilities.tensors import NestedTensor samples, targets = batch img = samples.tensors # [B, C, H, W] # Move Kornia modules to the batch device (no-op if already there). # nn.Module.to() is in-place; no reassignment needed. kornia_pipeline.to(img.device) kornia_normalize.to(img.device) boxes_padded, valid = collate_boxes(targets, img.device) if self.model_config.segmentation_head: image_height, image_width = img.shape[-2:] masks_padded = collate_masks( targets, img.device, n_max=valid.shape[1], image_height=image_height, image_width=image_width ) img_aug, boxes_aug, masks_aug = kornia_pipeline(img, boxes_padded, masks_padded) img_aug = kornia_normalize(img_aug) aug_height, aug_width = img_aug.shape[-2:] targets = unpack_boxes(boxes_aug, valid, targets, aug_height, aug_width, masks_aug=masks_aug) else: img_aug, boxes_aug = kornia_pipeline(img, boxes_padded) img_aug = kornia_normalize(img_aug) aug_height, aug_width = img_aug.shape[-2:] targets = unpack_boxes(boxes_aug, valid, targets, aug_height, aug_width) height, width = img_aug.shape[-2:] for target in targets: boxes = target["boxes"] if boxes.numel() == 0: continue scale = boxes.new_tensor([width, height, width, height]) target["boxes"] = box_xyxy_to_cxcywh(boxes) / scale batch = (NestedTensor(img_aug, samples.mask), targets) return batch # ------------------------------------------------------------------ # Properties # ------------------------------------------------------------------ @property def class_names(self) -> list[str] | None: """Class names from the training or validation dataset annotation file. Reads category names from the first available COCO-style dataset. Returns ``None`` if no dataset has been set up yet or the dataset does not expose COCO-style category information. Returns: Sorted list of class name strings, or ``None``. """ for dataset in (self._dataset_train, self._dataset_val, self._dataset_test): if dataset is None: continue coco = getattr(dataset, "coco", None) if coco is not None and hasattr(coco, "cats"): label2cat = getattr(dataset, "label2cat", None) if label2cat is None: label2cat = getattr(coco, "label2cat", None) if isinstance(label2cat, dict) and label2cat: max_label = max(label2cat) names = [""] * (max_label + 1) for label, category_id in sorted(label2cat.items()): category = coco.cats.get(category_id) if category is not None: names[label] = category["name"] return names return [coco.cats[k]["name"] for k in sorted(coco.cats.keys())] return None def transfer_batch_to_device( self, batch: tuple[Any, Any], device: torch.device, dataloader_idx: int ) -> tuple[Any, Any]: """Move a ``(NestedTensor, targets)`` batch to *device*. PTL's default iterates tuple elements and calls ``.to(device)``; that works for plain tensors but ``NestedTensor`` must be moved explicitly. Args: batch: Tuple of (NestedTensor samples, list of target dicts). device: Target device. dataloader_idx: Index of the dataloader providing this batch. Returns: Batch with all tensors on ``device``. """ samples, targets = batch non_blocking = device.type == "cuda" samples = samples.to(device, non_blocking=non_blocking) targets = [{k: v.to(device, non_blocking=non_blocking) for k, v in t.items()} for t in targets] return samples, targets