Train notebook

1_prepare_dataset.py

@@ -1,5 +1,5 @@
 """
-STEP 1 ─ prepare_dataset.py
+STEP 1 ─ 1_prepare_dataset.py
 ===========================
 Run this FIRST. It creates the required folder structure and
 a dataset.yaml that YOLOv8 expects.

2_augment.py

@@ -1,5 +1,5 @@
 """
-STEP 2 — augment.py
+STEP 2 — 2_augment.py
 ====================
 Expands your small labeled dataset using augmentations tuned
 specifically for retail shelf / packaged-goods detection.

4_detect.py

@@ -0,0 +1,224 @@
+"""
+SETP 4 - detect.py
+====================
+Run detection on a single image, a folder, or a webcam stream.
+Draw bounding boxes, prints a per-class count summary, and
+saves an annotated result image.
+
+Usage:
+    python detect.py --source scene.jpg
+    python detect.py --source images/
+    python detect.py --source 0
+    python detect.py --source scene.jpg --conf 0.5 --wieghts best.pt
+"""
+
+import argparse
+import sys
+import os
+import cv2
+import numpy as np
+
+from pathlib     import Path
+from collections import defaultdict
+from ultralytics import YOLO  
+
+# ── Defaults ──────────────────────────────────────────────────────────────────
+DEFAULT_WEIGHTS = "runs/detect/coffee_v1/weights/best.pt"
+DEFAULT_CONF    = .35   # lower = more detections and more false positives.
+DEFAULT_IOU     = .45   # NMS  IoU threshold.
+DEFAULT_IMGSZ   = 640
+OUTPUT_DIR      = Path("detections")
+
+
+# ── Colour palette (one BGR colour per class index) ──────────────────────────
+PALETTE = [
+    (  0, 200,   0),   # green
+    (  0, 120, 255),   # orange
+    (255,  50,  50),   # blue
+    (200,   0, 200),   # magenta
+    (  0, 200, 200),   # yellow
+    ( 80, 200,  80),
+    (200, 100,   0),
+    (  0,  80, 200),
+    (150,   0, 150),
+    (100, 200,   0),
+    (  0, 150, 150),
+    (200,  50, 100),
+    (  50,  50, 200),
+]
+
+def get_color(class_id: int):
+    return PALETTE[class_id % len(PALETTE)]
+
+# ── Drawing ──────────────────────────
+def draw_box(img, box, class_name, conf, class_id):
+    x1, y1, x2, y2 = map(int, box)
+    color          = get_color(class_id)
+
+    cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+    
+    label = f"{class_name} {conf: .0%}"
+    (tw, th), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 1)
+    cv2.rectangle(
+      img,
+      (x1, y1 - th - baseline - 6),
+      (x1 + tw + 4, y1),
+      color, cv2.FILLED
+    )
+
+    cv2.putText(
+      img,
+      label,
+      (x1 + 2, y1 - baseline -2),
+      cv2.FONT_HERSHEY_SIMPLEX,
+      0.55,
+      (0, 0, 0), 1, cv2.LINE_AA
+    )
+
+def draw_summary(img, counts: dict):
+    """ Overlay a product count table in the top-right corner."""
+    if not counts:
+        return
+
+    lines = ["── Count ──"] + [f"{n:>2}x  {name}" for name, n in sorted(counts.items())]
+
+    x = img.shape[1] - 200
+    y = 16
+    for line in lines:
+        (tw, th), _ = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        cv2.rectangle(img, (x - 4, y - th - 2), (x + tw + 4, y + 4), (30, 30, 30), cv2.FILLED)
+        cv2.putText(img, line, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (240, 240, 240), 1, cv2.LINE_AA)
+        y += th + 8
+
+
+# ── Single image inference ────────────────────────────────────────────────────
+def detect_image(model, img_path: Path, conf: float, iou: float, imgsz: int):
+    img = cv2.imread(str(img_path))
+    if img is None:
+        print(f"[WARN] Cannot read {img_path}")
+        return
+    results = model.predict(
+        source=str(img_path),
+        conf=conf,
+        iou=iou,
+        imgsz=imgsz,
+        verbose=False,
+    )[0]
+
+    counts = defaultdict(int)
+
+    for box in results.boxes:
+        cid = int(box.cls)
+        name = model.names[cid]
+        score = float(box.conf)
+        draw_box(img, box.xyxy[0].tolist(), name, score, cid)
+        counts[name] += 1
+
+    draw_summary(img, counts)
+
+    OUTPUT_DIR.mkdir(exist_ok=True)
+    out_path = OUTPUT_DIR / img_path.name
+    cv2.imwrite(str(out_path), img)
+
+    # Console Summary
+    print(f"\n{img_path.name}")
+    if counts:
+        for name, n in sorted(counts.items()):
+            print(f" {n:>2}x {name}")
+    else:
+        print(" (no detections)")
+    print(f" -> saved: {out_path}")
+    
+    return img, counts
+
+# ── Webcam / video stream ─────────────────────────────────────────────────────
+def detect_stream(model, source, conf: float, iou: float, imgsz: int):
+    cap = cv2.VideoCapture(int(source) if source.isdigit() else source)
+    if not cap.isOpened():
+        print(f"[ERROR] Cannot open source: {source}")
+        return
+
+    print("Streaming - press Q to quit")
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+
+        results = model.predict(
+            source  = frame,
+            conf    = conf,
+            iou     = iou,
+            imgsz   = imgsz,
+            verbose = False
+        )[0]
+
+        counts = defaultdict(int)
+        for box in results.boxes:
+            cid = int(box.cls)
+            name = model.names[cid]
+            draw_box(frame, box.xyxy[0].tolist(), name, float(box.conf), cid)
+            counts[name] += 1
+
+        draw_summary(frame, counts)
+        cv2.imshow("Coffee Detector", frame)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            break
+
+    cap.release()
+    cv2.destroyAllWindows()
+
+
+# ── Entry point ─────────────────────────────────────────────────────
+def main():
+    parser = argparse.ArgumentParser(description="Coffee Product Detector")
+    parser.add_argument("--source" , default="scene.jpg"    ,             help="Image path, folder, or 0 for webcam")
+    parser.add_argument("--weights", default=DEFAULT_WEIGHTS,             help="Path to *.pt")
+    parser.add_argument("--conf"   , default=DEFAULT_CONF   , type=float, help="Confidence threshold (0-1)")
+    parser.add_argument("--iou"    , default=DEFAULT_IOU    , type=float, help="NMS IoU threshold (0-1)")
+    parser.add_argument("--imgsz"  , default=DEFAULT_IMGSZ  , type=int  , help="Inference image size")
+
+    args = parser.parse_args()
+
+    if not Path(args.weights).exists():
+        print(f"[ERROR] Weights not found: {args.weights}")
+        print("         Run train.py first.")
+        sys.exit(1)
+    
+    print(f"Loading model: {args.weights}")
+    model = YOLO(args.weights)
+    print(f"Classes: {list(model.names.values())}\n")
+
+    source = Path(args.source)
+
+    # Webcam / video
+    if args.source.isdigit() or str(source).endswith((".mp4", ".avi", ".mov")):
+        detect_stream(model, args.source, args.conf, args.iou, args.imgsz)
+        return
+
+    # Folder
+    if source.is_dir():
+        exts   = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
+        images = sorted([p for p in source.iterdir() if p.suffix.lower() in exts])
+        print(f"Processing {len(images)} images from {source}...")
+        total_counts = defaultdict(int)
+        for img_path in images:
+            _, counts = detect_image(model, img_path, args.conf, args.iou, args.imgsz)
+            if counts:
+                for k, v in counts.items():
+                    total_counts[k] += v
+
+        print("\n── Total across all images ──")
+        for name, n in sorted(total_counts.items()):
+            print(f"   {n:>3}x {name}")
+        return
+
+    # Single image
+    if source.exists():
+        detect_image(model, source, args.conf, args.iou, args.imgsz)
+        return
+
+    print(f"[ERROR] Source not found: {source}")
+    sys.exit(1)
+
+if __name__ == "__main__":
+    main()