trigger_capture.py

#!/usr/bin/env python3
"""
Triggered capture: monitors the ESP32's serial output and takes a photo
with the UK1275 camera each time object detection causes the LEDs to turn on.

Usage:
    python trigger_capture.py [--port /dev/ttyUSB0] [--baud 115200]

The camera is initialized once at startup (~6 s) and stays armed so that
grab_one_frame() fires immediately on each trigger.

If consecutive captures fail (empty frame), re-running exposure_setup() is
attempted automatically before the next trigger.
"""

import argparse
import sys
import time
import threading
from pathlib import Path

import usb.core
import serial
from PIL import Image

from uk1275_camera import (
    open_device, init_camera, arm_capture, exposure_setup,
    grab_one_frame, decode_frame, get_device_address,
    FRAME_WIDTH, FRAME_HEIGHT,
    CMD_GET_STATUS_A2, EP_CMD_OUT, EP_STATUS_IN,
)

from read_barcodes import read_plate, reorient_plate, write_csv

CAPTURE_DIR = Path(__file__).parent / "captures"
CSV_DIR     = Path(__file__).parent / "csv"

# The plate's A1 notch exposes whichever sensor is at the A1 corner,
# so the triggering sensor always identifies where A1 is in the image.
# Set each value to the corner that sensor occupies in the camera frame.
# Options: "top-left", "bottom-right", "top-right", "bottom-left"
SENSOR_ORIENTATION = {
    "GPIO34": "bottom-left",
    "GPIO27": "top-right",
}
COOLDOWN_S  = 2.0   # minimum seconds between captures to avoid re-trigger noise
KEEPALIVE_S = 1.0   # how often to poll the camera while idle


def save_image(raw: bytes, stem: str) -> Path:
    CAPTURE_DIR.mkdir(exist_ok=True)
    arr = decode_frame(raw)
    img = Image.fromarray(arr, mode="L")
    out = CAPTURE_DIR / f"{stem}_{FRAME_WIDTH}x{FRAME_HEIGHT}_grey8.png"
    img.save(out)
    return out


def keepalive(dev):
    """Send a status poll to stop the camera's internal watchdog from firing."""
    try:
        dev.write(EP_CMD_OUT, CMD_GET_STATUS_A2)
        dev.read(EP_STATUS_IN, 91, timeout=500)
    except usb.core.USBError:
        pass


def init_camera_full(verbose=True):
    """Open + full init sequence. Returns the device handle."""
    dev = open_device()
    if verbose:
        print(f"  USB address {get_device_address(dev)}")
    init_camera(dev, verbose=False)
    arm_capture(dev, verbose=False)
    exposure_setup(dev, verbose=verbose)
    return dev


def recover_camera():
    """Re-open and fully re-init the camera after a USB error. Retries for up to 30 s."""
    print("  Camera lost — attempting recovery…")
    for attempt in range(6):
        # Short initial wait for USB re-enumeration; longer back-off on repeated failures.
        time.sleep(0.5 if attempt == 0 else 2)
        try:
            dev = init_camera_full(verbose=False)
            print("  Camera recovered.")
            return dev
        except Exception as e:
            print(f"  Recovery attempt {attempt + 1}/6 failed: {e}")
    raise RuntimeError("Could not recover camera after 6 attempts.")


def parse_args():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--port", default="/dev/ttyUSB0",
                   help="Serial port of the ESP32 (default: /dev/ttyUSB0)")
    p.add_argument("--baud", type=int, default=115200,
                   help="Baud rate (default: 115200)")
    return p.parse_args()


def main():
    args = parse_args()

    # ── Init camera once at startup ────────────────────────────────────────
    print("Opening camera…")
    try:
        dev = init_camera_full(verbose=True)
    except RuntimeError as e:
        print(f"ERROR: {e}")
        sys.exit(1)
    print("Camera ready.\n")

    # ── Open serial port ───────────────────────────────────────────────────
    print(f"Monitoring {args.port} at {args.baud} baud…")
    try:
        ser = serial.Serial(args.port, args.baud, timeout=0.4)
    except serial.SerialException as e:
        print(f"ERROR: cannot open serial port: {e}")
        sys.exit(1)

    last_capture   = 0.0
    last_keepalive = time.time()
    empty_frames   = 0

    print("Waiting for object detection (Ctrl-C to quit)…\n")
    try:
        while True:
            line = ser.readline().decode("ascii", errors="replace").strip()

            now = time.time()
            if (now - last_keepalive) >= KEEPALIVE_S:
                keepalive(dev)
                last_keepalive = now

            if not line:
                continue
            print(f"  ESP32: {line}")

            if "LEDs ON" not in line:
                continue

            # Parse which sensor fired; prefer GPIO34 when both trigger together.
            if "GPIO34" in line:
                triggered_sensor = "GPIO34"
            else:
                triggered_sensor = "GPIO27"

            now = time.time()
            if (now - last_capture) < COOLDOWN_S:
                print("  (cooldown — skipping capture)")
                continue

            last_capture = now
            last_keepalive = now  # reset so keepalive doesn't fire mid-capture
            print("  Triggering capture…")
            for attempt in range(2):
                try:
                    t0  = time.time()
                    raw = grab_one_frame(dev)
                    elapsed = time.time() - t0
                    break
                except usb.core.USBError as e:
                    print(f"  USB error during capture: {e}")
                    if attempt == 0:
                        try:
                            dev = recover_camera()
                        except RuntimeError as re:
                            print(f"  FATAL: {re}")
                            ser.close()
                            sys.exit(1)
                    else:
                        print("  Capture failed after recovery — waiting for next trigger.")
                        raw = b""
                        elapsed = 0

            if not raw:
                empty_frames += 1
                print(f"  WARNING: empty frame ({empty_frames} in a row)")
                if empty_frames >= 2:
                    try:
                        dev = recover_camera()
                    except RuntimeError as re:
                        print(f"  FATAL: {re}")
                        break
                    empty_frames = 0
                continue

            empty_frames = 0
            expected_bytes = FRAME_WIDTH * FRAME_HEIGHT
            print(f"  Got {len(raw):,} bytes in {elapsed:.2f} s"
                  + ("" if len(raw) == expected_bytes
                     else f"  ** WARNING: expected {expected_bytes:,} bytes **"))

            stem = f"uk1275_{time.strftime('%Y%m%d_%H%M%S')}"
            out  = save_image(raw, stem)
            print(f"  Saved → {out}")
            orientation = SENSOR_ORIENTATION.get(triggered_sensor, "top-left")
            CSV_DIR.mkdir(exist_ok=True)
            csv_path = CSV_DIR / f"{stem}.csv"

            def _scan(image_path, ori, out_path):
                print(f"  Reading barcodes (sensor {triggered_sensor}, A1 at {ori})…", flush=True)
                results = read_plate(image_path)
                results = reorient_plate(results, ori)
                write_csv(results, out_path)
                print()

            threading.Thread(target=_scan, args=(out, orientation, csv_path), daemon=True).start()

    except KeyboardInterrupt:
        print("\nInterrupted.")
    finally:
        ser.close()


if __name__ == "__main__":
    main()