""" STEP 2 — augment.py ==================== Expands your small labeled dataset using augmentations tuned specifically for retail shelf / packaged-goods detection. ~4 images → ~200 augmented images per class (configurable via AUG_PER_IMAGE below) Install: pip install albumentations opencv-python tqdm pyyaml Usage: python augment.py """ import os import cv2 import yaml import random import shutil import numpy as np from pathlib import Path from tqdm import tqdm import albumentations as A # ── Configuration ──────────────────────────────────────────────────────────────── SRC_IMAGES = Path("dataset/images/train") SRC_LABELS = Path("dataset/labels/train") OUT_IMAGES = Path("dataset/images/train") # Augmented added in-place OUT_LABELS = Path("dataset/labels/train") VAL_IMAGES = Path("dataset/images/val") VAL_LABELS = Path("dataset/labels/val") AUG_PER_IMAGE = 50 # augmented copies per original image VAL_SPLIT = 0.15 # fraction held out for validation SEED = 42 IMG_SIZE = 640 # ── Augmentation pipeline ───────────────────────────────────────────────────── # Chosen specifically for packaged-goods on retail shelves: # • Perspective/affine → products photographed at angles # • Brightness/contrast → fluorescent vs natural shelf lighting # • Hue/saturation → the Famico red looks different under warm vs cool light # • Blur → camera shake, shallow depth of field # • Noise → phone camera noise # • Shadow → shelf edge shadows across packages # • Cutout → partial occlusion by other products # • CLAHE → over/under-exposed shots def build_augmentation_pipeline() -> A.Compose: bbox_params = A.BboxParams( format = "yolo" , # cx, cy, w, h normalised 0-1 label_fields = ["class_ids"], min_area = 0.001 , # drop boxes that shrink below 0.1% of image min_visibility = 0.03 , # drop boxes that become >70% occluded ) return A.Compose([ # ── Geometry ──────────────────────────────────────────────────────── A.HorizontalFlip(p=0.5), A.ShiftScaleRotate( shift_limit = 0.08, scale_limit = 0.25, rotate_limit = 12, border_mode = cv2.BORDER_CONSTANT, value = 0, p = 0.8 ), A.Perspective( scale = (0.03, 0.10), keep_size = True, p = 0.5, ), A.RandomResizedCrop( size = (IMG_SIZE, IMG_SIZE), scale = (0.70, 1.00), ratio = (0.75, 1.33), p = 0.4, ), A.RandomBrightnessContrast( brightness_limit = 0.35, constrast_limit = 0.35, p = 0.8, ), A.HueSaturationValue( hue_shift_limit = 12, sat_shift_limit = 30, val_shit_limit = 25, p = 0.7, ), A.CLAHE( clip_limit = 4.0, tile_grid_size = (8, 8), p = 0.3, ), A.RGBShift( r_shift_limit = 15, g_shift_limit = 10, b_shift_limit = 10, p = 0.4, ), # ── Blur / Noise ───────────────────────────────────────────────────────────────── A.OneOf([ A.GaussianBlur(blur_limit=(3, 7)), A.MotionBlur (blur_limit=7), A.MedianBlur (blur_limit=5), ], p=0.5), A.GaussNoise( std_range = (0.01, 0.05), p = 0.4, ), A.ImageCompression( quality_range = (60, 95), p = 0.4, ), # ── Occlusion Simulation ────────────────────────────────────────────────────────── A.CoarseDropout( num_holes_range = (1, 4), hole_height_range = (0.04, 0.12), hole_width_range = (0.04, 0.12), fill = 0, p = 0.35, ), # ── Final Resize to YOLO input size ─────────────────────────────────────────────── A.LongestMaxSize(max_size=IMG_SIZE), A.PadIfNeeded( min_height = IMG_SIZE, min_width = IMG_SIZE, border_mode = cv2.BORDER_CONSTANT, value = 112, # Standard YOLO gray padding. ), ], bbox_params=bbox_params) # ── YOLO label helpers ──────────────────────────────────────────────────────── def read_yolo_label(label_path: Path): """ Retruns (class_ids: list[int], bboxes: list[tuple])""" class_ids, bboxes = [], [] if not label_path.exists(): return class_ids, bboxes with open(label_path) as f: for line in f: parts = line.strip().split() if len(parts) != 5: continue cid = int(parts[0]) bbox = tuple(float(x) for x in parts[1:]) class_ids.append(cid) bboxes .append(bbox) return class_ids, bboxes def write_yolo_label(label_path: Path, class_ids, bboxes): with open(label_path, "w") as f: for cid, (cx, cy, w, h) in zip(class_ids, bboxes): f.write(f"{cid} {cx:6f} {cy:6f} {w:.6f} {h:.6f}\n") def clamp_bbox(bbox): cx, cy, w, h = bbox cx = max(0.0, min(1.0, cx)) cy = max(0.0, min(1.0, cy)) w = max(0.001, min(1.0 - cx + w/2, w)) h = max(0.001, min(1.0 - cy + h/2, h)) return (cx, cy, w, h) # ── Train / val split ───────────────────────────────────────────────────────── def split_val(image_paths, val_ratio=VAL_SPLIT): random.seed(SEED) shuffled = list(image_paths) random.shuffle(shuffled) n_val = max(1, int(len(shuffled) * val_ratio)) return shuffled[n_val:], shuffled[:n_val] # train, val def copy_to_val(image_path: Path): stem = image_path.stem label_path = SRC_LABELS / f"{stem}.txt" VAL_IMAGES.mkdir(parents=True, exist_ok=True) VAL_LABELS.mkdir(parents=True, exist_ok=True) shutil.copy(image_path, VAL_IMAGES / image_path.name) if label_path.exists(): shutil.copy(label_path, VAL_LABELS / f"{stem}.txt") # ── Train / val split ───────────────────────────────────────────────────────── def main(): random.seed(SEED) np.random.seed(SEED) exts = { ".jpg", ".jpeg", ".png", ".bmp", ".webp"} originals = sorted([ p for p in SRC_IMAGES.iterdir() if p.suffix.lower() in exts and not p.stem.startswith("aug_") ]) if not originals: print(f"[ERROR] No images found in {SRC_IMAGES}") print(" Put your labeled images there first, then run this script.") return print(f"Found {len(originals)} original images.") # Split before augmenting - val set contains ONLY original photos train_paths, val_paths = split_val(originals, VAL_SPLIT) print(f"Validation set: {len(val_paths)} originals (never augmented)") print(f"Training set: {len(train_paths)} originals -> " f"{len(train_paths) * AUG_PER_IMAGE} augmented\n") for vp in val_paths: copy_to_val(vp) transform = build_augmentation_pipeline() total_written = 0 for img_path in tqdm(train_paths, desc="Augmenting"): image = cv2.imread(str(img_path)) if image is None: print(f" [WANR] Cannot read {img_path}, skipping.") continue image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) class_ids, bboxes = read_yolo_label( SRC_LABELS / f"{img_path.stem}.txt" ) if not bboxes: print(f" [WARN] No Labels for {img_path.name}, skipping.") continue for i in range(AUG_PER_IMAGE): try: result = transform( image = image, bboxes = bboxes, class_ids = class_ids, ) except Exception as e: print(f" [WARN] Augmentation failed for {img_path.name} #{i}: {e}") continue aug_bboxes = [clamp_bbox(b) for b in result["bboxes"]] aug_cids = result["class_ids"] if not aug_bboxes: continue # All Boxes fell out of frame. aug_name = f"aug_{img_path.stem}_{i:04d}" out_img = OUT_IMAGES / f"{aug_name}.jpg" out_label = OUT_LABELS / f"{aug_name}.txt" aug_img = cv2.cvtColor(result["image"], cv2.COLOR_RGB2BGR) cv2.imwrite(str(out_img), aug_img, [cv2.IMWRITE_JPEG_QUALITY, 92]) write_yolo_label(out_label, aug_cids, aug_bboxes) total_written += 1 print(f"\nDone - {total_written} augmented images written.") print(f" Train: {OUT_IMAGES} Val: {VAL_IMAGES}") print("\nRun next: python train.py") if __name__ == "__main__": main()