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