diff --git a/build/OrisysDashboard.spec b/build/OrisysDashboard.spec index 86917ef..a4add95 100644 --- a/build/OrisysDashboard.spec +++ b/build/OrisysDashboard.spec @@ -6,7 +6,7 @@ a = Analysis( pathex=[], binaries=[], datas=[('D:\\repo\\orisys-sdk\\config', 'config'), ('D:\\repo\\orisys-sdk\\assets', 'assets')], - hiddenimports=['PyQt5', 'pyqtgraph', 'open3d', 'cv2', 'orisys', 'gui', 'pipeline', 'viz', 'app_paths', 'plotly', '_plotly_utils', 'open3d.cpu.pybind'], + hiddenimports=['PyQt5', 'pyqtgraph', 'open3d', 'cv2', 'orisys', 'gui', 'pipeline', 'viz', 'gui.app_paths', 'plotly', '_plotly_utils', 'open3d.cpu.pybind'], hookspath=[], hooksconfig={}, runtime_hooks=[], diff --git a/build/build_gui_win_pyinstaller.py b/build/build_gui_win_pyinstaller.py index 9ca0284..b336c1e 100644 --- a/build/build_gui_win_pyinstaller.py +++ b/build/build_gui_win_pyinstaller.py @@ -33,7 +33,7 @@ HIDDEN_IMPORTS = [ "gui", "pipeline", "viz", - "app_paths", + "gui.app_paths", "plotly", "_plotly_utils", # sometimes needed for open3d: diff --git a/build/wheel_layout.py b/build/wheel_layout.py index 00f2ec9..8409458 100644 --- a/build/wheel_layout.py +++ b/build/wheel_layout.py @@ -82,7 +82,7 @@ def copy_bundle_items(repo_root: Path, pkg_dir: Path, items: list[dict]) -> dict def package_data_from_bundle(ext_suffix: str, items: list[dict]) -> dict[str, list[str]]: package_data: dict[str, list[str]] = { "orisys": [f"*{ext_suffix}"], - "orisys.examples": ["*.py", "README.md", "app_paths.py", "viz/**/*.py", "pipeline/**/*.py", "gui/**/*.py"], + "orisys.examples": ["*.py", "README.md", "viz/**/*.py", "pipeline/**/*.py", "gui/**/*.py"], } for item in items: if item["name"] == "examples": diff --git a/examples/README.md b/examples/README.md index 8d5a197..8de1a4d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -5,7 +5,6 @@ Top-level demo scripts under `examples/`. Run them from the repository root. ## Top-Level Files - `README.md` — this overview of the root-level example entry points -- `app_paths.py` — path helpers used by the packaged / frozen dashboard entry point - `basic_pipeline.py` — basic single-sensor OpenCV demo for the compound-eye pipeline - `plot_data.py` — single-sensor demo with OpenCV views plus PyQtGraph force charts - `two_sensor_pipeline.py` — two-sensor multiprocessing demo for parallel acquisition / display @@ -59,7 +58,7 @@ pip install -e ".[gui]" ## Notes -- `finger_pipeline_3dview.py` is the main interactive demo. It uses `app_paths.py` so the same entry point works in editable and packaged runs. +- `finger_pipeline_3dview.py` is the main interactive demo. It uses `gui/app_paths.py` so the same entry point works in editable and packaged runs. - Dashboard controls: click **启动** to start acquisition; `Q` quits, `R` resets, `S` starts recording, `E` stops recording, and `Esc` toggles fullscreen. - `finger_pipeline.py` is the leanest finger example if you want to validate UV unroll + force/contact processing without the full dashboard. - `finger_pipeline.py` uses `orisys.finger.FingerRuntime` as the high-level entry point, then runs tactile analysis through `runtime.tactile`. diff --git a/examples/finger_pipeline.py b/examples/finger_pipeline.py index d15be17..deeaddf 100644 --- a/examples/finger_pipeline.py +++ b/examples/finger_pipeline.py @@ -12,7 +12,6 @@ For the full demo (UV debug panels, dashboard, benchmarks) see: import argparse import os import sys -import time os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0") @@ -20,19 +19,6 @@ import cv2 import orisys from orisys.finger import FingerRuntime -from pipeline.profiling import StageProfiler, run_deformation_timed - -_FINGER_PIPELINE_STAGES = ( - "capture", - "uv_unroll", - "flow", - "hhd", - "contact", - "read_info", - "display", - "total", -) - def main(): parser = argparse.ArgumentParser(description="Minimum UV-unroll finger tactile pipeline") @@ -41,17 +27,6 @@ def main(): parser.add_argument("--config", "-c", default="./config/dd02-ov.json", help="Orisys SDK config") parser.add_argument("--show", action="store_true", default=True, help="Show unrolled image window") parser.add_argument("--verbose", action="store_true", help="SDK verbose logs") - parser.add_argument( - "--profile", - action="store_true", - help="Print per-stage timing (capture, UV, flow, HHD, contact, read_info, display)", - ) - parser.add_argument( - "--profile-every", - type=int, - default=30, - help="When --profile is set, print rolling averages every N frames (default: 30)", - ) args = parser.parse_args() try: @@ -77,48 +52,17 @@ def main(): runtime.warmup(color=False) tactile = runtime.tactile - profiler = ( - StageProfiler(report_every=args.profile_every, stage_order=_FINGER_PIPELINE_STAGES) - if args.profile - else None - ) print("Press 'q' to quit, 'r' to reset tracker.") - if args.profile: - print(f"Profiling enabled (report every {args.profile_every} frames).") try: while True: - frame_t0 = time.perf_counter() + ok, unrolled = runtime.grab_and_process(color=False) + if not ok: + break - if profiler is not None: - t0 = time.perf_counter() - ok, frame = runtime.read_source() - profiler.record("capture", (time.perf_counter() - t0) * 1000.0) - if not ok: - break - t0 = time.perf_counter() - unrolled = runtime.process_frame(frame, color=False) - profiler.record("uv_unroll", (time.perf_counter() - t0) * 1000.0) - else: - ok, unrolled = runtime.grab_and_process(color=False) - if not ok: - break - - if profiler is not None: - flow_ms, hhd_ms = run_deformation_timed( - tactile, check_motion=True, threshold=0 - ) - profiler.record("flow", flow_ms) - profiler.record("hhd", hhd_ms) - else: - tactile.compute_deformation(check_motion=True, threshold=0) - - t0 = time.perf_counter() + tactile.compute_deformation(check_motion=True, threshold=0) tactile.compute_contact() - if profiler is not None: - profiler.record("contact", (time.perf_counter() - t0) * 1000.0) - t0 = time.perf_counter() fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = tactile.read_info( tactile.info.FPS, tactile.info.VRAW, @@ -131,8 +75,6 @@ def main(): tactile.info.FSHEARX, tactile.info.FSHEARY, ) - if profiler is not None: - profiler.record("read_info", (time.perf_counter() - t0) * 1000.0) status = ( f"FPS={fps:.2f} normal={fnormal:.2f} shear=({fshearx:.2f}, {fsheary:.2f}) " @@ -140,7 +82,6 @@ def main(): ) if args.show: - t0 = time.perf_counter() arrows = orisys.util.draw_arrows( unrolled, flow, @@ -149,17 +90,7 @@ def main(): arrow_scale=1.0, ) cv2.imshow("UV Unrolled", arrows) - if profiler is not None: - profiler.record("display", (time.perf_counter() - t0) * 1000.0) - if profiler is not None: - total_ms = (time.perf_counter() - frame_t0) * 1000.0 - profiler.record("total", total_ms) - status += ( - f" | wall={total_ms:.1f}ms flow={profiler.last.get('flow', 0):.1f}ms" - f" hhd={profiler.last.get('hhd', 0):.1f}ms" - ) - profiler.tick() print(status, end="\r") key = cv2.waitKey(1) & 0xFF @@ -170,8 +101,6 @@ def main(): tactile.reset() print("\nTracker reset.") finally: - if profiler is not None and profiler.frame_count: - profiler.print_report(prefix="[profile final]") runtime.close() cv2.destroyAllWindows() diff --git a/examples/finger_pipeline_3dview.py b/examples/finger_pipeline_3dview.py index 0567d93..ca0362d 100644 --- a/examples/finger_pipeline_3dview.py +++ b/examples/finger_pipeline_3dview.py @@ -15,7 +15,7 @@ import sys os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0") -from app_paths import configure_frozen_runtime, default_config_path, ensure_examples_import_path +from gui.app_paths import configure_frozen_runtime, default_config_path, ensure_examples_import_path configure_frozen_runtime() ensure_examples_import_path() diff --git a/examples/app_paths.py b/examples/gui/app_paths.py similarity index 86% rename from examples/app_paths.py rename to examples/gui/app_paths.py index 3136362..e227e3b 100644 --- a/examples/app_paths.py +++ b/examples/gui/app_paths.py @@ -1,4 +1,4 @@ -"""Resolve bundled resource paths for dev runs and Nuitka standalone builds.""" +"""Resolve bundled resource paths for dev runs and standalone builds.""" from __future__ import annotations @@ -6,7 +6,8 @@ import os import sys from pathlib import Path -_EXAMPLES_DIR = Path(__file__).resolve().parent +_GUI_DIR = Path(__file__).resolve().parent +_EXAMPLES_DIR = _GUI_DIR.parent _REPO_ROOT = _EXAMPLES_DIR.parent diff --git a/examples/gui/dashboard.py b/examples/gui/dashboard.py index 5b79021..a05e0d4 100644 --- a/examples/gui/dashboard.py +++ b/examples/gui/dashboard.py @@ -19,7 +19,7 @@ from PyQt5.QtWidgets import ( from pipeline.config import PipelineConfig from pipeline.service import FingerPipeline -from app_paths import resource_path +from .app_paths import resource_path from .layout import ( DashboardLayout, LAYOUT, diff --git a/examples/viz/mesh_camera.py b/examples/viz/mesh_camera.py index 6dc4f2d..8d05524 100644 --- a/examples/viz/mesh_camera.py +++ b/examples/viz/mesh_camera.py @@ -39,7 +39,7 @@ def resolve_viz_pose_path(path: str | Path | None = None) -> Path: here.parents[2] / raw, ] try: - from app_paths import app_root + from gui.app_paths import app_root candidates.insert(0, app_root() / raw) except ImportError: diff --git a/pyproject.toml b/pyproject.toml index a3a204a..c627639 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,7 +51,7 @@ dev_install = true name = "examples" source = "examples" target = "examples" -patterns = ["*.py", "README.md", "app_paths.py", "viz/**/*.py", "pipeline/**/*.py", "gui/**/*.py"] +patterns = ["*.py", "README.md", "viz/**/*.py", "pipeline/**/*.py", "gui/**/*.py"] preserve_tree = true [[tool.orisys.bundle.items]] diff --git a/src/orisys/__init__.py b/src/orisys/__init__.py index 3069532..3e985bf 100644 --- a/src/orisys/__init__.py +++ b/src/orisys/__init__.py @@ -12,6 +12,7 @@ from .sensor import Sensor from . import util from . import configs from . import finger +from .frame_processors import FFCFrameProcessor CompoundEyeSensor = Sensor CameraSensor = Sensor @@ -35,4 +36,7 @@ __all__ = [ # configs utilities functions 'list_config', 'export_config' + + # test + 'FFCFrameProcessor' ] diff --git a/src/orisys/frame_processors.py b/src/orisys/frame_processors.py index 6c7fa4d..331ce1e 100644 --- a/src/orisys/frame_processors.py +++ b/src/orisys/frame_processors.py @@ -1,5 +1,8 @@ +from __future__ import annotations + import cv2 import numpy as np +from pathlib import Path from .unroll import UVRemap @@ -229,3 +232,180 @@ def create_frame_processor(sensor_type, stitcher=None, uv_remap=None, **uv_kwarg raise ValueError("UV unroll sensor requires a UVRemap instance.") return FrameUnwarper(uv_remap, **uv_kwargs) raise ValueError(f"Unsupported sensor type: {sensor_type}") + + +class FFCFrameProcessor: + """Flat-field correction with offline weight-map calibration and fast online correction.""" + + def __init__(self, blur_kernel=151, blur_sigma=0.0, epsilon=1e-6, masked_outside=True): + self.blur_kernel = self.ensure_odd(blur_kernel) + self.blur_sigma = float(blur_sigma) + self.epsilon = float(epsilon) + self.masked_outside = bool(masked_outside) + self.weight_map = None + self.weight_map_3ch = None + self.roi_mask = None + self._multiplied_buffer = None + self._output_buffer = None + self._buffer_shape = None + self._buffer_dtype = None + self._roi_bbox = None + + def calibrate(self, reference_image, roi_mask=None): + gray = self.to_gray_float32(reference_image) + roi_mask = self.normalize_mask(roi_mask, gray.shape) + active = roi_mask > 0 + if not np.any(active): + raise ValueError("ROI mask is empty; cannot calibrate FFC.") + + background = cv2.GaussianBlur( + gray, + (self.blur_kernel, self.blur_kernel), + sigmaX=self.blur_sigma, + sigmaY=self.blur_sigma, + borderType=cv2.BORDER_REPLICATE, + ) + mean_value = float(background[active].mean()) + weight = mean_value / (background + self.epsilon) + + outside_value = 0.0 if self.masked_outside else 1.0 + self.weight_map = np.where(active, weight, outside_value).astype(np.float32) + self.weight_map_3ch = None + self.roi_mask = roi_mask + self._roi_bbox = self._compute_roi_bbox(roi_mask) + self._clear_runtime_buffers() + return self.weight_map + + def correct(self, image): + if self.weight_map is None: + raise RuntimeError("FFC is not calibrated. Call calibrate() or load() first.") + if image.shape[:2] != self.weight_map.shape: + raise ValueError( + f"image shape {image.shape[:2]} does not match weight map {self.weight_map.shape}" + ) + if image.ndim not in (2, 3): + raise ValueError("image must have shape HxW or HxWxC") + + self._ensure_runtime_buffers(image) + if self.masked_outside: + self._output_buffer.fill(0) + + if self._roi_bbox is None: + src = image + weight = self.weight_map if image.ndim == 2 else self.weight_map_3ch + dst_view = self._output_buffer + else: + x1, y1, x2, y2 = self._roi_bbox + src = image[y1:y2, x1:x2] + if image.ndim == 2: + weight = self.weight_map[y1:y2, x1:x2] + else: + weight = self.weight_map_3ch[y1:y2, x1:x2] + dst_view = self._output_buffer[y1:y2, x1:x2] + + cv2.multiply(src, weight, dst=self._multiplied_buffer, dtype=cv2.CV_32F) + if np.issubdtype(image.dtype, np.uint8): + cv2.convertScaleAbs(self._multiplied_buffer, dst=dst_view, alpha=1.0, beta=0.0) + return self._output_buffer + + corrected = self.saturate_to_dtype(self._multiplied_buffer, image.dtype) + dst_view[...] = corrected + return self._output_buffer + + def save(self, path): + if self.weight_map is None: + raise RuntimeError("No weight map to save.") + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed( + path, + weight_map=self.weight_map, + roi_mask=self.roi_mask, + blur_kernel=np.array(self.blur_kernel, dtype=np.int32), + blur_sigma=np.array(self.blur_sigma, dtype=np.float32), + epsilon=np.array(self.epsilon, dtype=np.float32), + masked_outside=np.array(self.masked_outside, dtype=np.bool_), + ) + + def load(self, path): + data = np.load(path) + self.weight_map = data["weight_map"].astype(np.float32, copy=False) + self.weight_map_3ch = None + self.roi_mask = data["roi_mask"].astype(np.uint8, copy=False) + self._roi_bbox = self._compute_roi_bbox(self.roi_mask) + self.blur_kernel = int(data.get("blur_kernel", self.blur_kernel)) + self.blur_sigma = float(data.get("blur_sigma", self.blur_sigma)) + self.epsilon = float(data.get("epsilon", self.epsilon)) + self.masked_outside = bool(data.get("masked_outside", self.masked_outside)) + self._clear_runtime_buffers() + return self + + def _clear_runtime_buffers(self): + self._multiplied_buffer = None + self._output_buffer = None + self._buffer_shape = None + self._buffer_dtype = None + + @staticmethod + def _compute_roi_bbox(roi_mask): + ys, xs = np.nonzero(roi_mask > 0) + if xs.size == 0: + return None + return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1 + + def _ensure_runtime_buffers(self, image): + shape = image.shape + dtype = image.dtype + if self._buffer_shape == shape and self._buffer_dtype == dtype: + return + + self._buffer_shape = shape + self._buffer_dtype = dtype + self._output_buffer = np.empty(shape, dtype=dtype) + + if image.ndim == 3: + channels = image.shape[2] + self.weight_map_3ch = np.repeat(self.weight_map[:, :, None], channels, axis=2) + if self._roi_bbox is None: + self._multiplied_buffer = np.empty(shape, dtype=np.float32) + else: + x1, y1, x2, y2 = self._roi_bbox + self._multiplied_buffer = np.empty((y2 - y1, x2 - x1, channels), dtype=np.float32) + else: + self.weight_map_3ch = None + if self._roi_bbox is None: + self._multiplied_buffer = np.empty(shape, dtype=np.float32) + else: + x1, y1, x2, y2 = self._roi_bbox + self._multiplied_buffer = np.empty((y2 - y1, x2 - x1), dtype=np.float32) + + @staticmethod + def ensure_odd(value, minimum=3): + value = max(int(value), int(minimum)) + return value if value % 2 == 1 else value + 1 + + @staticmethod + def to_gray_float32(image): + if image.ndim == 2: + gray = image + elif image.ndim == 3: + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + else: + raise ValueError("image must have shape HxW or HxWxC") + return gray.astype(np.float32, copy=False) + + @staticmethod + def normalize_mask(mask, shape): + if mask is None: + return np.ones(shape, dtype=np.uint8) * 255 + mask = np.asarray(mask) + if mask.shape[:2] != shape: + mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST) + return np.where(mask > 0, 255, 0).astype(np.uint8) + + @staticmethod + def saturate_to_dtype(values, dtype): + if np.issubdtype(dtype, np.integer): + info = np.iinfo(dtype) + return np.clip(values, info.min, info.max).astype(dtype) + return values.astype(dtype, copy=False) diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index c6da59a..0000000 --- a/tests/conftest.py +++ /dev/null @@ -1,116 +0,0 @@ -import os -import sys -from pathlib import Path - -import pytest - -REPO_ROOT = Path(__file__).resolve().parents[1] -SRC_ROOT = REPO_ROOT / "src" -if str(SRC_ROOT) not in sys.path: - sys.path.insert(0, str(SRC_ROOT)) - -from orisys import Sensor -from orisys import configs as orisys_configs -from orisys.backend_availability import check_cupy_available, check_opencl_available - - -def _resolve_test_video() -> Path: - env_video = os.getenv("ORISYS_TEST_VIDEO") - if env_video: - raw = Path(env_video) - candidates = [ - raw, # as provided (absolute or cwd-relative) - REPO_ROOT / raw, # repo-relative - ] - # Common typo fallback: "test/..." -> "tests/..." - if env_video.startswith("test/") or env_video.startswith("test\\"): - fixed = Path("tests", *Path(env_video).parts[1:]) - candidates.append(REPO_ROOT / fixed) - - for candidate in candidates: - if candidate.exists(): - return candidate - return raw - - # Prefer tests/ fixture clip; fallback kept by caller exists check. - return REPO_ROOT / "tests" / "idle_1min.mp4" - - -@pytest.fixture(scope="session") -def test_video_path() -> Path: - path = _resolve_test_video() - if not path.exists(): - pytest.skip( - "Test video not found. Set ORISYS_TEST_VIDEO to a valid video path." - ) - return path - - -@pytest.fixture(scope="session") -def test_config_name() -> str: - requested = os.getenv("ORISYS_TEST_CONFIG") - - def _is_usable_config(name: str) -> bool: - cfg = orisys_configs.load_config(name) - return "box_size" in cfg.get("stitcher", {}) - - if requested: - if not _is_usable_config(requested): - raise RuntimeError( - f"ORISYS_TEST_CONFIG='{requested}' is missing required key " - "'stitcher.box_size'. Use a config like 'ddjx01' or provide a full path." - ) - return requested - - # Auto-pick a stable config for tests (avoid prototype variants missing keys). - for name in orisys_configs.list_config(): - try: - if _is_usable_config(name): - return name - except Exception: - continue - - raise RuntimeError( - "No usable config found for tests (missing 'stitcher.box_size'). " - "Set ORISYS_TEST_CONFIG to a valid config JSON/path." - ) - - -@pytest.fixture(scope="session") -def available_backends() -> list[str]: - backends = ["cpu"] - if check_cupy_available(): - backends.append("cupy") - if check_opencl_available(): - backends.append("opencl") - return backends - - -def make_sensor( - video_path: Path, - config_name: str = "prototype", - backend: str = "cpu", - verbose: bool = False, -) -> Sensor: - use_cuda = backend in ("cupy", "opencl") - return Sensor( - vid_src=str(video_path), - cuda=use_cuda, - isstitch=True, - verbose=verbose, - config_name=config_name, - backend=backend, - ) - - -@pytest.fixture -def sensor_factory(test_video_path, test_config_name): - def _factory(backend: str = "cpu", verbose: bool = False): - return make_sensor( - test_video_path, - config_name=test_config_name, - backend=backend, - verbose=verbose, - ) - - return _factory diff --git a/tests/test_ffc.py b/tests/test_ffc.py new file mode 100644 index 0000000..0d3f831 --- /dev/null +++ b/tests/test_ffc.py @@ -0,0 +1,143 @@ +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import cv2 + +REPO_ROOT = Path(__file__).resolve().parents[1] +SRC_ROOT = REPO_ROOT / "src" +for path in (REPO_ROOT, SRC_ROOT): + if str(path) not in sys.path: + sys.path.insert(0, str(path)) + +from orisys import FFCFrameProcessor +from tools.ffc.roi_polygon_tool import load_roi_mask, select_polygon_roi + + +def parse_source(value): + value = str(value).strip() + for prefix in ("camera:", "cam:", "webcam:"): + if value.lower().startswith(prefix): + value = value[len(prefix):].strip() + break + try: + return int(value) + except ValueError: + return value + + +def read_image_if_possible(source): + if isinstance(source, int): + return None + path = Path(source) + if not path.exists() or not path.is_file(): + return None + image = cv2.imread(str(path), cv2.IMREAD_COLOR) + return image + + +def draw_roi_outline(image, roi_mask): + view = image.copy() + contours, _ = cv2.findContours(roi_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + cv2.drawContours(view, contours, -1, (0, 255, 255), 1) + return view + + +def parse_args(): + parser = argparse.ArgumentParser(description="Test ROI drawing and flat-field correction.") + parser.add_argument("--source", default="0", help="Camera id, image path, or video path.") + parser.add_argument("--roi-config", default=None, help="Existing ROI JSON config.") + parser.add_argument("--roi-output", default="roi_polygon.json", help="Path to save ROI JSON.") + parser.add_argument("--weights", default=None, help="Existing FFC weights .npz to load.") + parser.add_argument("--save-weights", default="ffc_weights.npz", help="Path to save calibrated weights.") + parser.add_argument("--blur-kernel", type=int, default=151) + parser.add_argument("--blur-sigma", type=float, default=0.0) + parser.add_argument("--epsilon", type=float, default=1e-6) + parser.add_argument("--keep-outside", action="store_true", help="Keep pixels outside ROI.") + return parser.parse_args() + + +def prepare_corrector(frame, args): + ffc = FFCFrameProcessor( + blur_kernel=args.blur_kernel, + blur_sigma=args.blur_sigma, + epsilon=args.epsilon, + masked_outside=not args.keep_outside, + ) + + if args.weights: + ffc.load(args.weights) + return ffc, ffc.roi_mask + + if args.roi_config: + roi_mask = load_roi_mask(args.roi_config, frame.shape) + else: + roi = select_polygon_roi(frame, output=args.roi_output) + if roi is None: + raise RuntimeError("ROI selection cancelled.") + roi_mask = roi["mask"] + + ffc.calibrate(frame, roi_mask=roi_mask) + if args.save_weights: + ffc.save(args.save_weights) + return ffc, roi_mask + + +def run_single_image(image, args): + ffc, roi_mask = prepare_corrector(image, args) + corrected = ffc.correct(image) + cv2.imshow("Input ROI", draw_roi_outline(image, roi_mask)) + cv2.imshow("Corrected ROI", draw_roi_outline(corrected, roi_mask)) + cv2.waitKey(0) + cv2.destroyAllWindows() + + +def run_video(source, args): + cap = cv2.VideoCapture(source) + if not cap.isOpened(): + raise RuntimeError(f"Cannot open source: {source}") + + try: + ok, frame = cap.read() + if not ok: + raise RuntimeError("Failed to read first frame.") + + ffc, roi_mask = prepare_corrector(frame, args) + + while True: + corrected = ffc.correct(frame) + cv2.imshow("Input ROI", draw_roi_outline(frame, roi_mask)) + cv2.imshow("Corrected ROI", draw_roi_outline(corrected, roi_mask)) + + key = cv2.waitKey(1) & 0xFF + if key in (27, ord("q")): + break + if key == ord("r"): + ffc.calibrate(frame, roi_mask=roi_mask) + if args.save_weights: + ffc.save(args.save_weights) + print("Recalibrated from current frame.") + + ok, frame = cap.read() + if not ok: + break + finally: + cap.release() + cv2.destroyAllWindows() + + +def main(): + args = parse_args() + source = parse_source(args.source) + + image = read_image_if_possible(source) + if image is not None: + run_single_image(image, args) + else: + run_video(source, args) + + +if __name__ == "__main__": + main() diff --git a/tools/ffc/roi_polygon_tool.py b/tools/ffc/roi_polygon_tool.py new file mode 100644 index 0000000..a2f9439 --- /dev/null +++ b/tools/ffc/roi_polygon_tool.py @@ -0,0 +1,264 @@ +from __future__ import annotations + +import json +from datetime import datetime +from pathlib import Path + +import cv2 +import numpy as np + + +WINDOW_NAME = "ROI Polygon Editor" + + +def to_uint8_image(image): + array = np.asarray(image) + if array.dtype == np.uint8: + return array + array = np.nan_to_num(array.astype(np.float32), copy=False) + if array.size > 0 and float(np.max(array)) <= 1.0: + array = array * 255.0 + return np.clip(array, 0, 255).astype(np.uint8) + + +def normalize_editor_frame(frame): + if not isinstance(frame, np.ndarray): + raise TypeError("frame must be a numpy ndarray with shape HxW or HxWxC") + + frame = to_uint8_image(frame) + if frame.ndim == 2: + return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR) + if frame.ndim != 3: + raise ValueError("frame must have shape HxW or HxWxC") + + channels = frame.shape[2] + if channels == 1: + return cv2.cvtColor(frame[:, :, 0], cv2.COLOR_GRAY2BGR) + if channels == 3: + return np.ascontiguousarray(frame) + if channels == 4: + return cv2.cvtColor(frame, cv2.COLOR_BGRA2BGR) + raise ValueError("frame channel count must be 1, 3, or 4") + + +def polygon_to_mask(points, image_size): + width, height = int(image_size[0]), int(image_size[1]) + mask = np.zeros((height, width), dtype=np.uint8) + polygon = np.asarray(points, dtype=np.int32).reshape(-1, 1, 2) + cv2.fillPoly(mask, [polygon], 255) + return mask + + +def save_polygon_roi_config(path, points, image_size, mask_path=None, metadata=None): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + data = { + "version": "1.0", + "type": "polygon", + "image_size": [int(image_size[0]), int(image_size[1])], + "points": [[int(x), int(y)] for x, y in points], + "created_at": datetime.now().isoformat(timespec="seconds"), + } + if mask_path is not None: + data["mask_path"] = str(mask_path) + if metadata: + data["metadata"] = metadata + with path.open("w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + return data + + +def load_roi_config(path): + path = Path(path) + with path.open("r", encoding="utf-8") as f: + data = json.load(f) + data["_config_path"] = str(path) + return data + + +def scale_points(points, source_size, target_size): + source_w, source_h = int(source_size[0]), int(source_size[1]) + target_w, target_h = int(target_size[0]), int(target_size[1]) + pts = np.asarray(points, dtype=np.float32).reshape(-1, 2) + if source_w > 0 and source_h > 0 and (source_w, source_h) != (target_w, target_h): + pts[:, 0] *= target_w / float(source_w) + pts[:, 1] *= target_h / float(source_h) + pts[:, 0] = np.clip(pts[:, 0], 0, max(target_w - 1, 0)) + pts[:, 1] = np.clip(pts[:, 1], 0, max(target_h - 1, 0)) + return np.rint(pts).astype(np.int32) + + +def normalize_mask(mask, shape): + if mask is None: + return np.ones(shape, dtype=np.uint8) * 255 + mask = np.asarray(mask) + if mask.shape[:2] != shape: + mask = cv2.resize(mask, (shape[1], shape[0]), interpolation=cv2.INTER_NEAREST) + return np.where(mask > 0, 255, 0).astype(np.uint8) + + +def load_roi_mask(config_path, frame_shape): + config_path = Path(config_path) + data = load_roi_config(config_path) + + height, width = frame_shape[:2] + points = data.get("points") or data.get("polygon") or data.get("roi_polygon") + if points: + source_size = data.get("image_size", [width, height]) + polygon = scale_points(points, source_size, (width, height)).reshape(-1, 1, 2) + output = np.zeros((height, width), dtype=np.uint8) + cv2.fillPoly(output, [polygon], 255) + return output + + mask_path = data.get("mask_path") or data.get("roi_mask") + if mask_path: + mask_path = Path(mask_path) + if not mask_path.is_absolute(): + mask_path = config_path.parent / mask_path + loaded = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE) + if loaded is None: + raise FileNotFoundError(f"Failed to read ROI mask: {mask_path}") + return normalize_mask(loaded, (height, width)) + + raise ValueError(f"Unsupported ROI config: {config_path}") + + +class RoiPolygonEditor: + def __init__(self, frame): + self.frame = normalize_editor_frame(frame) + self.points = [] + + def set_frame(self, frame): + self.frame = normalize_editor_frame(frame) + self.points = [] + + def on_mouse(self, event, x, y, flags, param): + if event == cv2.EVENT_LBUTTONDOWN: + self.points.append((int(x), int(y))) + elif event == cv2.EVENT_RBUTTONDOWN and self.points: + self.points.pop() + + def clear(self): + self.points = [] + + def undo(self): + if self.points: + self.points.pop() + + def render(self): + view = self.frame.copy() + if len(self.points) >= 3: + overlay = view.copy() + polygon = np.asarray(self.points, dtype=np.int32).reshape(-1, 1, 2) + cv2.fillPoly(overlay, [polygon], (0, 128, 255)) + view = cv2.addWeighted(overlay, 0.25, view, 0.75, 0.0) + cv2.polylines(view, [polygon], True, (0, 255, 255), 2, cv2.LINE_AA) + elif len(self.points) >= 2: + cv2.polylines( + view, + [np.asarray(self.points, dtype=np.int32).reshape(-1, 1, 2)], + False, + (0, 255, 255), + 2, + cv2.LINE_AA, + ) + + for idx, (x, y) in enumerate(self.points): + cv2.circle(view, (x, y), 4, (0, 255, 0), -1, cv2.LINE_AA) + cv2.putText( + view, + str(idx + 1), + (x + 6, y - 6), + cv2.FONT_HERSHEY_SIMPLEX, + 0.45, + (0, 255, 0), + 1, + cv2.LINE_AA, + ) + + lines = [ + "LClick add RClick/U undo C clear R recapture S/Enter save Q/Esc quit", + f"points={len(self.points)} resolution={view.shape[1]}x{view.shape[0]}", + ] + for i, text in enumerate(lines): + y = 24 + i * 24 + cv2.putText(view, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 0), 3, cv2.LINE_AA) + cv2.putText(view, text, (10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA) + return view + + +def build_roi_result(points, image_size, mask, config=None, config_path=None, mask_path=None): + return { + "points": [[int(x), int(y)] for x, y in points], + "image_size": [int(image_size[0]), int(image_size[1])], + "mask": mask, + "config": config, + "config_path": str(config_path) if config_path is not None else None, + "mask_path": str(mask_path) if mask_path is not None else None, + } + + +def select_polygon_roi( + frame, + output=None, + window_name=WINDOW_NAME, + metadata=None, + save_mask=True, + destroy_window=True, + recapture_callback=None, +): + """Open an OpenCV window, draw a polygon ROI, and return points/mask/config info.""" + editor = RoiPolygonEditor(frame) + cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE) + cv2.setMouseCallback(window_name, editor.on_mouse) + + try: + while True: + cv2.imshow(window_name, editor.render()) + key = cv2.waitKey(16) & 0xFF + if key in (27, ord("q")): + return None + if key in (8, ord("u")): + editor.undo() + elif key == ord("c"): + editor.clear() + elif key == ord("r") and recapture_callback is not None: + editor.set_frame(recapture_callback()) + elif key in (13, ord("s")): + if len(editor.points) < 3: + print("Need at least 3 points to save ROI.") + continue + + image_size = (editor.frame.shape[1], editor.frame.shape[0]) + mask = polygon_to_mask(editor.points, image_size) + config = None + config_path = None + mask_path = None + + if output is not None: + config_path = Path(output) + config_path.parent.mkdir(parents=True, exist_ok=True) + saved_mask = None + if save_mask: + mask_path = config_path.with_name(config_path.stem + "_mask.png") + cv2.imwrite(str(mask_path), mask) + saved_mask = mask_path.name + config = save_polygon_roi_config( + config_path, + editor.points, + image_size, + mask_path=saved_mask, + metadata=metadata, + ) + + return build_roi_result( + editor.points, + image_size, + mask, + config=config, + config_path=config_path, + mask_path=mask_path, + ) + finally: + if destroy_window: + cv2.destroyWindow(window_name)