Skip to content

Data API

Datasets

modern_yolonas.data.coco.COCODetectionDataset

Bases: Dataset

COCO-format detection dataset.

Parameters:

Name Type Description Default
root str | Path

Path to image directory (e.g., coco/images/train2017).

 required
ann_file str | Path

Path to annotation JSON (e.g., coco/annotations/instances_train2017.json).

 required
transforms Transform | None

(image, targets) → (image, targets) callable.

 None
input_size int

Target input size (used by transforms).

 640
cat_id_to_label dict[int, int] | None

Category id → label index to use instead of deriving one from this file. Required when this dataset trains alongside another: the derived mapping only counts categories that have annotations, so a file missing one rare class would shift every later label by one.

 None
Source code in src/modern_yolonas/data/coco.py 
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
class COCODetectionDataset(Dataset):
    """COCO-format detection dataset.

    Args:
        root: Path to image directory (e.g., ``coco/images/train2017``).
        ann_file: Path to annotation JSON (e.g., ``coco/annotations/instances_train2017.json``).
        transforms: ``(image, targets) → (image, targets)`` callable.
        input_size: Target input size (used by transforms).
        cat_id_to_label: Category id → label index to use instead of deriving one
            from this file. Required when this dataset trains alongside another:
            the derived mapping only counts categories that have annotations, so a
            file missing one rare class would shift every later label by one.
    """

    def __init__(
        self,
        root: str | Path,
        ann_file: str | Path,
        transforms: Transform | None = None,
        input_size: int = 640,
        cache_annotations: bool = True,
        ignore_empty_annotations: bool = True,
        cat_id_to_label: dict[int, int] | None = None,
    ):
        from pycocotools.coco import COCO

        self.root = Path(root)
        self.transforms = transforms
        self.input_size = input_size

        self.coco = COCO(str(ann_file))
        self.ids = list(sorted(self.coco.imgs.keys()))

        # Build contiguous class mapping using only categories present in annotations.
        # Excludes unused categories (e.g. a "background" id=0 that carries no annotations)
        # so that the resulting label indices are always 0-indexed and contiguous.
        ann_cat_ids = {ann["category_id"] for ann in self.coco.dataset.get("annotations", [])}
        if cat_id_to_label is None:
            cat_ids = sorted(c for c in self.coco.getCatIds() if c in ann_cat_ids)
            self.cat_id_to_label = {cat_id: i for i, cat_id in enumerate(cat_ids)}
        else:
            unknown = {
                ann["category_id"] for ann in self.coco.dataset.get("annotations", [])
                if not ann.get("iscrowd", 0) and ann["category_id"] not in cat_id_to_label
            }
            if unknown:
                raise ValueError(
                    f"{ann_file} uses category ids missing from the given mapping: "
                    f"{sorted(unknown)}"
                )
            self.cat_id_to_label = dict(cat_id_to_label)
            cat_ids = sorted(self.cat_id_to_label, key=self.cat_id_to_label.__getitem__)

        # Ordered list of class names aligned with label indices 0..N-1
        cats = self.coco.loadCats(cat_ids)
        self.class_names: list[str] = [c["name"] for c in cats]

        # Pre-load all annotation arrays into memory to avoid repeated pycocotools lookups
        self._label_cache: list[np.ndarray] | None = None
        if cache_annotations:
            self._label_cache = [self._parse_anns(img_id) for img_id in self.ids]

        # Drop images with no annotations (background-only samples waste training batches)
        if ignore_empty_annotations:
            labels = self._label_cache if self._label_cache is not None else [self._parse_anns(img_id) for img_id in self.ids]
            keep = [i for i, t in enumerate(labels) if len(t) > 0]
            self.ids = [self.ids[i] for i in keep]
            if self._label_cache is not None:
                self._label_cache = [self._label_cache[i] for i in keep]

    def __len__(self) -> int:
        return len(self.ids)

    def _parse_anns(self, img_id: int) -> np.ndarray:
        """Return ``[N, 5]`` float32 targets for *img_id* (no image needed)."""
        img_info = self.coco.loadImgs(img_id)[0]
        h_img, w_img = img_info["height"], img_info["width"]
        ann_ids = self.coco.getAnnIds(imgIds=img_id)
        anns = self.coco.loadAnns(ann_ids)
        targets = []
        for ann in anns:
            x, y, bw, bh = ann["bbox"]
            if bw <= 0 or bh <= 0:
                # COCO train2017 carries two annotations with a zero-height box.
                # They are not trainable targets, and Albumentations rejects them
                # outright -- `y_max is less than or equal to y_min` -- which took
                # a training run down about a thousand steps in, once the images
                # happened to be drawn. pycocotools ignores zero-area annotations
                # when scoring, so nothing is lost by never emitting them. Applies
                # to crowd regions too: a zero-area one covers no anchors, so it
                # would only be a box waiting to crash a transform.
                continue
            # Crowd regions are kept, marked, and used by the loss to ignore the
            # anchors they cover. Dropping them taught the model that a crowd is
            # background; training on them would teach one box around many objects.
            cls = CROWD_CLASS if ann.get("iscrowd", 0) else self.cat_id_to_label[ann["category_id"]]
            xc = (x + bw / 2) / w_img
            yc = (y + bh / 2) / h_img
            nw = bw / w_img
            nh = bh / h_img
            targets.append([cls, xc, yc, nw, nh])
        return np.array(targets, dtype=np.float32).reshape(-1, 5) if targets else np.zeros((0, 5), dtype=np.float32)

    def load_raw(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        """Load image and targets without transforms."""
        img_id = self.ids[index]
        img_info = self.coco.loadImgs(img_id)[0]
        image = cv2.imread(str(self.root / img_info["file_name"]))

        if self._label_cache is not None:
            targets = self._label_cache[index]
        else:
            targets = self._parse_anns(img_id)

        return image, targets

    def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        if self.transforms is not None and hasattr(self.transforms, 'apply'):
            return self.transforms.apply(index, self.load_raw)
        image, targets = self.load_raw(index)
        if self.transforms is not None:
            image, targets = self.transforms(image, targets)
        return image, targets

load_raw(index)

Load image and targets without transforms.

Source code in src/modern_yolonas/data/coco.py 
121
122
123
124
125
126
127
128
129
130
131
132
def load_raw(self, index: int) -> tuple[np.ndarray, np.ndarray]:
    """Load image and targets without transforms."""
    img_id = self.ids[index]
    img_info = self.coco.loadImgs(img_id)[0]
    image = cv2.imread(str(self.root / img_info["file_name"]))

    if self._label_cache is not None:
        targets = self._label_cache[index]
    else:
        targets = self._parse_anns(img_id)

    return image, targets

modern_yolonas.data.yolo.YOLODetectionDataset

Bases: Dataset

Source code in src/modern_yolonas/data/yolo.py 
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class YOLODetectionDataset(Dataset):
    def __init__(
        self,
        root: str | Path,
        split: str = "train",
        transforms=None,
        input_size: int = 640,
        cache_annotations: bool = True,
        ignore_empty_annotations: bool = True,
    ):
        self.root = Path(root)
        self.transforms = transforms
        self.input_size = input_size

        img_dir = self.root / "images" / split
        label_dir = self.root / "labels" / split

        self.images = sorted(img_dir.glob("*.*"))
        self.images = [p for p in self.images if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".bmp")]
        self.label_dir = label_dir

        # Try to load class names from classes.txt or data.yaml at dataset root
        self.class_names: list[str] | None = self._load_class_names()

        # Pre-load all annotation arrays into memory (avoids repeated disk I/O during training)
        self._label_cache: list[np.ndarray] | None = None
        if cache_annotations:
            self._label_cache = [self._read_label(p) for p in self.images]

        # Drop images with no annotations (background-only samples waste training batches)
        if ignore_empty_annotations:
            labels = self._label_cache if self._label_cache is not None else [self._read_label(p) for p in self.images]
            keep = [i for i, t in enumerate(labels) if len(t) > 0]
            self.images = [self.images[i] for i in keep]
            if self._label_cache is not None:
                self._label_cache = [self._label_cache[i] for i in keep]

    def _load_class_names(self) -> list[str] | None:
        """Return class names from ``classes.txt`` or ``data.yaml`` if present."""
        classes_txt = self.root / "classes.txt"
        if classes_txt.exists():
            names = [label.strip() for label in classes_txt.read_text().splitlines() if label.strip()]
            return names if names else None

        for yaml_name in ("data.yaml", "dataset.yaml"):
            yaml_path = self.root / yaml_name
            if yaml_path.exists():
                import yaml
                data = yaml.safe_load(yaml_path.read_text())
                names = data.get("names")
                if isinstance(names, list) and names:
                    return [str(n) for n in names]
                if isinstance(names, dict):
                    return [str(names[k]) for k in sorted(names)]

        return None

    def __len__(self) -> int:
        return len(self.images)

    @property
    def num_classes(self) -> int:
        """Number of classes inferred from cached labels (or scanning all label files)."""
        labels = self._label_cache if self._label_cache is not None else [self._read_label(p) for p in self.images]
        all_cls = np.concatenate([t[:, 0] for t in labels if len(t)]) if any(len(t) for t in labels) else np.array([])
        return int(all_cls.max()) + 1 if len(all_cls) else 0

    def _label_path(self, img_path: Path) -> Path:
        return self.label_dir / (img_path.stem + ".txt")

    def _read_label(self, img_path: Path) -> np.ndarray:
        label_path = self._label_path(img_path)
        if label_path.exists():
            data = np.loadtxt(str(label_path), ndmin=2).reshape(-1, 5)
        else:
            data = np.zeros((0, 5))
        return data.astype(np.float32)

    def load_raw(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        """Load image and targets without transforms."""
        img_path = self.images[index]
        image = cv2.imread(str(img_path))
        targets = self._label_cache[index] if self._label_cache is not None else self._read_label(img_path)
        return image, targets

    def __getitem__(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        if self.transforms is not None and hasattr(self.transforms, 'apply'):
            return self.transforms.apply(index, self.load_raw)
        image, targets = self.load_raw(index)
        if self.transforms is not None:
            image, targets = self.transforms(image, targets)
        return image, targets

num_classes property

Number of classes inferred from cached labels (or scanning all label files).

load_raw(index)

Load image and targets without transforms.

Source code in src/modern_yolonas/data/yolo.py 
107
108
109
110
111
112
def load_raw(self, index: int) -> tuple[np.ndarray, np.ndarray]:
    """Load image and targets without transforms."""
    img_path = self.images[index]
    image = cv2.imread(str(img_path))
    targets = self._label_cache[index] if self._label_cache is not None else self._read_label(img_path)
    return image, targets

Transforms

modern_yolonas.data.transforms

Detection-aware data augmentations.

Each transform operates on (image, targets) where: - image: HWC uint8 BGR numpy array - targets: [N, 5] numpy array with [class_id, x_center, y_center, w, h] (normalized)

HSVAugment, HorizontalFlip, RandomAffine, RandomResizedCrop, and RandomChannelSwap are backed by Albumentations <https://albumentations.ai>_ (MIT license, v2+). Mosaic, Mixup, LetterboxResize, and Normalize use native implementations because they have no direct Albumentations equivalent.

Compose

Chain multiple transforms sequentially.

Parameters:

Name Type Description Default
transforms list

List of callables, each accepting (image, targets) and returning (image, targets).

 required
Source code in src/modern_yolonas/data/transforms.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
class Compose:
    """Chain multiple transforms sequentially.

    Args:
        transforms: List of callables, each accepting ``(image, targets)``
            and returning ``(image, targets)``.
    """

    def __init__(self, transforms: list):
        self.transforms = transforms

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        for t in self.transforms:
            image, targets = t(image, targets)
        return image, targets

    def disable_mosaic_mixup(self):
        """Turn off any Mosaic or Mixup in the chain (called by CloseMosaicCallback).

        ``TrainTransformPipeline`` holds them in named slots; a plain Compose just has
        them somewhere in the list, so find them by type. Idempotent.
        """
        for t in self.transforms:
            if isinstance(t, (Mosaic, Mixup)):
                t.enabled = False

disable_mosaic_mixup()

Turn off any Mosaic or Mixup in the chain (called by CloseMosaicCallback).

TrainTransformPipeline holds them in named slots; a plain Compose just has them somewhere in the list, so find them by type. Idempotent.

Source code in src/modern_yolonas/data/transforms.py 
72
73
74
75
76
77
78
79
80
def disable_mosaic_mixup(self):
    """Turn off any Mosaic or Mixup in the chain (called by CloseMosaicCallback).

    ``TrainTransformPipeline`` holds them in named slots; a plain Compose just has
    them somewhere in the list, so find them by type. Idempotent.
    """
    for t in self.transforms:
        if isinstance(t, (Mosaic, Mixup)):
            t.enabled = False

HSVAugment

Randomly shift hue, saturation and value.

Implemented with three 256-entry lookup tables rather than through Albumentations. The shift is constant across the image, so it is exactly a per-channel table lookup, and cv2.LUT applies one in a single pass. The Albumentations route cost two whole-image copies to swap BGR/RGB around a call that converted colour space again internally — 5.32 ms against 2.32 ms per 640x640 image, measured 2026-09-19.

The shift ranges are unchanged: OpenCV's uint8 HSV encoding puts hue in [0, 179] and saturation and value in [0, 255], which is the same convention HueSaturationValue used, so recipes keep their meaning.

Parameters:

Name Type Description Default
hgain int

Max hue shift, in OpenCV hue units. Matches the super-gradients hgain recipe param. Default: 18.

 18
sgain int

Max saturation shift in absolute units. Default: 30.

 30
vgain int

Max value shift in absolute units. Default: 30.

 30
p float

Probability of applying the transform.

 0.5
Source code in src/modern_yolonas/data/transforms.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
class HSVAugment:
    """Randomly shift hue, saturation and value.

    Implemented with three 256-entry lookup tables rather than through
    Albumentations. The shift is constant across the image, so it is exactly a
    per-channel table lookup, and ``cv2.LUT`` applies one in a single pass. The
    Albumentations route cost two whole-image copies to swap BGR/RGB around a call
    that converted colour space again internally — 5.32 ms against 2.32 ms per
    640x640 image, measured 2026-09-19.

    The shift ranges are unchanged: OpenCV's uint8 HSV encoding puts hue in
    ``[0, 179]`` and saturation and value in ``[0, 255]``, which is the same
    convention ``HueSaturationValue`` used, so recipes keep their meaning.

    Args:
        hgain: Max hue shift, in OpenCV hue units. Matches the super-gradients
              ``hgain`` recipe param. Default: 18.
        sgain: Max saturation shift in absolute units. Default: 30.
        vgain: Max value shift in absolute units. Default: 30.
        p: Probability of applying the transform.
    """

    def __init__(self, hgain: int = 18, sgain: int = 30, vgain: int = 30, p: float = 0.5):
        self.hgain = hgain
        self.sgain = sgain
        self.vgain = vgain
        self.p = p
        self._ramp = np.arange(256, dtype=np.int16)

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        if random.random() >= self.p:
            return image, targets

        dh, ds, dv = (
            random.uniform(-1, 1) * self.hgain,
            random.uniform(-1, 1) * self.sgain,
            random.uniform(-1, 1) * self.vgain,
        )

        hue, sat, val = cv2.split(cv2.cvtColor(image, cv2.COLOR_BGR2HSV))
        x = self._ramp
        # Hue is an angle, so it wraps; saturation and value saturate at the ends.
        lut_h = ((x + dh) % 180).astype(np.uint8)
        lut_s = np.clip(x + ds, 0, 255).astype(np.uint8)
        lut_v = np.clip(x + dv, 0, 255).astype(np.uint8)

        merged = cv2.merge((cv2.LUT(hue, lut_h), cv2.LUT(sat, lut_s), cv2.LUT(val, lut_v)))
        return cv2.cvtColor(merged, cv2.COLOR_HSV2BGR), targets

HorizontalFlip

Randomly flip the image and bounding boxes horizontally via Albumentations.

Parameters:

Name Type Description Default
p float

Probability of applying the flip.

 0.5
Source code in src/modern_yolonas/data/transforms.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
class HorizontalFlip:
    """Randomly flip the image and bounding boxes horizontally via Albumentations.

    Args:
        p: Probability of applying the flip.
    """

    def __init__(self, p: float = 0.5):
        self.p = p
        self._aug = A.Compose([A.HorizontalFlip(p=1.0)], bbox_params=_BBOX_PARAMS)

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        if random.random() >= self.p:
            return image, targets
        bboxes, labels = _to_albu(targets)
        r = self._aug(image=image, bboxes=bboxes, class_labels=labels)
        return r["image"], _from_albu(r["bboxes"], r["class_labels"], targets.dtype)

RandomChannelShuffle

Randomly permute BGR channels (equivalent to super-gradients DetectionRGB2BGR).

Source code in src/modern_yolonas/data/transforms.py 
165
166
167
168
169
170
171
172
173
174
175
class RandomChannelShuffle:
    """Randomly permute BGR channels (equivalent to super-gradients DetectionRGB2BGR)."""

    def __init__(self, p: float = 0.5):
        self.p = p

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        if random.random() < self.p:
            perm = np.random.permutation(3)
            image = image[:, :, perm].copy()
        return image, targets

RandomCrop

Random crop with bbox clipping and re-normalization.

Source code in src/modern_yolonas/data/transforms.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
class RandomCrop:
    """Random crop with bbox clipping and re-normalization."""

    def __init__(self, min_scale: float = 0.3, max_scale: float = 1.0, p: float = 1.0):
        self.min_scale = min_scale
        self.max_scale = max_scale
        self.p = p

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        if random.random() >= self.p:
            return image, targets

        h, w = image.shape[:2]

        # Sample crop dimensions
        crop_h = int(h * random.uniform(self.min_scale, self.max_scale))
        crop_w = int(w * random.uniform(self.min_scale, self.max_scale))
        crop_h = max(1, crop_h)
        crop_w = max(1, crop_w)

        # Sample crop position
        y0 = random.randint(0, h - crop_h)
        x0 = random.randint(0, w - crop_w)

        image = image[y0 : y0 + crop_h, x0 : x0 + crop_w].copy()

        if len(targets):
            targets = targets.copy()
            # Convert normalized xywh to pixel xyxy
            x1 = (targets[:, 1] - targets[:, 3] / 2) * w
            y1 = (targets[:, 2] - targets[:, 4] / 2) * h
            x2 = (targets[:, 1] + targets[:, 3] / 2) * w
            y2 = (targets[:, 2] + targets[:, 4] / 2) * h

            # Shift to crop coordinates and clip
            x1 = np.clip(x1 - x0, 0, crop_w)
            y1 = np.clip(y1 - y0, 0, crop_h)
            x2 = np.clip(x2 - x0, 0, crop_w)
            y2 = np.clip(y2 - y0, 0, crop_h)

            # Filter out boxes that are too small
            box_w = x2 - x1
            box_h = y2 - y1
            valid = (box_w > 2) & (box_h > 2)

            targets = targets[valid]
            x1, y1, x2, y2 = x1[valid], y1[valid], x2[valid], y2[valid]

            if len(targets):
                targets[:, 1] = ((x1 + x2) / 2) / crop_w
                targets[:, 2] = ((y1 + y2) / 2) / crop_h
                targets[:, 3] = (x2 - x1) / crop_w
                targets[:, 4] = (y2 - y1) / crop_h

        return image, targets

RandomAffine

Apply random rotation, scale, translation, and shear via Albumentations.

Parameters:

Name Type Description Default
degrees float

Maximum rotation in degrees.

 0.0
translate float

Maximum translation as a fraction of image size.

 0.25
scale tuple[float, float]

Scale range (min, max).

 (0.5, 1.5)
shear float

Maximum shear in degrees.

 0.0
Source code in src/modern_yolonas/data/transforms.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
class RandomAffine:
    """Apply random rotation, scale, translation, and shear via Albumentations.

    Args:
        degrees: Maximum rotation in degrees.
        translate: Maximum translation as a fraction of image size.
        scale: Scale range ``(min, max)``.
        shear: Maximum shear in degrees.
    """

    def __init__(
        self,
        degrees: float = 0.0,
        translate: float = 0.25,
        scale: tuple[float, float] = (0.5, 1.5),
        shear: float = 0.0,
    ):
        self._aug = A.Compose(
            [
                A.Affine(
                    scale=scale,
                    translate_percent={"x": (-translate, translate), "y": (-translate, translate)},
                    rotate=(-degrees, degrees),
                    shear=(-shear, shear),
                    border_mode=cv2.BORDER_CONSTANT,
                    fill=114,
                    p=1.0,
                )
            ],
            bbox_params=_BBOX_PARAMS,
        )

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        bboxes, labels = _to_albu(targets)
        r = self._aug(image=image, bboxes=bboxes, class_labels=labels)
        return r["image"], _from_albu(r["bboxes"], r["class_labels"], targets.dtype)

RandomResizedCrop

Randomly crop a region of the image and resize it to size via Albumentations.

Mirrors torchvision.transforms.RandomResizedCrop but is bounding-box aware. Boxes whose area falls below min_width / min_height pixels after cropping are automatically discarded by _BBOX_PARAMS.

Parameters:

Name Type Description Default
size int

Output square side length in pixels.

 640
scale tuple[float, float]

Range of fraction of the original image area to crop. Default (0.08, 1.0) matches the torchvision default.

 (0.08, 1.0)
ratio tuple[float, float]

Range of aspect ratio of the crop. Default (0.75, 1.333) matches the torchvision default.

 (0.75, 1.333)
interpolation int

OpenCV interpolation flag (default cv2.INTER_LINEAR).

 INTER_LINEAR
p float

Probability of applying the transform.

 1.0
Source code in src/modern_yolonas/data/transforms.py 
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
class RandomResizedCrop:
    """Randomly crop a region of the image and resize it to ``size`` via Albumentations.

    Mirrors ``torchvision.transforms.RandomResizedCrop`` but is bounding-box
    aware.  Boxes whose area falls below ``min_width`` / ``min_height`` pixels
    after cropping are automatically discarded by ``_BBOX_PARAMS``.

    Args:
        size: Output square side length in pixels.
        scale: Range of fraction of the original image area to crop.
              Default ``(0.08, 1.0)`` matches the torchvision default.
        ratio: Range of aspect ratio of the crop.
              Default ``(0.75, 1.333)`` matches the torchvision default.
        interpolation: OpenCV interpolation flag (default ``cv2.INTER_LINEAR``).
        p: Probability of applying the transform.
    """

    def __init__(
        self,
        size: int = 640,
        scale: tuple[float, float] = (0.08, 1.0),
        ratio: tuple[float, float] = (0.75, 1.333),
        interpolation: int = cv2.INTER_LINEAR,
        p: float = 1.0,
    ):
        self._aug = A.Compose(
            [
                A.RandomResizedCrop(
                    size=(size, size),
                    scale=scale,
                    ratio=ratio,
                    interpolation=interpolation,
                    p=1.0,
                )
            ],
            bbox_params=_BBOX_PARAMS,
        )
        self.p = p

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        if random.random() >= self.p:
            return image, targets
        bboxes, labels = _to_albu(targets)
        r = self._aug(image=image, bboxes=bboxes, class_labels=labels)
        return r["image"], _from_albu(r["bboxes"], r["class_labels"], targets.dtype)

RandomResizedCropFlipAffine

RandomResizedCrop → HorizontalFlip → RandomAffine, as one warp.

The three transforms it replaces are each affine, so their matrices multiply and a single cv2.warpAffine produces the same output. That is worth doing twice over: it resamples the image once instead of three times, and it drops two full passes over a 640x640 array plus two Albumentations bbox round-trips.

The parameters are drawn by calling Albumentations' own samplers, so the augmentation distribution is identical by construction rather than by reimplementation — A.Affine samples x and y scale independently, for one, which a hand-rolled version would be unlikely to reproduce.

Two matrices, deliberately. OpenCV's resize and warpAffine map pixel centres: the source coordinate for output pixel d is (d + 0.5) * s - 0.5. Albumentations' bbox crop is geometric — (src - x1) * S / crop with no half-pixel term — which is also why A.Affine returns a separate bbox_matrix. Composing the image chain with the geometric convention shifts the picture by 0.5 * (1 - S / crop) pixels, half a pixel when the crop is upscaled 2x, while the boxes stay put. Nothing raises; the labels simply stop matching the pixels. Measured: with the geometric matrix the 99th percentile of the difference against the chain is up to 14 grey levels, with the pixel-centre matrix it is 1.

So _image_matrix carries the half-pixel terms and _box_matrix does not. They are not meant to agree.

Only axis-aligned chains are supported. With rotation or shear, clipping a box once at the end is not the same as clipping it after each step, and this class would silently produce different boxes; it raises instead.

Parameters:

Name Type Description Default
size int

Output square side.

 640
scale tuple[float, float]

Crop area fraction range, as A.RandomResizedCrop.

 (0.05, 0.8)
ratio tuple[float, float]

Crop aspect ratio range, as A.RandomResizedCrop.

 (0.75, 1.33)
flip_prob float

Probability of the horizontal flip.

 0.5
translate float

Affine translation as a fraction of output size.

 0.25
affine_scale tuple[float, float]

Affine scale range.

 (0.5, 1.5)
degrees float

Rotation. Must be 0.

 0.0
shear float

Shear. Must be 0.

 0.0
pad_value int

Fill for pixels pulled in from outside the source.

 114
min_box_size float

Drop boxes thinner than this, in output pixels. Matches the min_width / min_height of _BBOX_PARAMS.

 2.0
Source code in src/modern_yolonas/data/transforms.py 
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
class RandomResizedCropFlipAffine:
    """``RandomResizedCrop`` → ``HorizontalFlip`` → ``RandomAffine``, as one warp.

    The three transforms it replaces are each affine, so their matrices multiply and
    a single ``cv2.warpAffine`` produces the same output. That is worth doing twice
    over: it resamples the image once instead of three times, and it drops two full
    passes over a 640x640 array plus two Albumentations bbox round-trips.

    The parameters are drawn by calling Albumentations' own samplers, so the
    augmentation distribution is identical by construction rather than by
    reimplementation — ``A.Affine`` samples x and y scale *independently*, for one,
    which a hand-rolled version would be unlikely to reproduce.

    **Two matrices, deliberately.** OpenCV's ``resize`` and ``warpAffine`` map pixel
    *centres*: the source coordinate for output pixel ``d`` is ``(d + 0.5) * s - 0.5``.
    Albumentations' bbox crop is geometric — ``(src - x1) * S / crop`` with no half-pixel
    term — which is also why ``A.Affine`` returns a separate ``bbox_matrix``. Composing
    the image chain with the geometric convention shifts the picture by
    ``0.5 * (1 - S / crop)`` pixels, half a pixel when the crop is upscaled 2x, while
    the boxes stay put. Nothing raises; the labels simply stop matching the pixels.
    Measured: with the geometric matrix the 99th percentile of the difference against
    the chain is up to 14 grey levels, with the pixel-centre matrix it is 1.

    So ``_image_matrix`` carries the half-pixel terms and ``_box_matrix`` does not.
    They are not meant to agree.

    Only axis-aligned chains are supported. With rotation or shear, clipping a box once
    at the end is not the same as clipping it after each step, and this class would
    silently produce different boxes; it raises instead.

    Args:
        size: Output square side.
        scale: Crop area fraction range, as ``A.RandomResizedCrop``.
        ratio: Crop aspect ratio range, as ``A.RandomResizedCrop``.
        flip_prob: Probability of the horizontal flip.
        translate: Affine translation as a fraction of output size.
        affine_scale: Affine scale range.
        degrees: Rotation. Must be 0.
        shear: Shear. Must be 0.
        pad_value: Fill for pixels pulled in from outside the source.
        min_box_size: Drop boxes thinner than this, in output pixels. Matches the
            ``min_width`` / ``min_height`` of ``_BBOX_PARAMS``.
    """

    def __init__(
        self,
        size: int = 640,
        scale: tuple[float, float] = (0.05, 0.8),
        ratio: tuple[float, float] = (0.75, 1.33),
        flip_prob: float = 0.5,
        translate: float = 0.25,
        affine_scale: tuple[float, float] = (0.5, 1.5),
        degrees: float = 0.0,
        shear: float = 0.0,
        pad_value: int = 114,
        min_box_size: float = 2.0,
    ):
        if degrees != 0.0 or shear != 0.0:
            raise ValueError(
                "RandomResizedCropFlipAffine fuses an axis-aligned chain only. "
                f"Got degrees={degrees}, shear={shear}. Clipping boxes once at the end "
                "is equivalent to clipping them at each step only while the transform "
                "stays axis-aligned; under rotation or shear it is not, and the fused "
                "boxes would differ from the chain's without any error being raised. "
                "Use RandomResizedCrop + HorizontalFlip + RandomAffine for those."
            )

        self.size = size
        self.flip_prob = flip_prob
        self.pad_value = pad_value
        self.min_box_size = min_box_size

        # Sampling is delegated to Albumentations so the distributions cannot drift.
        self._crop = A.RandomResizedCrop(size=(size, size), scale=scale, ratio=ratio, p=1.0)
        self._affine = A.Affine(
            scale=affine_scale,
            translate_percent={"x": (-translate, translate), "y": (-translate, translate)},
            rotate=(0, 0),
            shear=(0, 0),
            border_mode=cv2.BORDER_CONSTANT,
            fill=pad_value,
            p=1.0,
        )

    # -- matrices -----------------------------------------------------------

    def _crop_matrix(self, crop: tuple[int, int, int, int], *, pixel_centre: bool) -> np.ndarray:
        x1, y1, x2, y2 = crop
        sx, sy = self.size / (x2 - x1), self.size / (y2 - y1)
        if pixel_centre:
            tx, ty = (0.5 - x1) * sx - 0.5, (0.5 - y1) * sy - 0.5
        else:
            tx, ty = -sx * x1, -sy * y1
        return np.array([[sx, 0.0, tx], [0.0, sy, ty], [0.0, 0.0, 1.0]])

    def _flip_matrix(self, *, pixel_centre: bool) -> np.ndarray:
        # cv2.flip sends column x to W-1-x; a normalised box coordinate goes to 1-x,
        # which in pixels is W-x. One pixel apart, and each is right for its target.
        offset = self.size - 1 if pixel_centre else self.size
        return np.array([[-1.0, 0.0, offset], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])

    # -- api ----------------------------------------------------------------

    def sample_params(self, height: int, width: int) -> dict:
        """Draw one set of parameters. Split out so tests can drive both paths alike."""
        crop = self._crop.get_params_dependent_on_data({"shape": (height, width, 3)}, {})["crop_coords"]
        affine = self._affine.get_params_dependent_on_data({"shape": (self.size, self.size, 3)}, {})
        return {
            "crop": tuple(crop),
            "flip": random.random() < self.flip_prob,
            "matrix": np.asarray(affine["matrix"], dtype=np.float64),
            "bbox_matrix": np.asarray(affine["bbox_matrix"], dtype=np.float64),
        }

    def image_matrix(self, params: dict) -> np.ndarray:
        eye = np.eye(3)
        return (
            params["matrix"]
            @ (self._flip_matrix(pixel_centre=True) if params["flip"] else eye)
            @ self._crop_matrix(params["crop"], pixel_centre=True)
        )

    def box_matrix(self, params: dict) -> np.ndarray:
        eye = np.eye(3)
        return (
            params["bbox_matrix"]
            @ (self._flip_matrix(pixel_centre=False) if params["flip"] else eye)
            @ self._crop_matrix(params["crop"], pixel_centre=False)
        )

    def apply(self, image: np.ndarray, targets: np.ndarray, params: dict) -> tuple[np.ndarray, np.ndarray]:
        h, w = image.shape[:2]
        out = cv2.warpAffine(
            image,
            self.image_matrix(params)[:2],
            (self.size, self.size),
            flags=cv2.INTER_LINEAR,
            borderMode=cv2.BORDER_CONSTANT,
            borderValue=(self.pad_value,) * image.shape[2] if image.ndim == 3 else self.pad_value,
        )

        if not len(targets):
            return out, targets

        m = self.box_matrix(params)
        cx, cy, bw, bh = targets[:, 1] * w, targets[:, 2] * h, targets[:, 3] * w, targets[:, 4] * h
        x1, y1, x2, y2 = cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2

        # Map both corners. Axis-aligned, so the transformed corners are still the
        # extremes; min/max guards against a negative scale (the flip).
        px = m[0, 0] * np.stack([x1, x2]) + m[0, 2]
        py = m[1, 1] * np.stack([y1, y2]) + m[1, 2]
        nx1, nx2 = px.min(0), px.max(0)
        ny1, ny2 = py.min(0), py.max(0)

        np.clip(nx1, 0, self.size, out=nx1)
        np.clip(nx2, 0, self.size, out=nx2)
        np.clip(ny1, 0, self.size, out=ny1)
        np.clip(ny2, 0, self.size, out=ny2)

        keep = ((nx2 - nx1) >= self.min_box_size) & ((ny2 - ny1) >= self.min_box_size)
        if not keep.any():
            return out, np.zeros((0, 5), dtype=targets.dtype)

        kept = np.empty((int(keep.sum()), 5), dtype=targets.dtype)
        kept[:, 0] = targets[keep, 0]
        kept[:, 1] = (nx1[keep] + nx2[keep]) / 2 / self.size
        kept[:, 2] = (ny1[keep] + ny2[keep]) / 2 / self.size
        kept[:, 3] = (nx2[keep] - nx1[keep]) / self.size
        kept[:, 4] = (ny2[keep] - ny1[keep]) / self.size
        return out, kept

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        return self.apply(image, targets, self.sample_params(*image.shape[:2]))

sample_params(height, width)

Draw one set of parameters. Split out so tests can drive both paths alike.

Source code in src/modern_yolonas/data/transforms.py 
423
424
425
426
427
428
429
430
431
432
def sample_params(self, height: int, width: int) -> dict:
    """Draw one set of parameters. Split out so tests can drive both paths alike."""
    crop = self._crop.get_params_dependent_on_data({"shape": (height, width, 3)}, {})["crop_coords"]
    affine = self._affine.get_params_dependent_on_data({"shape": (self.size, self.size, 3)}, {})
    return {
        "crop": tuple(crop),
        "flip": random.random() < self.flip_prob,
        "matrix": np.asarray(affine["matrix"], dtype=np.float64),
        "bbox_matrix": np.asarray(affine["bbox_matrix"], dtype=np.float64),
    }

RandomChannelSwap

Randomly swap BGR channel order to RGB (and vice-versa) via Albumentations.

Adds photometric variety without touching bounding boxes.

Parameters:

Name Type Description Default
p float

Probability of swapping channels.

 0.5
Source code in src/modern_yolonas/data/transforms.py 
496
497
498
499
500
501
502
503
504
505
506
507
508
509
class RandomChannelSwap:
    """Randomly swap BGR channel order to RGB (and vice-versa) via Albumentations.

    Adds photometric variety without touching bounding boxes.

    Args:
        p: Probability of swapping channels.
    """

    def __init__(self, p: float = 0.5):
        self._aug = A.ChannelShuffle(p=p)

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        return self._aug(image=image)["image"], targets

CenterCrop

Crop the center of the image to size × size pixels.

Bounding boxes that fall outside the cropped region are discarded; those that overlap are clipped to the new canvas by _BBOX_PARAMS.

Parameters:

Name Type Description Default
size int

Output square side length in pixels.

 640
Source code in src/modern_yolonas/data/transforms.py 
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
class CenterCrop:
    """Crop the center of the image to ``size`` × ``size`` pixels.

    Bounding boxes that fall outside the cropped region are discarded;
    those that overlap are clipped to the new canvas by ``_BBOX_PARAMS``.

    Args:
        size: Output square side length in pixels.
    """

    def __init__(self, size: int = 640):
        self._aug = A.Compose(
            [A.CenterCrop(height=size, width=size, p=1.0)],
            bbox_params=_BBOX_PARAMS,
        )

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        bboxes, labels = _to_albu(targets)
        r = self._aug(image=image, bboxes=bboxes, class_labels=labels)
        return r["image"], _from_albu(r["bboxes"], r["class_labels"], targets.dtype)

Mosaic

4-image mosaic augmentation.

Source code in src/modern_yolonas/data/transforms.py 
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
class Mosaic:
    """4-image mosaic augmentation."""

    def __init__(self, dataset: BaseDetectionDataset, input_size: int = 640, prob: float = 1.0):
        self.dataset = dataset
        self.input_size = input_size
        self.prob = prob
        self.enabled = True

    def __call__(self, index: int) -> tuple[np.ndarray, np.ndarray]:
        if not self.enabled or random.random() >= self.prob:
            return self.dataset.load_raw(index)

        s = self.input_size
        yc, xc = (int(random.uniform(s * 0.5, s * 1.5)) for _ in range(2))

        indices = [index] + [random.randint(0, len(self.dataset) - 1) for _ in range(3)]
        mosaic_img = np.full((s * 2, s * 2, 3), 114, dtype=np.uint8)
        all_targets = []

        for i, idx in enumerate(indices):
            img, targets = self.dataset.load_raw(idx)
            h, w = img.shape[:2]

            if i == 0:
                x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc
                x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h
            elif i == 1:
                x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
                x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
            elif i == 2:
                x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
                x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
            else:
                x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
                x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)

            mosaic_img[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]
            pad_w = x1a - x1b
            pad_h = y1a - y1b

            if len(targets):
                targets = targets.copy()
                # Convert to pixel coords, offset, then back to normalized
                targets[:, 1] = (targets[:, 1] * w + pad_w) / (s * 2)
                targets[:, 2] = (targets[:, 2] * h + pad_h) / (s * 2)
                targets[:, 3] = targets[:, 3] * w / (s * 2)
                targets[:, 4] = targets[:, 4] * h / (s * 2)
                all_targets.append(targets)

        targets = np.concatenate(all_targets, 0) if all_targets else np.zeros((0, 5))

        # Crop to input_size
        crop_x = int(random.uniform(0, s))
        crop_y = int(random.uniform(0, s))
        mosaic_img = mosaic_img[crop_y : crop_y + s, crop_x : crop_x + s]

        if len(targets):
            targets = targets.copy()
            # Targets are normalized to the 2s x 2s canvas; the crop is s x s, so every
            # coordinate is renormalized by 2. Widths and heights need that factor just
            # as much as the centres do — omitting it emits every box at half size.
            targets[:, 1] = targets[:, 1] * 2 - crop_x / s
            targets[:, 2] = targets[:, 2] * 2 - crop_y / s
            targets[:, 3] = targets[:, 3] * 2
            targets[:, 4] = targets[:, 4] * 2

            # Clip to the crop window, then drop boxes whose centre left it or that
            # survive only as a sliver.
            x1 = np.clip(targets[:, 1] - targets[:, 3] / 2, 0.0, 1.0)
            y1 = np.clip(targets[:, 2] - targets[:, 4] / 2, 0.0, 1.0)
            x2 = np.clip(targets[:, 1] + targets[:, 3] / 2, 0.0, 1.0)
            y2 = np.clip(targets[:, 2] + targets[:, 4] / 2, 0.0, 1.0)
            centre_inside = (
                (targets[:, 1] > 0) & (targets[:, 1] < 1)
                & (targets[:, 2] > 0) & (targets[:, 2] < 1)
            )
            targets[:, 1] = (x1 + x2) / 2
            targets[:, 2] = (y1 + y2) / 2
            targets[:, 3] = x2 - x1
            targets[:, 4] = y2 - y1

            valid = centre_inside & (targets[:, 3] > 0.002) & (targets[:, 4] > 0.002)
            targets = targets[valid]

        return mosaic_img, targets

Mixup

Mixup augmentation for detection.

Should be placed in the pipeline after LetterboxResize and before Normalize, so both images are already square uint8.

Parameters:

Name Type Description Default
dataset BaseDetectionDataset

Dataset exposing a load_raw(index) method.

 required
alpha float

Beta distribution alpha parameter.

 1.5
beta float

Beta distribution beta parameter.

 1.5
prob float

Per-sample probability of applying mixup.

 0.5
Source code in src/modern_yolonas/data/transforms.py 
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
class Mixup:
    """Mixup augmentation for detection.

    Should be placed in the pipeline **after** ``LetterboxResize`` and
    **before** ``Normalize``, so both images are already square uint8.

    Args:
        dataset: Dataset exposing a ``load_raw(index)`` method.
        alpha: Beta distribution ``alpha`` parameter.
        beta: Beta distribution ``beta`` parameter.
        prob: Per-sample probability of applying mixup.
    """

    def __init__(
        self,
        dataset: BaseDetectionDataset,
        alpha: float = 1.5,
        beta: float = 1.5,
        prob: float = 0.5,
    ):
        self.dataset = dataset
        self.alpha = alpha
        self.beta = beta
        self.prob = prob
        self.enabled = True
        self.inner_transforms: Compose | None = None

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        if not self.enabled or random.random() >= self.prob:
            return image, targets

        idx2 = random.randint(0, len(self.dataset) - 1)
        img2, targets2 = self.dataset.load_raw(idx2)

        # Per-image augmentations for the mixed-in image, so it is not a pristine
        # sample pasted onto an augmented one.
        if self.inner_transforms is not None:
            img2, targets2 = self.inner_transforms(img2, targets2)

        # Letterbox the second image to match the (already resized) first image
        target_h, target_w = image.shape[:2]
        h2, w2 = img2.shape[:2]
        if (h2, w2) != (target_h, target_w):
            scale = min(target_h / h2, target_w / w2)
            new_h, new_w = int(round(h2 * scale)), int(round(w2 * scale))
            img2 = cv2.resize(img2, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
            top  = (target_h - new_h) // 2
            left = (target_w - new_w) // 2
            padded = np.full((target_h, target_w, 3), 114, dtype=np.uint8)
            padded[top:top + new_h, left:left + new_w] = img2
            img2 = padded
            if len(targets2):
                targets2 = targets2.copy()
                targets2[:, 1] = (targets2[:, 1] * new_w + left) / target_w
                targets2[:, 2] = (targets2[:, 2] * new_h + top)  / target_h
                targets2[:, 3] = targets2[:, 3] * new_w / target_w
                targets2[:, 4] = targets2[:, 4] * new_h / target_h

        r = float(np.random.beta(self.alpha, self.beta))
        # cv2 blends in one SIMD pass over uint8. Promoting both images to float32
        # first allocated two 4.9 MB arrays per call and cost 1.47 ms against 0.30 ms,
        # measured 2026-09-19. The results differ by at most one grey level, and in
        # cv2's favour: it rounds, where the float32 cast truncated.
        mixed = cv2.addWeighted(image, r, img2, 1.0 - r, 0.0)

        if len(targets) and len(targets2):
            combined = np.concatenate([targets, targets2], 0)
        elif len(targets):
            combined = targets
        else:
            combined = targets2

        return mixed, combined

LetterboxResize

Resize with aspect ratio preservation and center padding.

Parameters:

Name Type Description Default
target_size int

Output square dimension.

 640
pad_value int

Pixel value for padding (default 114, matching YOLO convention).

 114
Source code in src/modern_yolonas/data/transforms.py 
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
class LetterboxResize:
    """Resize with aspect ratio preservation and center padding.

    Args:
        target_size: Output square dimension.
        pad_value: Pixel value for padding (default 114, matching YOLO convention).
    """

    def __init__(self, target_size: int = 640, pad_value: int = 114):
        self.target_size = target_size
        self.pad_value = pad_value

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        h, w = image.shape[:2]
        scale = self.target_size / max(h, w)
        new_h, new_w = int(round(h * scale)), int(round(w * scale))
        image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

        pad_h = self.target_size - new_h
        pad_w = self.target_size - new_w
        top = pad_h // 2
        left = pad_w // 2

        padded = np.full((self.target_size, self.target_size, 3), self.pad_value, dtype=np.uint8)
        padded[top : top + new_h, left : left + new_w] = image

        if len(targets):
            targets = targets.copy()
            # Adjust for padding (targets are normalized)
            targets[:, 1] = (targets[:, 1] * new_w + left) / self.target_size
            targets[:, 2] = (targets[:, 2] * new_h + top) / self.target_size
            targets[:, 3] = targets[:, 3] * new_w / self.target_size
            targets[:, 4] = targets[:, 4] * new_h / self.target_size

        return padded, targets

Normalize

Convert HWC BGR uint8 to CHW RGB, optionally leaving the scaling to the GPU.

With dtype="uint8" the array stays 1.23 MB per 640x640 sample instead of 4.92 MB, which is what then crosses the collate function, the pin_memory thread and the PCIe bus; YoloNASLightningModule.on_after_batch_transfer divides by 255 on the device, where it is free. Measured 2026-09-19: 3.23 ms to 0.28 ms per sample, and the reconstructed float32 is bit-identical.

The default stays "float32" so that every consumer that reads batches directly — yolonas eval, yolonas quantize, the examples — is unaffected.

Parameters:

Name Type Description Default
dtype str

"float32" to divide by 255 here, "uint8" to defer it.

 'float32'
Source code in src/modern_yolonas/data/transforms.py 
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
class Normalize:
    """Convert HWC BGR uint8 to CHW RGB, optionally leaving the scaling to the GPU.

    With ``dtype="uint8"`` the array stays 1.23 MB per 640x640 sample instead of
    4.92 MB, which is what then crosses the collate function, the ``pin_memory``
    thread and the PCIe bus; ``YoloNASLightningModule.on_after_batch_transfer``
    divides by 255 on the device, where it is free. Measured 2026-09-19:
    3.23 ms to 0.28 ms per sample, and the reconstructed float32 is bit-identical.

    The default stays ``"float32"`` so that every consumer that reads batches
    directly — ``yolonas eval``, ``yolonas quantize``, the examples — is unaffected.

    Args:
        dtype: ``"float32"`` to divide by 255 here, ``"uint8"`` to defer it.
    """

    def __init__(self, dtype: str = "float32"):
        if dtype not in ("float32", "uint8"):
            raise ValueError(f"dtype must be 'float32' or 'uint8', got {dtype!r}")
        self.dtype = dtype

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        # BGR → RGB and HWC → CHW in one go; ascontiguousarray materialises the
        # result once, where the old chain copied for the flip and again for the cast.
        image = np.ascontiguousarray(image[:, :, ::-1].transpose(2, 0, 1))
        if self.dtype == "float32":
            image = image.astype(np.float32) / 255.0
        return image, targets

TrainTransformPipeline

Full training augmentation pipeline with Mosaic and Mixup support.

Pipeline order (matching super-gradients): Mosaic → per_image_transforms (RandomAffine, RandomChannelShuffle, HSV, HorizontalFlip) → Mixup → final_transforms (LetterboxResize, Normalize)

Source code in src/modern_yolonas/data/transforms.py 
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
class TrainTransformPipeline:
    """Full training augmentation pipeline with Mosaic and Mixup support.

    Pipeline order (matching super-gradients):
    Mosaic → per_image_transforms (RandomAffine, RandomChannelShuffle, HSV, HorizontalFlip)
    → Mixup → final_transforms (LetterboxResize, Normalize)
    """

    def __init__(
        self,
        mosaic: Mosaic | None,
        per_image_transforms: Compose,
        mixup: Mixup | None,
        final_transforms: Compose,
    ):
        self.mosaic = mosaic
        self.per_image_transforms = per_image_transforms
        self.mixup = mixup
        self.final_transforms = final_transforms

    def apply(self, index: int, load_raw_fn) -> tuple[np.ndarray, np.ndarray]:
        """Dataset-aware pipeline entry point called from __getitem__."""
        # 1. Mosaic (or fallback to load_raw)
        if self.mosaic is not None:
            image, targets = self.mosaic(index)
        else:
            image, targets = load_raw_fn(index)

        # 2. Per-image transforms
        image, targets = self.per_image_transforms(image, targets)

        # 3. Mixup
        if self.mixup is not None:
            image, targets = self.mixup(image, targets)

        # 4. Final transforms (resize + normalize)
        image, targets = self.final_transforms(image, targets)

        return image, targets

    def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Fallback: per-image + final only (no Mosaic/Mixup)."""
        image, targets = self.per_image_transforms(image, targets)
        image, targets = self.final_transforms(image, targets)
        return image, targets

    def disable_mosaic_mixup(self):
        """Disable Mosaic and Mixup (called by CloseMosaicCallback)."""
        if self.mosaic is not None:
            self.mosaic.enabled = False
        if self.mixup is not None:
            self.mixup.enabled = False

apply(index, load_raw_fn)

Dataset-aware pipeline entry point called from getitem.

Source code in src/modern_yolonas/data/transforms.py 
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
def apply(self, index: int, load_raw_fn) -> tuple[np.ndarray, np.ndarray]:
    """Dataset-aware pipeline entry point called from __getitem__."""
    # 1. Mosaic (or fallback to load_raw)
    if self.mosaic is not None:
        image, targets = self.mosaic(index)
    else:
        image, targets = load_raw_fn(index)

    # 2. Per-image transforms
    image, targets = self.per_image_transforms(image, targets)

    # 3. Mixup
    if self.mixup is not None:
        image, targets = self.mixup(image, targets)

    # 4. Final transforms (resize + normalize)
    image, targets = self.final_transforms(image, targets)

    return image, targets

__call__(image, targets)

Fallback: per-image + final only (no Mosaic/Mixup).

Source code in src/modern_yolonas/data/transforms.py 
804
805
806
807
808
def __call__(self, image: np.ndarray, targets: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Fallback: per-image + final only (no Mosaic/Mixup)."""
    image, targets = self.per_image_transforms(image, targets)
    image, targets = self.final_transforms(image, targets)
    return image, targets

disable_mosaic_mixup()

Disable Mosaic and Mixup (called by CloseMosaicCallback).

Source code in src/modern_yolonas/data/transforms.py 
810
811
812
813
814
815
def disable_mosaic_mixup(self):
    """Disable Mosaic and Mixup (called by CloseMosaicCallback)."""
    if self.mosaic is not None:
        self.mosaic.enabled = False
    if self.mixup is not None:
        self.mixup.enabled = False