craic/tools/camera_capture.py

#!/usr/bin/env python3
"""Capture a JPEG frame from the CRAIC ESP32-S3 camera via HTTP MJPEG stream.

Usage:
    python tools/camera_capture.py                        # auto-detect or default IP
    python tools/camera_capture.py --ip 192.168.x.x       # specify IP
    python tools/camera_capture.py -o frame.jpg           # custom output
    python tools/camera_capture.py --scan                 # scan local subnet for ESP32
"""

from __future__ import annotations

import argparse
import os
import socket
import sys
import time
from datetime import datetime


def find_esp32_on_subnet(subnet_prefix: str) -> list[str]:
    """Ping sweep a /24 subnet and return IPs with port 80 open."""
    candidates: list[str] = []
    print(f"Scanning {subnet_prefix}.0/24 ...")
    for i in range(1, 255):
        ip = f"{subnet_prefix}.{i}"
        r = os.system(f"ping -c 1 -W 1 {ip} >/dev/null 2>&1")
        if r == 0:
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(0.5)
                result = sock.connect_ex((ip, 80))
                sock.close()
                if result == 0:
                    candidates.append(ip)
            except:
                pass
        if i % 50 == 0:
            print(f"  ... scanned {i}/254")
    return candidates


def probe_esp32_status(ip: str, timeout: float = 3) -> bool:
    """Check if an IP serves the ESP32 camera JSON status endpoint."""
    try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(timeout)
        sock.connect((ip, 80))
        sock.sendall(
            b"GET /status HTTP/1.1\r\n"
            b"Host: " + ip.encode() + b"\r\n"
            b"User-Agent: Mozilla/5.0\r\n"
            b"Accept: */*\r\n"
            b"Connection: close\r\n\r\n"
        )
        data = b""
        while True:
            try:
                chunk = sock.recv(4096)
                if not chunk:
                    break
                data += chunk
            except socket.timeout:
                break
            except:
                break
        sock.close()
        # ESP32 /status returns JSON with these keys
        return b"capture_fps" in data or b"has_client" in data
    except:
        return False


def grab_jpeg_frame(ip: str, port: int = 80, timeout: float = 10) -> bytes | None:
    """Connect to MJPEG stream, read raw bytes, extract first valid JPEG frame."""
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(timeout)
    try:
        sock.connect((ip, port))
        sock.sendall(
            f"GET /stream HTTP/1.1\r\n"
            f"Host: {ip}\r\n"
            f"User-Agent: Mozilla/5.0\r\n"
            f"Accept: */*\r\n"
            f"Connection: close\r\n\r\n".encode()
        )

        buf = b""
        start = time.time()
        boundary = b"--frame"

        while time.time() - start < timeout:
            try:
                chunk = sock.recv(4096)
                if not chunk:
                    break
                buf += chunk
            except socket.timeout:
                continue
            except:
                break

            # Scan buffer for a complete JPEG frame
            # Format: --frame\r\n...headers...\r\n\r\n<JPEG>\r\n--frame
            while True:
                # Find JPEG start-of-image marker
                soi = buf.find(b"\xff\xd8")
                if soi == -1:
                    break
                # Find JPEG end-of-image marker
                eoi = buf.find(b"\xff\xd9", soi + 2)
                if eoi == -1:
                    break

                jpeg = buf[soi : eoi + 2]

                # Verify it's realistically sized (>= 1KB) and preceded by boundary
                if len(jpeg) >= 1024:
                    return jpeg

                # False positive — keep scanning past this SOI
                buf = buf[soi + 2:]

        return None

    except socket.timeout:
        print(f"  Timeout connecting to {ip}:{port}", file=sys.stderr)
        return None
    except ConnectionRefusedError:
        return None
    except Exception as e:
        print(f"  Socket error: {e}", file=sys.stderr)
        return None
    finally:
        try:
            sock.close()
        except:
            pass


def resolve_camera() -> str | None:
    """Try to find the ESP32 camera on local networks."""
    # Get local IP
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
    except:
        local_ip = "127.0.0.1"

    prefix = ".".join(local_ip.split(".")[:3])

    # Collect candidate IPs to try
    candidates: list[str] = []

    # 1. Gateway (common for ESP32 AP mode)
    candidates.append(f"{prefix}.1")
    # 2. Common DHCP range
    for i in range(100, 111):
        candidates.append(f"{prefix}.{i}")
    # 3. Some other common patterns
    for i in [50, 51, 20, 30, 200, 201]:
        candidates.append(f"{prefix}.{i}")

    # Deduplicate
    seen: set[str] = set()
    unique: list[str] = []
    for ip in candidates:
        if ip not in seen:
            seen.add(ip)
            unique.append(ip)

    print(f"Probing {len(unique)} likely IPs on {prefix}.0/24 ...")
    for ip in unique:
        print(f"  {ip} ... ", end="", flush=True)
        if probe_esp32_status(ip):
            print("ESP32 camera found!")
            return ip
        print("no")

    # Broader subnet scan
    print(f"\nNo camera at common IPs. Scanning full subnet {prefix}.0/24 ...")
    port80_hosts = find_esp32_on_subnet(prefix)
    for ip in port80_hosts:
        if probe_esp32_status(ip):
            print(f"ESP32 camera found at {ip}")
            return ip

    return None


def main() -> int:
    parser = argparse.ArgumentParser(
        description="从 CRAIC ESP32-S3 摄像头抓取 JPEG 帧"
    )
    parser.add_argument("--ip", help="ESP32 IP 地址 (默认自动检测)")
    parser.add_argument("--port", type=int, default=80)
    parser.add_argument("-o", "--output", help="输出文件路径")
    parser.add_argument("--scan", action="store_true", help="扫描子网查找 ESP32")
    args = parser.parse_args()

    # ---- resolve IP ----
    ip = args.ip
    if args.scan:
        ip = resolve_camera()
    elif ip is None:
        ip = resolve_camera()

    if ip is None:
        print(
            "ERROR: 找不到 ESP32 摄像头，请确保:\n"
            "  1. ESP32 已上电\n"
            "  2. ESP32 已连接 WiFi (STA 模式, SSID: FS)\n"
            "  3. 本机与 ESP32 在同一子网\n"
            "  或使用 --ip 直接指定地址",
            file=sys.stderr,
        )
        return 1

    # ---- capture ----
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    out = args.output or f"capture_{ts}.jpg"

    print(f"Connecting to http://{ip}:{args.port}/stream ...")
    jpeg = grab_jpeg_frame(ip, args.port)

    if jpeg is None:
        print("ERROR: 无法获取 JPEG 帧。", file=sys.stderr)
        return 1

    with open(out, "wb") as f:
        f.write(jpeg)
    print(f"Saved {out}  ({len(jpeg)} bytes)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())