169 lines
5.8 KiB
Python
169 lines
5.8 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
import heapq
|
|
import threading
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9_.-]+")
|
|
|
|
|
|
class CameraSnapshots:
|
|
"""Small optional OpenCV webcam wrapper for bench-test snapshots."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
camera_index: int,
|
|
output_dir: Path,
|
|
warmup_seconds: float = 0.5,
|
|
width: int | None = None,
|
|
height: int | None = None,
|
|
) -> None:
|
|
self.camera_index = camera_index
|
|
self.output_dir = output_dir
|
|
self.warmup_seconds = max(0.0, warmup_seconds)
|
|
self.width = width
|
|
self.height = height
|
|
self.cv2: Any | None = None
|
|
self.device: Any | None = None
|
|
self._condition = threading.Condition()
|
|
self._tasks: list[tuple[float, int, dict[str, Any]]] = []
|
|
self._records: list[dict[str, Any]] = []
|
|
self._thread: threading.Thread | None = None
|
|
self._seq = 0
|
|
self._closing = False
|
|
|
|
def open(self) -> None:
|
|
try:
|
|
import cv2
|
|
except ImportError as exc: # pragma: no cover - depends on local bench environment.
|
|
raise SystemExit(
|
|
"OpenCV is required for webcam snapshots. Install it with: "
|
|
".\\.venv\\Scripts\\python.exe -m pip install opencv-python"
|
|
) from exc
|
|
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
self.cv2 = cv2
|
|
self.device = cv2.VideoCapture(self.camera_index)
|
|
if not self.device.isOpened():
|
|
raise SystemExit(f"could not open webcam index {self.camera_index}")
|
|
if self.width:
|
|
self.device.set(cv2.CAP_PROP_FRAME_WIDTH, self.width)
|
|
if self.height:
|
|
self.device.set(cv2.CAP_PROP_FRAME_HEIGHT, self.height)
|
|
self._warm_up()
|
|
self._thread = threading.Thread(target=self._run, name="camera-snapshots", daemon=True)
|
|
self._thread.start()
|
|
|
|
def close(self) -> None:
|
|
with self._condition:
|
|
self._closing = True
|
|
self._condition.notify_all()
|
|
if self._thread is not None:
|
|
self._thread.join()
|
|
self._thread = None
|
|
if self.device is not None:
|
|
self.device.release()
|
|
self.device = None
|
|
|
|
def schedule(
|
|
self,
|
|
*,
|
|
label: str,
|
|
frame_text: str,
|
|
phase: str,
|
|
delay_seconds: float = 0.0,
|
|
step_index: int | None = None,
|
|
) -> dict[str, Any]:
|
|
if self.cv2 is None or self.device is None:
|
|
raise RuntimeError("camera snapshotter is not open")
|
|
|
|
due_monotonic = time.monotonic() + max(0.0, delay_seconds)
|
|
due_timestamp = datetime.now() + timedelta(seconds=max(0.0, delay_seconds))
|
|
timestamp = due_timestamp.strftime("%Y%m%d-%H%M%S-%f")[:-3]
|
|
step_part = f"step{step_index:03d}_" if step_index is not None else ""
|
|
frame_part = frame_text.replace(" ", "")
|
|
filename = _safe_name(f"{timestamp}_{step_part}{phase}_{label}_{frame_part}.jpg")
|
|
path = self.output_dir / filename
|
|
record: dict[str, Any] = {
|
|
"path": str(path),
|
|
"scheduled_timestamp": timestamp,
|
|
"camera_index": self.camera_index,
|
|
"step_index": step_index,
|
|
"phase": phase,
|
|
"delay_seconds": delay_seconds,
|
|
"label": label,
|
|
"frame": frame_text,
|
|
"status": "scheduled",
|
|
}
|
|
with self._condition:
|
|
self._seq += 1
|
|
heapq.heappush(self._tasks, (due_monotonic, self._seq, record))
|
|
self._records.append(record)
|
|
self._condition.notify_all()
|
|
return record
|
|
|
|
def records(self) -> list[dict[str, Any]]:
|
|
with self._condition:
|
|
return [dict(record) for record in self._records]
|
|
|
|
def _run(self) -> None:
|
|
while True:
|
|
with self._condition:
|
|
while not self._tasks and not self._closing:
|
|
self._condition.wait()
|
|
if not self._tasks and self._closing:
|
|
return
|
|
due_monotonic, _seq, record = self._tasks[0]
|
|
wait_seconds = due_monotonic - time.monotonic()
|
|
if wait_seconds > 0:
|
|
self._condition.wait(wait_seconds)
|
|
continue
|
|
heapq.heappop(self._tasks)
|
|
self._write_snapshot(record)
|
|
|
|
def _write_snapshot(self, record: dict[str, Any]) -> None:
|
|
if self.cv2 is None or self.device is None:
|
|
record["status"] = "error"
|
|
record["error"] = "camera closed before capture"
|
|
return
|
|
|
|
image = None
|
|
ok = False
|
|
for _attempt in range(3):
|
|
ok, image = self.device.read()
|
|
if ok:
|
|
break
|
|
time.sleep(0.020)
|
|
if not ok or image is None:
|
|
record["status"] = "error"
|
|
record["error"] = f"webcam index {self.camera_index} did not return an image"
|
|
return
|
|
|
|
path = Path(str(record["path"]))
|
|
if not self.cv2.imwrite(str(path), image):
|
|
record["status"] = "error"
|
|
record["error"] = f"failed to write webcam snapshot {path}"
|
|
return
|
|
record["status"] = "written"
|
|
record["captured_timestamp"] = datetime.now().strftime("%Y%m%d-%H%M%S-%f")[:-3]
|
|
|
|
def _warm_up(self) -> None:
|
|
if self.device is None:
|
|
return
|
|
deadline = time.monotonic() + self.warmup_seconds
|
|
while time.monotonic() < deadline:
|
|
self.device.read()
|
|
time.sleep(0.020)
|
|
self.device.read()
|
|
|
|
|
|
def _safe_name(text: str) -> str:
|
|
cleaned = _SAFE_NAME_RE.sub("_", text.strip()).strip("._")
|
|
return cleaned or "snapshot.jpg"
|