first commit
This commit is contained in:
ykcheungab
@@ -0,0 +1,3 @@
|
||||
"""Sensor Bench: operator tooling for raw sensor inspection and recording."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Orisys SDK-owned fixed-format acquisition for Sensor Bench.
|
||||
|
||||
Camera open/read/process is owned by ``orisys.Sensor`` using the DDJX01
|
||||
config (1920x1080 @ 30 FPS MJPG by default). Sensor Bench only displays
|
||||
``sensor.frame`` (raw) and ``get_img()`` (processed) plus force metrics.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FixedCameraMode:
|
||||
"""Stream mode taken from the Orisys config JSON."""
|
||||
|
||||
width: int
|
||||
height: int
|
||||
fps: float
|
||||
config_ref: str
|
||||
stitch_width: int = 0
|
||||
stitch_height: int = 0
|
||||
|
||||
def label(self) -> str:
|
||||
return f"{self.width}x{self.height} @{self.fps:g} (DDJX01 / MJPG)"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SdkPacket:
|
||||
"""One acquisition tick from the Orisys Sensor."""
|
||||
|
||||
raw_frame: np.ndarray
|
||||
processed_frame: np.ndarray
|
||||
timestamp: float
|
||||
fps: float
|
||||
process_ms: float
|
||||
width: int
|
||||
height: int
|
||||
processed_width: int
|
||||
processed_height: int
|
||||
normal_force: float = 0.0
|
||||
shear_x: float = 0.0
|
||||
shear_y: float = 0.0
|
||||
session_id: int = 0
|
||||
mode: Optional[FixedCameraMode] = None
|
||||
|
||||
|
||||
def orisys_importable() -> tuple[bool, str]:
|
||||
try:
|
||||
import orisys # noqa: F401
|
||||
|
||||
return True, getattr(__import__("orisys"), "__version__", "unknown")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
return False, str(exc)
|
||||
|
||||
|
||||
def load_fixed_mode(config_ref: str | Path) -> FixedCameraMode:
|
||||
ref = str(config_ref).strip()
|
||||
path = Path(ref)
|
||||
if path.exists():
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
resolved_ref = str(path.resolve())
|
||||
else:
|
||||
try:
|
||||
import orisys
|
||||
|
||||
data = orisys.configs.load_config(ref)
|
||||
resolved_ref = ref
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise FileNotFoundError(
|
||||
f"Orisys config not found as a file or bundled config name: {ref}"
|
||||
) from exc
|
||||
camera = data.get("camera") or {}
|
||||
raw = camera.get("resolution_raw") or [1920, 1080]
|
||||
stitch = camera.get("resolution_stitch") or [0, 0]
|
||||
fps = float(camera.get("fps") or 30.0)
|
||||
if not isinstance(raw, (list, tuple)) or len(raw) < 2:
|
||||
raise ValueError(f"Invalid camera.resolution_raw in {ref}")
|
||||
return FixedCameraMode(
|
||||
width=int(raw[0]),
|
||||
height=int(raw[1]),
|
||||
fps=fps,
|
||||
config_ref=resolved_ref,
|
||||
stitch_width=int(stitch[0]) if len(stitch) >= 2 else 0,
|
||||
stitch_height=int(stitch[1]) if len(stitch) >= 2 else 0,
|
||||
)
|
||||
|
||||
|
||||
def _ensure_bgr(img: np.ndarray) -> np.ndarray:
|
||||
if img.ndim == 2:
|
||||
return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
|
||||
if img.ndim == 3 and img.shape[2] == 1:
|
||||
return cv2.cvtColor(img[:, :, 0], cv2.COLOR_GRAY2BGR)
|
||||
if img.ndim == 3 and img.shape[2] == 4:
|
||||
return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
|
||||
return img
|
||||
|
||||
|
||||
class SdkAcquisitionWorker:
|
||||
"""Background loop: ``Sensor.get_img()`` + deformation + contact."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
on_error: Optional[Callable[[str], None]] = None,
|
||||
on_status: Optional[Callable[[str], None]] = None,
|
||||
):
|
||||
self._on_error = on_error
|
||||
self._on_status = on_status
|
||||
self._thread: Optional[threading.Thread] = None
|
||||
self._stop = threading.Event()
|
||||
self._lock = threading.Lock()
|
||||
self._latest: Optional[SdkPacket] = None
|
||||
self._session_id = 0
|
||||
self._sensor = None
|
||||
self._mode: Optional[FixedCameraMode] = None
|
||||
self._video_source: Optional[str | int] = None
|
||||
self._cal_path: Optional[str] = None
|
||||
self._backend = "auto"
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self._thread is not None and self._thread.is_alive()
|
||||
|
||||
@property
|
||||
def session_id(self) -> int:
|
||||
return self._session_id
|
||||
|
||||
@property
|
||||
def mode(self) -> Optional[FixedCameraMode]:
|
||||
return self._mode
|
||||
|
||||
def latest(self) -> Optional[SdkPacket]:
|
||||
with self._lock:
|
||||
return self._latest
|
||||
|
||||
def start(
|
||||
self,
|
||||
video_source: str | int = 0,
|
||||
*,
|
||||
config_path: str | Path,
|
||||
cal_path: Optional[str | Path] = None,
|
||||
backend: str = "auto",
|
||||
) -> FixedCameraMode:
|
||||
self.stop()
|
||||
ok, detail = orisys_importable()
|
||||
if not ok:
|
||||
raise RuntimeError(
|
||||
"orisys SDK is not importable. Install/build orisys-sdk and ensure "
|
||||
f"it is on PYTHONPATH. Detail: {detail}"
|
||||
)
|
||||
mode = load_fixed_mode(config_path)
|
||||
cal = Path(cal_path) if cal_path else None
|
||||
if cal is not None and not cal.exists():
|
||||
raise FileNotFoundError(f"Calibration not found: {cal}")
|
||||
|
||||
self._mode = mode
|
||||
self._video_source = video_source
|
||||
self._cal_path = str(cal) if cal is not None else None
|
||||
self._backend = backend or "auto"
|
||||
self._latest = None
|
||||
self._session_id += 1
|
||||
session = self._session_id
|
||||
self._stop.clear()
|
||||
self._thread = threading.Thread(
|
||||
target=self._run,
|
||||
name=f"SdkAcquisition-{session}",
|
||||
daemon=True,
|
||||
args=(session,),
|
||||
)
|
||||
self._thread.start()
|
||||
return mode
|
||||
|
||||
def stop(self) -> None:
|
||||
self._stop.set()
|
||||
thread = self._thread
|
||||
if thread is not None and thread.is_alive():
|
||||
thread.join(timeout=5.0)
|
||||
self._thread = None
|
||||
self._close_sensor()
|
||||
with self._lock:
|
||||
self._latest = None
|
||||
|
||||
def _close_sensor(self) -> None:
|
||||
sensor = self._sensor
|
||||
self._sensor = None
|
||||
if sensor is None:
|
||||
return
|
||||
try:
|
||||
sensor.disconnect()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
def _status(self, message: str) -> None:
|
||||
if self._on_status is not None:
|
||||
self._on_status(message)
|
||||
|
||||
def _run(self, session: int) -> None:
|
||||
import orisys
|
||||
|
||||
mode = self._mode
|
||||
source = self._video_source
|
||||
if mode is None or source is None:
|
||||
return
|
||||
try:
|
||||
self._status(f"Opening Orisys Sensor ({mode.label()})…")
|
||||
kwargs = {
|
||||
"config_name": mode.config_ref,
|
||||
"verbose": False,
|
||||
"backend": self._backend,
|
||||
}
|
||||
if self._cal_path:
|
||||
kwargs["cal_path"] = self._cal_path
|
||||
self._sensor = orisys.Sensor(source, **kwargs)
|
||||
self._status(f"Streaming {mode.label()}")
|
||||
except Exception as exc: # noqa: BLE001
|
||||
if self._on_error is not None and session == self._session_id:
|
||||
self._on_error(str(exc))
|
||||
return
|
||||
|
||||
frame_times: list[float] = []
|
||||
while not self._stop.is_set() and session == self._session_id:
|
||||
try:
|
||||
t0 = time.perf_counter()
|
||||
processed = self._sensor.get_img()
|
||||
if processed is None:
|
||||
time.sleep(0.005)
|
||||
continue
|
||||
raw = getattr(self._sensor, "frame", None)
|
||||
if raw is None:
|
||||
time.sleep(0.005)
|
||||
continue
|
||||
|
||||
try:
|
||||
self._sensor.compute_deformation(decompose=True, check_motion=False)
|
||||
self._sensor.compute_contact()
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
t1 = time.perf_counter()
|
||||
frame_times.append(t1)
|
||||
cutoff = t1 - 1.0
|
||||
frame_times = [t for t in frame_times if t >= cutoff]
|
||||
fps = float(len(frame_times)) if frame_times else 0.0
|
||||
|
||||
fn = fx = fy = 0.0
|
||||
try:
|
||||
fn, fx, fy = self._sensor.read_info(
|
||||
self._sensor.info.FNORMAL,
|
||||
self._sensor.info.FSHEARX,
|
||||
self._sensor.info.FSHEARY,
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
|
||||
raw_bgr = _ensure_bgr(raw)
|
||||
proc_bgr = _ensure_bgr(processed)
|
||||
packet = SdkPacket(
|
||||
raw_frame=raw_bgr,
|
||||
processed_frame=proc_bgr,
|
||||
timestamp=time.time(),
|
||||
fps=fps,
|
||||
process_ms=(t1 - t0) * 1000.0,
|
||||
width=int(raw_bgr.shape[1]),
|
||||
height=int(raw_bgr.shape[0]),
|
||||
processed_width=int(proc_bgr.shape[1]),
|
||||
processed_height=int(proc_bgr.shape[0]),
|
||||
normal_force=float(fn),
|
||||
shear_x=float(fx),
|
||||
shear_y=float(fy),
|
||||
session_id=session,
|
||||
mode=mode,
|
||||
)
|
||||
with self._lock:
|
||||
self._latest = packet
|
||||
except Exception as exc: # noqa: BLE001
|
||||
if self._on_error is not None and session == self._session_id:
|
||||
self._on_error(str(exc))
|
||||
break
|
||||
|
||||
self._close_sensor()
|
||||
self._status("Acquisition stopped")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,74 @@
|
||||
"""Load JSON configuration for sensor options and workflow steps."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OptionItem:
|
||||
id: str
|
||||
label: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SensorOptions:
|
||||
sensor_types: list[OptionItem]
|
||||
elastomer_types: list[OptionItem]
|
||||
cmos_types: list[OptionItem]
|
||||
marker_types: list[OptionItem]
|
||||
indentor_types: list[OptionItem]
|
||||
platform_types: list[OptionItem]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StepItem:
|
||||
id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkflowSteps:
|
||||
steps: list[StepItem]
|
||||
|
||||
|
||||
def _parse_options(raw: list[Any]) -> list[OptionItem]:
|
||||
items: list[OptionItem] = []
|
||||
for entry in raw:
|
||||
if isinstance(entry, str):
|
||||
items.append(OptionItem(id=entry, label=entry))
|
||||
continue
|
||||
item_id = str(entry.get("id", "")).strip()
|
||||
label = str(entry.get("label", item_id)).strip() or item_id
|
||||
if item_id:
|
||||
items.append(OptionItem(id=item_id, label=label))
|
||||
return items
|
||||
|
||||
|
||||
def load_sensor_options(path: Path) -> SensorOptions:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
return SensorOptions(
|
||||
sensor_types=_parse_options(data.get("sensor_types", [])),
|
||||
elastomer_types=_parse_options(data.get("elastomer_types", [])),
|
||||
cmos_types=_parse_options(data.get("cmos_types", [])),
|
||||
marker_types=_parse_options(data.get("marker_types", [])),
|
||||
indentor_types=_parse_options(data.get("indentor_types", [])),
|
||||
platform_types=_parse_options(
|
||||
data.get("platform_types", data.get("platoform_types", []))
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def load_workflow_steps(path: Path) -> WorkflowSteps:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
steps: list[StepItem] = []
|
||||
for entry in data.get("steps", []):
|
||||
if isinstance(entry, str):
|
||||
step_id = entry.strip()
|
||||
else:
|
||||
step_id = str(entry.get("id", f"step_{len(steps)+1}")).strip()
|
||||
if step_id:
|
||||
steps.append(StepItem(id=step_id))
|
||||
return WorkflowSteps(steps=steps)
|
||||
@@ -0,0 +1,77 @@
|
||||
"""Simple JSON-backed localization for Sensor Bench."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any, Optional
|
||||
|
||||
|
||||
DEFAULT_LOCALE = "zh_CN"
|
||||
FALLBACK_LOCALE = "en_US"
|
||||
|
||||
|
||||
class Translator:
|
||||
"""Lookup translated strings with English fallback."""
|
||||
|
||||
def __init__(self, locales_dir: Path, *, locale: str = DEFAULT_LOCALE):
|
||||
self.locales_dir = Path(locales_dir)
|
||||
self._catalogs: dict[str, dict[str, Any]] = {}
|
||||
self.locale = locale
|
||||
self.reload()
|
||||
|
||||
def available_locales(self) -> list[str]:
|
||||
found = sorted(p.stem for p in self.locales_dir.glob("*.json"))
|
||||
return found or [DEFAULT_LOCALE, FALLBACK_LOCALE]
|
||||
|
||||
def reload(self) -> None:
|
||||
self._catalogs.clear()
|
||||
for path in self.locales_dir.glob("*.json"):
|
||||
try:
|
||||
self._catalogs[path.stem] = json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception: # noqa: BLE001
|
||||
continue
|
||||
if self.locale not in self._catalogs and DEFAULT_LOCALE in self._catalogs:
|
||||
self.locale = DEFAULT_LOCALE
|
||||
elif self.locale not in self._catalogs and FALLBACK_LOCALE in self._catalogs:
|
||||
self.locale = FALLBACK_LOCALE
|
||||
|
||||
def set_locale(self, locale: str) -> None:
|
||||
self.locale = locale
|
||||
if locale not in self._catalogs:
|
||||
self.reload()
|
||||
|
||||
def t(self, key: str, **kwargs: Any) -> str:
|
||||
text = self._lookup(self.locale, key)
|
||||
if text is None:
|
||||
text = self._lookup(FALLBACK_LOCALE, key)
|
||||
if text is None:
|
||||
text = key
|
||||
if kwargs:
|
||||
try:
|
||||
return str(text).format(**kwargs)
|
||||
except Exception: # noqa: BLE001
|
||||
return str(text)
|
||||
return str(text)
|
||||
|
||||
def step_title(self, step_id: str) -> str:
|
||||
return self.t(f"workflow.steps.{step_id}.title")
|
||||
|
||||
def step_body(self, step_id: str) -> str:
|
||||
return self.t(f"workflow.steps.{step_id}.body")
|
||||
|
||||
def workflow_title(self) -> str:
|
||||
return self.t("workflow.title")
|
||||
|
||||
def _lookup(self, locale: str, key: str) -> Optional[str]:
|
||||
data: Any = self._catalogs.get(locale)
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
node: Any = data
|
||||
for part in key.split("."):
|
||||
if not isinstance(node, dict) or part not in node:
|
||||
return None
|
||||
node = node[part]
|
||||
if isinstance(node, str):
|
||||
return node
|
||||
return None
|
||||
@@ -0,0 +1,24 @@
|
||||
"""Entry point for the Sensor Bench desktop application."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from PyQt5.QtWidgets import QApplication
|
||||
|
||||
from .app import MainWindow
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = list(sys.argv if argv is None else argv)
|
||||
qt_app = QApplication(args)
|
||||
qt_app.setApplicationName("Sensor Bench")
|
||||
qt_app.setOrganizationName("Orisys-Tech")
|
||||
|
||||
window = MainWindow()
|
||||
window.show()
|
||||
return qt_app.exec_()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,80 @@
|
||||
"""Resolve project resource paths for editable and packaged runs."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def is_frozen() -> bool:
|
||||
return bool(getattr(sys, "frozen", False))
|
||||
|
||||
|
||||
def app_root() -> Path:
|
||||
"""Return the editable repo root or frozen executable directory."""
|
||||
if is_frozen():
|
||||
return Path(sys.executable).resolve().parent
|
||||
return Path(__file__).resolve().parents[2]
|
||||
|
||||
|
||||
def resource_root() -> Path:
|
||||
"""Return bundled resource root for editable and PyInstaller runs."""
|
||||
if is_frozen():
|
||||
return Path(getattr(sys, "_MEIPASS", app_root())).resolve()
|
||||
return app_root()
|
||||
|
||||
|
||||
def project_root() -> Path:
|
||||
"""Return the editable repo root or frozen executable directory."""
|
||||
return app_root()
|
||||
|
||||
|
||||
def default_options_path() -> Path:
|
||||
external = app_root() / "config" / "options.json"
|
||||
if external.exists():
|
||||
return external
|
||||
return resource_root() / "config" / "options.json"
|
||||
|
||||
|
||||
def default_steps_path() -> Path:
|
||||
external = app_root() / "steps" / "default.json"
|
||||
if external.exists():
|
||||
return external
|
||||
return resource_root() / "steps" / "default.json"
|
||||
|
||||
|
||||
def default_locales_dir() -> Path:
|
||||
external = app_root() / "config" / "locales"
|
||||
if external.exists():
|
||||
return external
|
||||
return resource_root() / "config" / "locales"
|
||||
|
||||
|
||||
def default_recordings_dir() -> Path:
|
||||
path = app_root() / "recordings"
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def default_orisys_sdk_root() -> Path:
|
||||
"""Sibling ``orisys-sdk`` checkout when sensor-bench lives next to it."""
|
||||
sibling = app_root().parent / "orisys-sdk"
|
||||
if sibling.exists():
|
||||
return sibling
|
||||
return Path(r"D:\repo\orisys-sdk")
|
||||
|
||||
|
||||
def default_ddjx01_config() -> str:
|
||||
if is_frozen():
|
||||
return "ddjx01"
|
||||
path = default_orisys_sdk_root() / "config" / "ddjx01.json"
|
||||
if path.exists():
|
||||
return str(path)
|
||||
return "ddjx01"
|
||||
|
||||
|
||||
def default_ddjx01_calibration() -> Path:
|
||||
external = app_root() / "config" / "ddjx01.npy"
|
||||
if external.exists():
|
||||
return external
|
||||
return default_orisys_sdk_root() / "config" / "ddjx01.npy"
|
||||
@@ -0,0 +1,85 @@
|
||||
"""Simple rolling history for live performance curves."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlotState:
|
||||
times: np.ndarray
|
||||
values: np.ndarray # shape (n, channels)
|
||||
x_range: tuple[float, float]
|
||||
y_ranges: tuple[tuple[float, float], ...]
|
||||
|
||||
|
||||
class MetricHistory:
|
||||
"""Ring-style buffer for FPS / brightness (or other scalar metrics)."""
|
||||
|
||||
def __init__(self, npoints: int = 120, *, channels: int = 2, window_sec: float = 4.0):
|
||||
self.npoints = int(npoints)
|
||||
self.channels = int(channels)
|
||||
self.window_sec = float(window_sec)
|
||||
self.reset()
|
||||
|
||||
def reset(self) -> None:
|
||||
self._data = np.full((self.npoints, 1 + self.channels), np.nan, dtype=np.float64)
|
||||
self.sample_count = 0
|
||||
self._t0: Optional[float] = None
|
||||
|
||||
def push(self, timestamp: float, *values: float) -> None:
|
||||
if len(values) != self.channels:
|
||||
raise ValueError(f"Expected {self.channels} values, got {len(values)}")
|
||||
if self._t0 is None:
|
||||
self._t0 = float(timestamp)
|
||||
t = float(timestamp) - self._t0
|
||||
row = (t, *map(float, values))
|
||||
if self.sample_count < self.npoints:
|
||||
self._data[self.sample_count, :] = row
|
||||
self.sample_count += 1
|
||||
return
|
||||
self._data[:-1, :] = self._data[1:, :]
|
||||
self._data[-1, :] = row
|
||||
|
||||
def state(self) -> Optional[PlotState]:
|
||||
n = self.sample_count
|
||||
if n == 0:
|
||||
return None
|
||||
block = self._data[:n]
|
||||
times = block[:, 0].copy()
|
||||
values = block[:, 1:].copy()
|
||||
if n < self.npoints:
|
||||
x_range = (0.0, self.window_sec)
|
||||
else:
|
||||
x0 = float(times[0])
|
||||
x1 = float(times[-1])
|
||||
if x1 <= x0:
|
||||
x1 = x0 + self.window_sec
|
||||
x_range = (x0, x1)
|
||||
|
||||
y_ranges = []
|
||||
for i in range(self.channels):
|
||||
col = values[:, i]
|
||||
ymin = float(np.nanmin(col))
|
||||
ymax = float(np.nanmax(col))
|
||||
if not np.isfinite(ymin) or not np.isfinite(ymax):
|
||||
ymin, ymax = 0.0, 1.0
|
||||
if abs(ymax - ymin) < 1e-6:
|
||||
pad = 1.0 if abs(ymin) < 1e-6 else abs(ymin) * 0.1
|
||||
ymin -= pad
|
||||
ymax += pad
|
||||
else:
|
||||
pad = (ymax - ymin) * 0.12
|
||||
ymin -= pad
|
||||
ymax += pad
|
||||
y_ranges.append((ymin, ymax))
|
||||
|
||||
return PlotState(
|
||||
times=times,
|
||||
values=values,
|
||||
x_range=x_range,
|
||||
y_ranges=tuple(y_ranges),
|
||||
)
|
||||
@@ -0,0 +1,82 @@
|
||||
"""Raw video recording helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
class RawRecorder:
|
||||
"""Write raw camera frames to a timestamped MP4 file."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._writer: Optional[cv2.VideoWriter] = None
|
||||
self.path: Optional[Path] = None
|
||||
self.frame_count = 0
|
||||
self._size: Optional[tuple[int, int]] = None
|
||||
|
||||
@property
|
||||
def is_recording(self) -> bool:
|
||||
return self._writer is not None
|
||||
|
||||
def start(self, output_dir: Path, *, fps: float, frame: np.ndarray) -> Path:
|
||||
self.stop()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
path = output_dir / f"raw_{stamp}.mp4"
|
||||
|
||||
write_frame = self._to_bgr(frame)
|
||||
height, width = write_frame.shape[:2]
|
||||
fps_out = float(fps) if fps and fps > 1.0 else 30.0
|
||||
writer = cv2.VideoWriter(
|
||||
str(path),
|
||||
cv2.VideoWriter_fourcc(*"mp4v"),
|
||||
fps_out,
|
||||
(width, height),
|
||||
)
|
||||
if not writer.isOpened():
|
||||
raise RuntimeError(f"Unable to create recording file: {path}")
|
||||
|
||||
self._writer = writer
|
||||
self.path = path
|
||||
self.frame_count = 0
|
||||
self._size = (width, height)
|
||||
self.write(frame)
|
||||
return path
|
||||
|
||||
def write(self, frame: np.ndarray) -> None:
|
||||
if self._writer is None:
|
||||
return
|
||||
write_frame = self._to_bgr(frame)
|
||||
height, width = write_frame.shape[:2]
|
||||
if self._size is not None and (width, height) != self._size:
|
||||
write_frame = cv2.resize(write_frame, self._size, interpolation=cv2.INTER_AREA)
|
||||
self._writer.write(write_frame)
|
||||
self.frame_count += 1
|
||||
|
||||
def stop(self) -> tuple[Optional[Path], int]:
|
||||
path = self.path
|
||||
count = self.frame_count
|
||||
if self._writer is not None:
|
||||
self._writer.release()
|
||||
self._writer = None
|
||||
self.path = None
|
||||
self.frame_count = 0
|
||||
self._size = None
|
||||
return path, count
|
||||
|
||||
@staticmethod
|
||||
def _to_bgr(frame: np.ndarray) -> np.ndarray:
|
||||
if frame.ndim == 2:
|
||||
return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
|
||||
if frame.ndim == 3 and frame.shape[2] == 1:
|
||||
return cv2.cvtColor(frame[:, :, 0], cv2.COLOR_GRAY2BGR)
|
||||
if frame.ndim == 3 and frame.shape[2] == 4:
|
||||
return cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR)
|
||||
if frame.ndim == 3 and frame.shape[2] == 3:
|
||||
return frame
|
||||
raise ValueError(f"Unsupported frame shape for recording: {frame.shape}")
|
||||
Reference in New Issue
Block a user