Updated data/ + added new scripts

.gitignore

@@ -1,3 +1,6 @@
+.label_cache/
+.color_cache
+
 # Auto Label
 .autolabel_staging/
 _rejected_/

3_inference.py

@@ -47,7 +47,7 @@ from tqdm import tqdm
 def parse_args():
     p = argparse.ArgumentParser(description="YOLOv8 detect + classify inference")
     p.add_argument("--detector_weights",   default="detector/best.pt", help="YOLOv8n .pt (fine-tuned)")
-    p.add_argument("--classifier_weights", default="best.pt", help="Classifier best.pt from script 2")
+    p.add_argument("--classifier_weights", default="runs/classify/best.pt", help="Classifier best.pt from script 2")
     p.add_argument("--source",             required=True, help="Image / folder / video / webcam index")
     p.add_argument("--det_conf",    type=float, default=0.30,  help="YOLO detection confidence")
     p.add_argument("--det_iou",     type=float, default=0.45,  help="YOLO NMS IoU")

api_server.py

@@ -64,7 +64,7 @@ from ultralytics import YOLO
 # ─────────────────────────── config ──────────────────────────────────────────
 
 DETECTOR_WEIGHTS   = os.getenv("DETECTOR_WEIGHTS",   "detector/best.pt")
-CLASSIFIER_WEIGHTS = os.getenv("CLASSIFIER_WEIGHTS", "best.pt")
+CLASSIFIER_WEIGHTS = os.getenv("CLASSIFIER_WEIGHTS", "runs/classify/best.pt")
 DET_CONF  = float(os.getenv("DET_CONF",  "0.25"))
 CLS_CONF  = float(os.getenv("CLS_CONF",  "0.75"))
 DET_IOU   = float(os.getenv("DET_IOU",   "0.45"))

auto_label.py

@@ -55,7 +55,7 @@ def parse_args():
     p = argparse.ArgumentParser(description="Auto-label crops and review in browser")
     p.add_argument("--source",               required=True)
     p.add_argument("--detector_weights",     default="detector/best.pt")
-    p.add_argument("--classifier_weights",   default="best.pt")
+    p.add_argument("--classifier_weights",   default="runs/classify/best.pt")
     p.add_argument("--output_dir",           default="data")
     p.add_argument("--staging_dir",          default=".autolabel_staging",
                    help="Temp folder for staged crops before commit (hidden by default)")

balance_and_augment.py

@@ -0,0 +1,453 @@
+"""
+balance_and_augment.py — Dataset Balancer & Augmentor
+=======================================================
+Analyses your crops_dataset/ (or data/) for class imbalance,
+then generates augmented images until every class hits the
+same target count.
+
+Strategy
+--------
+1. Count images per class across train split
+2. Find the target count  (--target_count or auto = max class × scale)
+3. For each under-represented class, generate augmented copies
+   until it reaches the target
+4. Augmented images are written beside the originals — originals
+   are never touched
+
+Augmentation pipeline (albumentations)
+---------------------------------------
+Light  (always applied):
+  HorizontalFlip, slight rotation, brightness/contrast jitter,
+  hue/saturation shift
+
+Medium (random subset):
+  Perspective, GridDistortion, GaussianBlur, Sharpen,
+  CoarseDropout (occlusion), JPEG compression noise
+
+Heavy  (low probability):
+  RandomSunFlare, RandomShadow, RandomFog, ChannelShuffle
+
+Usage
+-----
+python balance_and_augment.py \
+    --data_dir      crops_dataset \
+    --split         train \
+    --target_count  300 \
+    --max_scale     3.0 \
+    --workers       8 \
+    --dry_run
+
+Arguments
+---------
+--data_dir      Root of the split dataset (contains train/ val/ test/)
+--split         Which split to augment (default: train)
+--target_count  Explicit target per class. Omit to use auto.
+--max_scale     Auto target = max_class_count × max_scale  (default 2.0)
+                Capped so dominant classes are not inflated further.
+--min_count     Skip augmentation for classes already above this fraction
+                of the target (default 0.95 — within 5% is fine)
+--workers       Parallel workers (default: cpu_count - 1)
+--dry_run       Print the plan without writing any files
+--val_split     Also balance val split proportionally (0 = skip, default 0)
+--seed          Random seed for reproducibility (default 42)
+--suffix        Suffix added to augmented file stems (default: _aug)
+--quality       JPEG save quality for augmented images (default: 92)
+"""
+
+import argparse
+import math
+import multiprocessing
+import random
+import shutil
+from collections import Counter
+from pathlib import Path
+
+import cv2
+import numpy as np
+
+try:
+    import albumentations as A
+except ImportError:
+    raise ImportError("Run:  pip install albumentations")
+
+
+# ─────────────────────────── args ────────────────────────────────────────────
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Balance and augment a classification dataset",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    p.add_argument("--data_dir",      default="crops_dataset")
+    p.add_argument("--split",         default="train")
+    p.add_argument("--target_count",  type=int,   default=None,
+                   help="Target images per class. Auto if omitted.")
+    p.add_argument("--max_scale",     type=float, default=2.0,
+                   help="Auto target = largest_class × max_scale")
+    p.add_argument("--min_count",     type=float, default=0.95,
+                   help="Classes already at ≥ this fraction of target are skipped")
+    p.add_argument("--workers",       type=int,
+                   default=max(1, multiprocessing.cpu_count() - 1))
+    p.add_argument("--dry_run",       action="store_true",
+                   help="Print plan only, write nothing")
+    p.add_argument("--val_split",     type=float, default=0.0,
+                   help="Fraction of augmented images to copy to val/ as well")
+    p.add_argument("--seed",          type=int,   default=42)
+    p.add_argument("--suffix",        default="_aug",
+                   help="Suffix added to augmented image filenames")
+    p.add_argument("--quality",       type=int,   default=92)
+    return p.parse_args()
+
+
+# ─────────────────────────── augmentation pipeline ───────────────────────────
+
+def build_pipeline(img_size: int) -> A.Compose:
+    """
+    Three-tier augmentation pipeline tuned for product crops.
+    Each tier is applied with a probability so every generated
+    image is unique but still realistic.
+    """
+    return A.Compose([
+
+        # ── Tier 1 — always active ─────────────────────────────────────────
+        A.HorizontalFlip(p=0.5),
+
+        A.ShiftScaleRotate(
+            shift_limit=0.06,
+            scale_limit=0.12,
+            rotate_limit=15,
+            border_mode=cv2.BORDER_REFLECT_101,
+            p=0.8,
+        ),
+
+        A.RandomBrightnessContrast(
+            brightness_limit=0.30,
+            contrast_limit=0.30,
+            p=0.85,
+        ),
+
+        A.HueSaturationValue(
+            hue_shift_limit=14,
+            sat_shift_limit=25,
+            val_shift_limit=20,
+            p=0.75,
+        ),
+
+        # ── Tier 2 — medium augmentations (random subset) ─────────────────
+        A.OneOf([
+            A.Perspective(scale=(0.04, 0.10), p=1.0),
+            A.GridDistortion(num_steps=4, distort_limit=0.25, p=1.0),
+            A.ElasticTransform(alpha=60, sigma=8, p=1.0),
+        ], p=0.40),
+
+        A.OneOf([
+            A.GaussianBlur(blur_limit=(3, 5), p=1.0),
+            A.MotionBlur(blur_limit=5, p=1.0),
+            A.Sharpen(alpha=(0.1, 0.4), p=1.0),
+        ], p=0.35),
+
+        A.OneOf([
+            A.GaussNoise(var_limit=(5, 30), p=1.0),
+            A.ISONoise(color_shift=(0.01, 0.05), intensity=(0.05, 0.20), p=1.0),
+        ], p=0.30),
+
+        # Occlusion: simulate a product partially hidden by another
+        A.CoarseDropout(
+            max_holes=4,
+            max_height=int(img_size * 0.20),
+            max_width=int(img_size * 0.20),
+            min_holes=1,
+            fill_value=0,
+            p=0.25,
+        ),
+
+        # JPEG compression artefacts (common with phone cameras)
+        A.ImageCompression(quality_lower=55, quality_upper=90, p=0.25),
+
+        # ── Tier 3 — lighting / environment (low probability) ─────────────
+        A.OneOf([
+            A.RandomShadow(shadow_roi=(0, 0, 1, 1),
+                           num_shadows_lower=1, num_shadows_upper=2,
+                           shadow_dimension=4, p=1.0),
+            A.RandomSunFlare(flare_roi=(0, 0, 1, 0.5),
+                             angle_lower=0.5, num_flare_circles_lower=2,
+                             num_flare_circles_upper=6,
+                             src_radius=100, p=1.0),
+        ], p=0.15),
+
+        A.RandomFog(fog_coef_lower=0.05, fog_coef_upper=0.20,
+                    alpha_coef=0.06, p=0.10),
+
+        # Colour channel shuffle (catches model over-reliance on colour)
+        A.ChannelShuffle(p=0.08),
+
+        # Final mild colour normalisation drift
+        A.RGBShift(r_shift_limit=12, g_shift_limit=12, b_shift_limit=12, p=0.30),
+    ])
+
+
+# ─────────────────────────── worker ──────────────────────────────────────────
+
+def _augment_worker(task: dict) -> dict:
+    """
+    Subprocess worker — generates `n_needed` augmented copies of one class.
+
+    task = {
+        "class_name": str,
+        "src_paths":  [Path, ...],
+        "dst_dir":    Path,
+        "n_needed":   int,
+        "suffix":     str,
+        "quality":    int,
+        "seed":       int,
+        "img_size":   int,
+    }
+    Returns {"class_name": str, "generated": int, "errors": [str]}
+    """
+    rng = random.Random(task["seed"])
+    np.random.seed(task["seed"] % (2**32))
+
+    pipeline   = build_pipeline(task["img_size"])
+    src_paths  = task["src_paths"]
+    dst_dir    = Path(task["dst_dir"])
+    suffix     = task["suffix"]
+    quality    = task["quality"]
+    n_needed   = task["n_needed"]
+    generated  = 0
+    errors     = []
+
+    # Cycle through source images, applying a different random transform each time
+    for i in range(n_needed):
+        src_path = src_paths[i % len(src_paths)]
+        img = cv2.imread(str(src_path))
+        if img is None:
+            errors.append(f"Cannot read {src_path}")
+            continue
+
+        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        try:
+            result  = pipeline(image=img_rgb)
+            aug_rgb = result["image"]
+            aug_bgr = cv2.cvtColor(aug_rgb, cv2.COLOR_RGB2BGR)
+        except Exception as e:
+            errors.append(f"Augmentation failed for {src_path}: e")
+            continue
+
+        # Build a unique output name: original_stem + suffix + counter
+        stem     = Path(src_path).stem
+        out_name = f"{stem}{suffix}{i:05d}.jpg"
+        out_path = dst_dir / out_name
+
+        # Avoid overwriting an existing file
+        if out_path.exists():
+            out_name = f"{stem}{suffix}{i:05d}_{rng.randint(0,9999):04d}.jpg"
+            out_path = dst_dir / out_name
+
+        cv2.imwrite(
+            str(out_path), aug_bgr,
+            [cv2.IMWRITE_JPEG_QUALITY, quality],
+        )
+        generated += 1
+
+    return {
+        "class_name": task["class_name"],
+        "generated":  generated,
+        "errors":     errors,
+    }
+
+
+# ─────────────────────────── helpers ─────────────────────────────────────────
+
+IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"}
+
+
+def get_images(cls_dir: Path) -> list[Path]:
+    return [p for p in cls_dir.iterdir() if p.suffix.lower() in IMAGE_EXTS]
+
+
+def detect_img_size(cls_dir: Path) -> int:
+    """Sample the first image to get its smaller dimension (proxy for crop size)."""
+    for p in cls_dir.iterdir():
+        if p.suffix.lower() in IMAGE_EXTS:
+            img = cv2.imread(str(p))
+            if img is not None:
+                return min(img.shape[:2])
+    return 112   # fallback
+
+
+def print_table(rows: list[dict], target: int):
+    """Print a nicely formatted class statistics table."""
+    col = max(len(r["class"]) for r in rows) + 2
+    header = f"  {'Class':<{col}} {'Current':>9} {'Target':>9} {'To add':>9}  Status"
+    sep    = "─" * len(header)
+    print(sep)
+    print(header)
+    print(sep)
+    for r in rows:
+        need   = max(0, target - r["current"])
+        status = "✓ ok" if need == 0 else f"+ {need:,}"
+        bar    = "█" * min(20, int(20 * r["current"] / target))
+        print(f"  {r['class']:<{col}} {r['current']:>9,} {target:>9,} {need:>9,}  {status}  {bar}")
+    print(sep)
+    total_now = sum(r["current"] for r in rows)
+    total_add = sum(max(0, target - r["current"]) for r in rows)
+    print(f"  {'TOTAL':<{col}} {total_now:>9,}           {total_add:>9,}")
+    print(sep)
+
+
+# ─────────────────────────── main ────────────────────────────────────────────
+
+def main():
+    args = parse_args()
+    random.seed(args.seed)
+
+    split_dir = Path(args.data_dir) / args.split
+    if not split_dir.exists():
+        raise FileNotFoundError(f"Split directory not found: {split_dir}")
+
+    # ── Scan class directories ────────────────────────────────────────────────
+    class_dirs = sorted(
+        d for d in split_dir.iterdir()
+        if d.is_dir() and not d.name.startswith("_")
+    )
+    if not class_dirs:
+        raise RuntimeError(f"No class folders found in {split_dir}")
+
+    counts = {d.name: len(get_images(d)) for d in class_dirs}
+    if not any(counts.values()):
+        raise RuntimeError("All class folders appear empty.")
+
+    max_count = max(counts.values())
+    min_count = min(counts.values())
+
+    # ── Determine target ──────────────────────────────────────────────────────
+    if args.target_count:
+        target = args.target_count
+        print(f"\n[plan] Explicit target: {target:,} images per class")
+    else:
+        target = min(int(max_count * args.max_scale), max_count * 3)
+        print(f"\n[plan] Auto target: {max_count:,} (max) × {args.max_scale} = {target:,} per class")
+
+    # ── Print plan table ──────────────────────────────────────────────────────
+    rows = [{"class": name, "current": cnt} for name, cnt in sorted(counts.items())]
+    print()
+    print_table(rows, target)
+
+    needs_aug = [
+        d for d in class_dirs
+        if counts[d.name] < target * args.min_count
+    ]
+
+    if not needs_aug:
+        print("\n✓ All classes already meet the target — nothing to do.")
+        return
+
+    print(f"\n  Classes to augment: {len(needs_aug)}/{len(class_dirs)}")
+    print(f"  Imbalance ratio  : {max_count/max(min_count,1):.1f}x  →  1.0x after augmentation")
+
+    if args.dry_run:
+        print("\n  [dry-run] No files written. Remove --dry_run to apply.\n")
+        return
+
+    # ── Sample image size (for CoarseDropout sizing) ──────────────────────────
+    img_size = detect_img_size(class_dirs[0])
+
+    # ── Build worker tasks ────────────────────────────────────────────────────
+    tasks = []
+    for cls_dir in needs_aug:
+        name      = cls_dir.name
+        src_imgs  = get_images(cls_dir)
+        n_needed  = target - len(src_imgs)
+        if n_needed <= 0:
+            continue
+        tasks.append({
+            "class_name": name,
+            "src_paths":  [str(p) for p in src_imgs],
+            "dst_dir":    str(cls_dir),
+            "n_needed":   n_needed,
+            "suffix":     args.suffix,
+            "quality":    args.quality,
+            "seed":       args.seed + abs(hash(name)) % 10000,
+            "img_size":   img_size,
+        })
+
+    total_to_generate = sum(t["n_needed"] for t in tasks)
+    print(f"\n[augment] Generating {total_to_generate:,} images across "
+          f"{len(tasks)} classes using {args.workers} worker(s)…\n")
+
+    # ── Run (parallel per class) ──────────────────────────────────────────────
+    results = []
+    if args.workers > 1:
+        with multiprocessing.Pool(processes=args.workers) as pool:
+            for i, res in enumerate(pool.imap_unordered(_augment_worker, tasks), 1):
+                results.append(res)
+                done  = sum(r["generated"] for r in results)
+                pct   = done / total_to_generate * 100
+                print(f"  [{i}/{len(tasks)}] {res['class_name']:<30} "
+                      f"+{res['generated']:,}  "
+                      f"({pct:.0f}% total)", flush=True)
+    else:
+        for i, task in enumerate(tasks, 1):
+            res = _augment_worker(task)
+            results.append(res)
+            done = sum(r["generated"] for r in results)
+            pct  = done / total_to_generate * 100
+            print(f"  [{i}/{len(tasks)}] {res['class_name']:<30} "
+                  f"+{res['generated']:,}  "
+                  f"({pct:.0f}% total)", flush=True)
+
+    # ── Collect errors ────────────────────────────────────────────────────────
+    all_errors = [(r["class_name"], e) for r in results for e in r["errors"]]
+    total_gen  = sum(r["generated"] for r in results)
+
+    # ── Optional val propagation ──────────────────────────────────────────────
+    val_copied = 0
+    if args.val_split > 0:
+        val_dir = Path(args.data_dir) / "val"
+        if val_dir.exists():
+            print(f"\n[val] Copying {args.val_split:.0%} of augmented images to val/…")
+            for res in results:
+                cls_name  = res["class_name"]
+                aug_files = sorted(
+                    p for p in (split_dir / cls_name).iterdir()
+                    if args.suffix in p.stem and p.suffix.lower() in IMAGE_EXTS
+                )
+                n_copy    = max(1, int(len(aug_files) * args.val_split))
+                dst_cls   = val_dir / cls_name
+                dst_cls.mkdir(exist_ok=True)
+                for src in aug_files[:n_copy]:
+                    shutil.copy2(str(src), str(dst_cls / src.name))
+                    val_copied += 1
+        else:
+            print(f"\n  ⚠  val/ not found at {val_dir} — skipping val propagation")
+
+    # ── Final report ──────────────────────────────────────────────────────────
+    print("\n" + "═" * 60)
+    print("  Augmentation complete!")
+    print(f"  Images generated : {total_gen:,}")
+    if val_copied:
+        print(f"  Copied to val/   : {val_copied:,}")
+    if all_errors:
+        print(f"\n  ⚠  {len(all_errors)} error(s):")
+        for cls, err in all_errors[:10]:
+            print(f"     [{cls}] {err}")
+        if len(all_errors) > 10:
+            print(f"     … and {len(all_errors)-10} more")
+
+    # ── Verify final counts ───────────────────────────────────────────────────
+    print("\n  Final class counts:")
+    col = max(len(d.name) for d in class_dirs) + 2
+    print(f"  {'Class':<{col}} {'Before':>8} {'After':>8}")
+    print("  " + "─" * (col + 20))
+    for d in sorted(class_dirs, key=lambda x: x.name):
+        before = counts[d.name]
+        after  = len(get_images(d))
+        delta  = after - before
+        flag   = f"  +{delta:,}" if delta > 0 else ""
+        print(f"  {d.name:<{col}} {before:>8,} {after:>8,}{flag}")
+    print("═" * 60)
+    print(f"\n  ➜  Next step: python 2_train_classifier.py --data_dir {args.data_dir}")
+
+
+if __name__ == "__main__":
+    main()

class_names.json

@@ -2,31 +2,37 @@
   "0": "1001",
   "1": "1001 Gold",
   "2": "A100",
-  "3": "Aroma",
+  "3": "Afia",
-  "4": "Aroma Espresso",
+  "4": "Aroma",
-  "5": "Aroma Familial",
+  "5": "Aroma Espresso",
-  "6": "Aroma Gold",
+  "6": "Aroma Familial",
-  "7": "Bonal",
+  "7": "Aroma Gold",
-  "8": "Bonal Excellence",
+  "8": "Bonal",
-  "9": "Cafe D Or Caps",
+  "9": "Bonal Excellence",
-  "10": "Canastel",
+  "10": "Boukhari",
-  "11": "Caps",
+  "11": "Cafe D Or Caps",
-  "12": "CityOne Instant Coffee",
+  "12": "Canastel",
-  "13": "Dozia",
+  "13": "Caps",
-  "14": "El Kabir",
+  "14": "CityOne Instant Coffee",
-  "15": "Facto",
+  "15": "Dozia",
-  "16": "Facto no",
+  "16": "El Kabir",
-  "17": "Famico",
+  "17": "Elio",
-  "18": "Famico Caps",
+  "18": "Facto",
-  "19": "Famico Exclusive",
+  "19": "Facto no",
-  "20": "Molino",
+  "20": "Famico",
-  "21": "Mundo",
+  "21": "Famico Caps",
-  "22": "Mundo Caps",
+  "22": "Famico Exclusive",
-  "23": "Nizier",
+  "23": "Gosto",
-  "24": "Nouara",
+  "24": "Molino",
-  "25": "Oscar",
+  "25": "Mundo",
-  "26": "Primo",
+  "26": "Mundo Caps",
-  "27": "Ricamar Thon",
+  "27": "Nizier",
-  "28": "Siglo",
+  "28": "Nouara",
-  "29": "Skor"
+  "29": "Oscar",
+  "30": "Poseidon Thon",
+  "31": "Primo",
+  "32": "Proto",
+  "33": "Ricamar Thon",
+  "34": "Siglo",
+  "35": "Skor"
 }