Webcam Data Capture: A Camera Trap for Your Dataset

dataset/
  player/  player_20261105-101500_0001.jpg ...
  ball/    ball_20261105-101500_0001.jpg ...
  capture_log.csv

import csv
import os
import time
from datetime import datetime
 
import cv2
import numpy as np
 
# ---------------- SETTINGS: change these, and justify them in your field brief ----------------
CLASS_NAME = "player"        # player | ball | goal | background
DATASET_ROOT = "dataset"
BLUR_THRESHOLD = 80.0        # lower it outdoors (faster motion, wind, distance)
CAPTURE_INTERVAL = 0.8       # seconds between automatic tries. 0 = SPACE only
MIN_DIFF = 500               # pixels that must differ from the last saved frame
CAMERA_INDEX = 0
SAVE_QUALITY = 92
MOTION_ONLY = False          # True: only try to save when something is moving (lesson 9)
 
TARGET_COUNTS = {"player": 100, "ball": 100, "goal": 50, "background": 50}
# ----------------------------------------------------------------------------------------------
 
CLASSES = list(TARGET_COUNTS.keys())
SMALL = (160, 120)           # frames are compared at this small size
 
 
def blur_score(gray):
    return cv2.Laplacian(gray, cv2.CV_64F).var()
 
 
def pixels_changed(a, b):
    return int(np.count_nonzero(cv2.absdiff(a, b) > 25))
 
 
def new_session_id():
    return datetime.now().strftime("%Y%m%d-%H%M%S")
 
 
def count_existing(cls):
    folder = os.path.join(DATASET_ROOT, cls)
    return len([f for f in os.listdir(folder) if f.endswith(".jpg")]) if os.path.isdir(folder) else 0
 
 
def main():
    cap = cv2.VideoCapture(CAMERA_INDEX)
    if not cap.isOpened():
        raise SystemExit("Cannot open the camera. Try CAMERA_INDEX = 1 or 2.")
 
    os.makedirs(DATASET_ROOT, exist_ok=True)
    for c in CLASSES:
        os.makedirs(os.path.join(DATASET_ROOT, c), exist_ok=True)
 
    log_path = os.path.join(DATASET_ROOT, "capture_log.csv")
    new_log = not os.path.exists(log_path)
    log = open(log_path, "a", newline="")
    writer = csv.writer(log)
    if new_log:
        writer.writerow(["time", "session", "class", "file", "blur", "pixels_changed", "accepted", "reason"])
 
    cls = CLASS_NAME if CLASS_NAME in CLASSES else CLASSES[0]
    counts = {c: count_existing(c) for c in CLASSES}
    session = new_session_id()
    paused = False
    last_try = 0.0
    last_saved_small = None
    recent = []                                     # last 50 attempts: True or False
    message = "ready"
    bg = cv2.createBackgroundSubtractorMOG2(history=200, varThreshold=25, detectShadows=False)
    kernel = np.ones((3, 3), np.uint8)
 
    print("SPACE capture | C class | N new session | P pause | Q quit")
 
    while True:
        ret, frame = cap.read()
        if not ret:
            print("No frame received.")
            break
 
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        score = blur_score(gray)
        small = cv2.resize(gray, SMALL)
 
        moving = True
        if MOTION_ONLY:
            fg = cv2.morphologyEx(bg.apply(frame), cv2.MORPH_OPEN, kernel)
            moving = np.count_nonzero(fg) > 0.01 * fg.size
 
        now = time.time()
        key = cv2.waitKey(1) & 0xFF
        manual = key == ord(" ")
        auto = (CAPTURE_INTERVAL > 0 and not paused and moving and now - last_try >= CAPTURE_INTERVAL)
 
        if key == ord("q"):
            break
        if key == ord("c"):
            cls = CLASSES[(CLASSES.index(cls) + 1) % len(CLASSES)]
            message = f"class: {cls}"
        if key == ord("n"):
            session = new_session_id()
            message = f"new session {session}"
        if key == ord("p"):
            paused = not paused
 
        if manual or auto:
            last_try = now
            changed = pixels_changed(small, last_saved_small) if last_saved_small is not None else 10 ** 9
            if score < BLUR_THRESHOLD:
                accepted, reason = False, "blurry"
            elif changed < MIN_DIFF:
                accepted, reason = False, "too similar"
            else:
                accepted, reason = True, "ok"
 
            name = ""
            if accepted:
                counts[cls] += 1
                name = f"{cls}_{session}_{counts[cls]:04d}.jpg"
                path = os.path.join(DATASET_ROOT, cls, name)
                cv2.imwrite(path, frame, [cv2.IMWRITE_JPEG_QUALITY, SAVE_QUALITY])
                last_saved_small = small
            writer.writerow([datetime.now().isoformat(timespec="seconds"), session, cls, name,
                             round(score, 1), changed if changed < 10 ** 9 else "", accepted, reason])
            recent = (recent + [accepted])[-50:]
            message = f"saved {name}" if accepted else f"rejected: {reason}"
 
        # ---- on-screen display ----
        h, w = frame.shape[:2]
        rate = 100 * sum(recent) / len(recent) if recent else 0
        target = TARGET_COUNTS[cls]
        done = counts[cls] >= target
        cv2.rectangle(frame, (0, 0), (w, 70), (90, 30, 10), -1)
        cv2.putText(frame, f"class: {cls}   {counts[cls]}/{target}" + ("  TARGET REACHED: press C" if done else ""),
                    (10, 22), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0) if not done else (0, 255, 255), 2, cv2.LINE_AA)
        blur_colour = (0, 255, 0) if score >= BLUR_THRESHOLD else (0, 0, 255)
        cv2.putText(frame, f"blur {score:5.0f} (min {BLUR_THRESHOLD:.0f})   accept rate {rate:3.0f}%",
                    (10, 46), cv2.FONT_HERSHEY_SIMPLEX, 0.55, blur_colour, 1, cv2.LINE_AA)
        cv2.putText(frame, message + ("   PAUSED" if paused else ""), (10, 64),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (220, 220, 220), 1, cv2.LINE_AA)
        cv2.imshow("capture.py", frame)
 
    cap.release()
    log.close()
    cv2.destroyAllWindows()
    print("Counts:", counts)
 
 
if __name__ == "__main__":
    main()