From 40e77d1034613da59bad3c400091ef86d0de63cc Mon Sep 17 00:00:00 2001 From: Hamza Zakaria Date: Mon, 27 Apr 2026 10:32:17 +0100 Subject: [PATCH] added augment.py --- .gitignore | 2 + augment.py | 289 +++++++++++++++++++++++++++++++++++++++++++++ dataset.yml | 18 +++ enviroment.yml | 1 + prepare_dataset.py | 10 +- 5 files changed, 315 insertions(+), 5 deletions(-) create mode 100644 augment.py create mode 100644 dataset.yml diff --git a/.gitignore b/.gitignore index 93db21b..8a74c77 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ +# Dataset dataset/ +dataset.yaml diff --git a/augment.py b/augment.py new file mode 100644 index 0000000..e5b419d --- /dev/null +++ b/augment.py @@ -0,0 +1,289 @@ +""" +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() + diff --git a/dataset.yml b/dataset.yml new file mode 100644 index 0000000..f1b4652 --- /dev/null +++ b/dataset.yml @@ -0,0 +1,18 @@ +path: C:\Users\TRIZ\clones\YOLO-Coffee-Object-Detection\dataset +train: images/train +val: images/val +nc: 13 +names: +- '1001' +- Aroma +- Aroma Gold +- Bonal +- Facto +- Famico +- Famico Exclusive +- Fegalo +- Gosto +- Molino +- Oscar +- Primo +- Siglo diff --git a/enviroment.yml b/enviroment.yml index 1268744..65af4b8 100644 --- a/enviroment.yml +++ b/enviroment.yml @@ -8,6 +8,7 @@ dependencies: - opencv - tqdm - pyyaml + - simsimd #- torch #- numpy #- opencv-contrib-python diff --git a/prepare_dataset.py b/prepare_dataset.py index f7fa30d..8d9b1d1 100644 --- a/prepare_dataset.py +++ b/prepare_dataset.py @@ -60,16 +60,16 @@ def Structure_Create(): print("[1/3] Folder structure created:") for d in DIRS: - print(" {d}/") + print(f" {d}/") def YAML_Write(): cfg = { - "path": os.path.abspath(DATASET_ROOT), + "path" : os.path.abspath(DATASET_ROOT), "train": "images/train", - "val": "images/val", - "nc": len(CLASS_NAMES), + "val" : "images/val", + "nc" : len(CLASS_NAMES), "names": CLASS_NAMES, - } + } yaml_path = "dataset.yml" with open(yaml_path, "w") as f: