ykcheungab
412 lines
15 KiB
Python
412 lines
15 KiB
Python
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)
|