integrated finger process pipeline on remapped 2d image
This commit is contained in:
ykcheungab
@@ -1,6 +1,8 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from .unroll import UVRemap
|
||||
|
||||
|
||||
class CompoundEyeFrameProcessor:
|
||||
"""Convert a raw compound-eye frame into one stitched BGR tactile image."""
|
||||
@@ -25,7 +27,196 @@ class SingleCameraFrameProcessor:
|
||||
return frame
|
||||
|
||||
|
||||
def create_frame_processor(sensor_type, stitcher=None):
|
||||
class UVUnrollFrameProcessor:
|
||||
"""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:
|
||||
@@ -33,4 +224,8 @@ def create_frame_processor(sensor_type, stitcher=None):
|
||||
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 UVUnrollFrameProcessor(uv_remap, **uv_kwargs)
|
||||
raise ValueError(f"Unsupported sensor type: {sensor_type}")
|
||||
|
||||
Reference in New Issue
Block a user