diff --git a/src/orisys/__init__.py b/src/orisys/__init__.py index 0b7eb55..622b754 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 034169d..0147700 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 UVUnrollFrameProcessor(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)