from __future__ import annotations import cv2 import numpy as np from pathlib import Path from .unroll import UVRemap class CompoundEyeFrameProcessor: """Convert a raw compound-eye frame into one stitched BGR tactile image.""" def __init__(self, stitcher): self.stitcher = stitcher def process(self, frame): self.stitcher.get_stitch_unit(frame) self.stitcher.extract_tiles() return self.stitcher.stitch_rt_optimized() class SingleCameraFrameProcessor: """Convert a normal camera frame into one BGR tactile image.""" def process(self, frame): if frame.ndim == 2: return cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR) if frame.shape[2] == 1: return cv2.cvtColor(frame[:, :, 0], cv2.COLOR_GRAY2BGR) return frame class FrameUnwarper: """Remap a raw camera frame into a flattened UV-domain BGR image.""" def __init__( self, remap: UVRemap, camera_matrix=None, dist_coeffs=None, source_width=None, source_height=None, flip_source_vertical=True, border_value=(0, 0, 0), ): self.remap = remap self.camera_matrix = camera_matrix self.dist_coeffs = dist_coeffs meta = remap.metadata self.source_width = source_width if source_width is not None else int(meta.get("source_width", 0) or 0) self.source_height = source_height if source_height is not None else int(meta.get("source_height", 0) or 0) self.flip_source_vertical = flip_source_vertical self.border_value = border_value self._flip_buffer = None self._gray_source_buffer = None self._remap_dst_bgr = None self._remap_dst_gray = None self.runtime_remap = remap self._use_composed_source_map = False self._build_runtime_remap() @property def valid_mask(self): """(H, W) bool mask in unrolled coordinates; pass to Sensor.get_img(valid_mask=...).""" return self.runtime_remap.valid_mask def _build_runtime_remap(self): """Compose flip/undistort into the UV remap when source geometry is known.""" can_compose = self.source_width > 0 and self.source_height > 0 and ( self.flip_source_vertical or (self.camera_matrix is not None and self.dist_coeffs is not None) ) if not can_compose: self.runtime_remap.ensure_fast_maps() return lookup_x = self.remap.map_x.astype(np.float32, copy=True) lookup_y = self.remap.map_y.astype(np.float32, copy=True) if self.flip_source_vertical: lookup_y = (self.source_height - 1.0) - lookup_y if self.camera_matrix is not None and self.dist_coeffs is not None: undistort_x, undistort_y = cv2.initUndistortRectifyMap( self.camera_matrix, self.dist_coeffs, None, self.camera_matrix, (self.source_width, self.source_height), cv2.CV_32FC1, ) composed_x = cv2.remap( undistort_x, lookup_x, lookup_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=-1, ) composed_y = cv2.remap( undistort_y, lookup_x, lookup_y, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=-1, ) else: composed_x = lookup_x composed_y = lookup_y invalid = ~self.remap.valid_mask composed_x[invalid] = -1.0 composed_y[invalid] = -1.0 self.runtime_remap = UVRemap( map_x=composed_x, map_y=composed_y, valid_mask=self.remap.valid_mask, metadata=dict(self.remap.metadata), ) self.runtime_remap.ensure_fast_maps() self._use_composed_source_map = True def _flip_vertical(self, frame_bgr): if not self.flip_source_vertical: return frame_bgr if ( self._flip_buffer is None or self._flip_buffer.shape != frame_bgr.shape or self._flip_buffer.dtype != frame_bgr.dtype ): self._flip_buffer = np.empty_like(frame_bgr) cv2.flip(frame_bgr, 0, dst=self._flip_buffer) return self._flip_buffer def prepare_source(self, frame_bgr): """Resize, undistort, and orient raw frame before UV remap. frame_bgr: (H_src, W_src, 3) raw camera BGR returns: (H_src, W_src, 3) prepared BGR (not unrolled) """ if frame_bgr is None: return None frame = frame_bgr if self.source_width > 0 and self.source_height > 0: if frame.shape[1] != self.source_width or frame.shape[0] != self.source_height: frame = cv2.resize( frame, (self.source_width, self.source_height), interpolation=cv2.INTER_LINEAR, ) if self._use_composed_source_map: return frame if self.camera_matrix is not None and self.dist_coeffs is not None: frame = cv2.undistort(frame, self.camera_matrix, self.dist_coeffs) return self._flip_vertical(frame) def prepare_source_gray(self, frame_bgr): """Resize/orient source frame and convert once to grayscale.""" frame = self.prepare_source(frame_bgr) if frame is None: return None if frame.ndim == 2: return frame if frame.ndim == 3 and frame.shape[2] == 1: return frame[:, :, 0] gray_shape = frame.shape[:2] if self._gray_source_buffer is None or self._gray_source_buffer.shape != gray_shape: self._gray_source_buffer = np.empty(gray_shape, dtype=np.uint8) cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY, dst=self._gray_source_buffer) return self._gray_source_buffer def _remap_dst(self, prepared): if prepared.ndim == 2: if ( self._remap_dst_gray is None or self._remap_dst_gray.shape != self.runtime_remap.valid_mask.shape or self._remap_dst_gray.dtype != prepared.dtype ): self._remap_dst_gray = np.empty(self.runtime_remap.valid_mask.shape, dtype=prepared.dtype) return self._remap_dst_gray shape = self.runtime_remap.valid_mask.shape + (prepared.shape[2],) if ( self._remap_dst_bgr is None or self._remap_dst_bgr.shape != shape or self._remap_dst_bgr.dtype != prepared.dtype ): self._remap_dst_bgr = np.empty(shape, dtype=prepared.dtype) return self._remap_dst_bgr def flatten_prepared(self, prepared_bgr): """Remap an already prepared source frame into UV flatten coordinates.""" if prepared_bgr is None: return None return self.runtime_remap.flatten( prepared_bgr, border_value=self.border_value, dst=self._remap_dst(prepared_bgr), ) def flatten_prepared_gray(self, prepared_gray): """Remap an already prepared grayscale frame into UV coordinates.""" if prepared_gray is None: return None return self.runtime_remap.flatten( prepared_gray, border_value=0, dst=self._remap_dst(prepared_gray), ) def process(self, frame_bgr): """Remap raw camera frame into UV flatten coordinates. frame_bgr: (H_src, W_src, 3) raw camera BGR returns: (H_out, W_out, 3) unrolled BGR """ return self.flatten_prepared(self.prepare_source(frame_bgr)) def create_frame_processor(sensor_type, stitcher=None, uv_remap=None, **uv_kwargs): sensor_type = (sensor_type or "compound_eye").lower() if sensor_type in ("compound_eye", "compound-eye", "compound"): if stitcher is None: raise ValueError("Compound-eye sensor requires a Stitcher instance.") return CompoundEyeFrameProcessor(stitcher) if sensor_type in ("single_camera", "camera", "mono"): return SingleCameraFrameProcessor() if sensor_type in ("uv_unroll", "unroll"): if uv_remap is None: 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)