225 lines
7.3 KiB
Python
225 lines
7.3 KiB
Python
"""
|
|
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()
|