slr_handshapes/webcam_capture.py

#!/usr/bin/env python3
"""
capture_webcam.py
Show webcam preview and, given --letter L, count down 5s, then capture frames
every --interval seconds until --count images are saved.
Saves PNGs to ./captures as L001.PNG, L002.PNG, ...

Usage:
  python capture_webcam.py --letter A
  python capture_webcam.py --letter B --camera 1
  python capture_webcam.py --letter C --count 10 --interval 1

   # Default: 5 captures at 2s spacing, 640x480
   python capture_webcam.py --letter A

   # Ten captures, 1s apart
   python capture_webcam.py --letter B --count 10 --interval 1

   # USB camera index 1, HD override
   python capture_webcam.py --letter C --camera 1 --width 1280 --height 720 --count 8 --interval 1.5
"""

import argparse
import os
import re
import time
from pathlib import Path

import cv2

COUNTDOWN_SECONDS = 5

def next_sequence_number(captures_dir: Path, letter: str) -> int:
    """Return next available sequence number for files like 'A001.PNG'."""
    pattern = re.compile(rf"^{re.escape(letter)}(\d{{3}})\.PNG$", re.IGNORECASE)
    max_idx = 0
    if captures_dir.exists():
        for name in os.listdir(captures_dir):
            m = pattern.match(name)
            if m:
                try:
                    idx = int(m.group(1))
                    if idx > max_idx:
                        max_idx = idx
                except ValueError:
                    pass
    return max_idx + 1

def draw_text(img, text, org, scale=1.4, color=(0, 255, 0), thickness=2):
    cv2.putText(img, text, org, cv2.FONT_HERSHEY_SIMPLEX, scale, color, thickness, cv2.LINE_AA)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--letter", required=True, help="Target letter A–Z. Output files like A001.PNG")
    ap.add_argument("--camera", type=int, default=0, help="OpenCV camera index (default: 0)")
    ap.add_argument("--width", type=int, default=640, help="Requested capture width (default: 640)")
    ap.add_argument("--height", type=int, default=480, help="Requested capture height (default: 480)")
    ap.add_argument("--count", type=int, default=5, help="Number of captures to take (default: 5)")
    ap.add_argument("--interval", type=float, default=2.0, help="Seconds between captures (default: 2.0)")
    args = ap.parse_args()

    letter = args.letter.upper().strip()
    if not (len(letter) == 1 and "A" <= letter <= "Z"):
        raise SystemExit("Please pass a single letter A–Z to --letter (e.g., --letter A)")
    if args.count <= 0:
        raise SystemExit("--count must be >= 1")
    if args.interval <= 0:
        raise SystemExit("--interval must be > 0")

    captures_dir = Path("./captures")
    captures_dir.mkdir(parents=True, exist_ok=True)
    start_idx = next_sequence_number(captures_dir, letter)

    cap = cv2.VideoCapture(args.camera)
    if not cap.isOpened():
        raise SystemExit(f"❌ Could not open camera index {args.camera}")

    # Try to set resolution (best-effort)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, args.width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, args.height)

    window_title = f"Capture {letter} (press 'q' to quit)"
    print(f"Showing webcam. Countdown {COUNTDOWN_SECONDS}s, then capturing {args.count} frame(s) every {args.interval}s...")
    print(f"Saving to: {captures_dir.resolve()} as {letter}NNN.PNG starting at index {start_idx:03d}")

    countdown_done_at = time.time() + COUNTDOWN_SECONDS
    # Absolute times when we want to capture (after countdown)
    capture_times = [countdown_done_at + i * args.interval for i in range(args.count)]
    capture_taken = [False] * args.count
    captures_made = 0
    idx = start_idx

    while True:
        ok, frame = cap.read()
        if not ok:
            print("⚠️  Frame grab failed; ending.")
            break

        now = time.time()

        # Countdown overlay
        if now < countdown_done_at:
            remaining = int(round(countdown_done_at - now))
            overlay = frame.copy()
            draw_text(overlay, f"Starting in: {remaining}s", (30, 60), scale=2.0, color=(0, 255, 255), thickness=3)
            draw_text(overlay, f"Letter: {letter}", (30, 120), scale=1.2, color=(0, 255, 0), thickness=2)
            cv2.imshow(window_title, overlay)
        else:
            # Check if it's time for any pending captures
            for i, tcap in enumerate(capture_times):
                if (not capture_taken[i]) and now >= tcap:
                    filename = f"{letter}{idx:03d}.PNG"
                    out_path = captures_dir / filename
                    cv2.imwrite(str(out_path), frame)
                    capture_taken[i] = True
                    captures_made += 1
                    idx += 1
                    print(f"📸 Saved {out_path.name}")

            # Overlay progress
            elapsed_after = now - countdown_done_at
            total_duration = args.interval * (args.count - 1) if args.count > 1 else 0
            remaining_after = max(0.0, total_duration - elapsed_after)
            overlay = frame.copy()
            draw_text(overlay, f"Capturing {letter}… {captures_made}/{args.count}", (30, 60),
                      scale=1.5, color=(0, 255, 0), thickness=3)
            draw_text(overlay, f"Time left: {int(round(remaining_after))}s", (30, 110),
                      scale=1.2, color=(0, 255, 255), thickness=2)
            cv2.imshow(window_title, overlay)

            # If finished all captures, keep preview up until user quits
            if captures_made >= args.count:
                draw_text(overlay, "Done! Press 'q' to close.", (30, 160),
                          scale=1.2, color=(0, 200, 255), thickness=2)
                cv2.imshow(window_title, overlay)

        # Quit on 'q'
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()