added augment.py

.gitignore

@@ -1 +1,3 @@
+# Dataset
 dataset/
+dataset.yaml

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()
+

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

enviroment.yml

@@ -8,6 +8,7 @@ dependencies:
   - opencv
   - tqdm
   - pyyaml
+  - simsimd
   #- torch
   #- numpy
   #- opencv-contrib-python

prepare_dataset.py

@@ -60,7 +60,7 @@ def Structure_Create():
     print("[1/3] Folder structure created:")
 
     for d in DIRS:
-        print("         {d}/")
+        print(f"         {d}/")
 
 def YAML_Write():
     cfg = {