""" Sensor - Core sensor class for optical tactile sensing. This module provides the Sensor class for capturing, processing, and analyzing optical tactile sensor data with support for stitching, optical flow tracking, and Helmholtz-Hodge decomposition. """ import os import numpy as np import cv2 import time import threading from enum import Enum, auto from typing import Tuple # from threading import Thread, Event # from queue import Queue, Empty from .Stitcher import Stitcher from .tracker import Tracker from .hhd import nHHD from . import configs from .get_depth import Contact3D from .backend_availability import (check_cupy_available, check_opencl_available) from .frame_processors import create_frame_processor class Sensor: """ Optical tactile sensor interface with deformation analysis. Provides image capture, stitching, optical flow tracking, and force decomposition. Supports both synchronous and asynchronous (pipelined) processing modes. """ # ========================================================================= # Info Enum (for read_info API) # ========================================================================= class info(Enum): """Data types available through read_info().""" VRAW = auto() VNORMAL = auto() VSHEAR = auto() FPS = auto() TIMESTAMP = auto() FNORMAL = auto() FSHEARX = auto() FSHEARY = auto() CENTROID = auto() CONTOUR = auto() DEPTH = auto() IMG = auto() # ========================================================================= # Initialization # ========================================================================= def __init__(self, vid_src, cuda=True, isstitch=True, verbose=True, debug=False, config_name="prototype", cal_path=None, override=False, box=None, mirror=True, backend='auto'): """ Initialize Sensor with video source and configuration. Args: vid_src: Video source (int for camera ID, str for file path) cuda: Enable CUDA acceleration (default: True) isstitch: Enable stitching mode for multi-tile sensor (default: False) verbose: Print initialization logs (default: True) config_name: Configuration name (default: "prototype") cal_path: Path to calibration file (default: None) override: Skip auto-detection, use provided box (default: False) box: Manual bounding boxes when override=True (default: None) backend: Poisson backend (default: "auto") "auto": Automatically select the best backend based on availability "cupy": Use CuPy for GPU acceleration "opencl": Use OpenCL for GPU acceleration "cpu": Use CPU for computation hhd_downsample: Downsample factor for flow field before HHD (default: 1, no downsampling). Factor of 2 halves each dimension, giving ~4x speedup in Poisson solve. """ self.CUPY_AVAILABLE = check_cupy_available() self.OPENCL_AVAILABLE = check_opencl_available() if vid_src is None: raise Exception("User must specify a video source.") self._external_only = isinstance(vid_src, str) and vid_src.lower() == "external" # Store settings self.vid_src = vid_src self.isstitch = isstitch self.cuda = cuda self.verbose = verbose self.debug = debug self.config_name = config_name self.cal_path = cal_path self.mirror = mirror # Load configuration self.__config_cam = None self.__config_stitcher = None self.__config_contact = None self._load_config() # State variables self.stitch_valid = False self.compute_deformation_valid = False self.compute_contact_valid = False # Image Data Variables (public) # self.stitch_calibration = [] self.frame = self.img = None self.initial_img = None self._img_bgr_raw = None self._img_gray_buf = None shape = (self.__config_cam["resolution_stitch"][1], self.__config_cam["resolution_stitch"][0]) self.img_size_raw = None self.img_filtered = None self.tail_frames_count = 0 # Optical Flow Data Variables (public) self.flow = np.zeros(shape + (2,), dtype=np.float32) self.vtemplate = np.zeros_like(self.flow) self.vnormal = self.vtemplate self.vshear = self.vtemplate self.n_boxes = None self.fnormal = self.fshearx = self.fsheary = self.ftor = 0.0 # Contact Data Variables (public) self.centroid = (0, 0) self.centroid_found = False self.contour = np.empty(0) self.depth_map = np.zeros_like(self.img) # Timing self.last_log_time = time.perf_counter() self._fps = 0 self._last_capture_time = 0.0 self._last_processing_time = 0.0 self._last_stage_times = { "extract_tiles": 0.0, "stitch": 0.0, "resize": 0.0, "mirror": 0.0, "astype": 0.0, } # Initialize components self.stitcher = ( Stitcher(isstitch=isstitch, config_stitcher=self.__config_stitcher, verbose=self.verbose) if self._is_compound_eye() else None ) self.frame_processor = create_frame_processor(self.sensor_type, stitcher=self.stitcher) self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, config_tracker=self.__config_tracker) self.__contact3d = Contact3D(config_contact=self.__config_contact, verbose=self.verbose) self.__hhd = None # Integer >= 1: run HHD on flow downsampled by this factor for speed (1 = full res). _raw_ds = self.__config_tracker.get("hhd_downsample", 1) try: self._hhd_downsample = max(1, int(_raw_ds)) if _raw_ds is not None else 1 except (TypeError, ValueError): self._hhd_downsample = 1 self._hhd_grid_shape = None cfg_backend = self.__config_tracker.get("backend", None) self.poisson_backend = cfg_backend if cfg_backend is not None else backend self._use_pyav = False self._av_container = None self._av_stream = None self._av_decoder = None self.cap = None self._is_live_camera = ( isinstance(self.vid_src, int) or (isinstance(self.vid_src, str) and self.vid_src.startswith("/dev/video")) ) and not self._external_only self._capture_thread = None self._capture_running = False self._latest_frame = None self._latest_frame_seq = 0 self._last_consumed_seq = -1 self._latest_frame_lock = threading.Lock() # Setup video capture unless this sensor is fed externally. if self._external_only: self.stitch_valid = False if self.verbose: print("Sensor initialized in external-frame mode.") else: self._open_video_source(override, box) if self.verbose: print(f" Video source: {vid_src}") print(f" Stitch mode: {self.stitch_valid}") print(f" Optical flow algorithm: {self.__config_tracker['algorithm']}") print(f" Poisson backend: {self.poisson_backend}") # print(f" CUDA acceleration: {self.cuda}") print(f" CuPy available: {self.CUPY_AVAILABLE}") print(f" OpenCL available: {self.OPENCL_AVAILABLE}") # ========================================================================= # Public API - Core Methods (Most Frequently Used) # ========================================================================= def get_img(self, frame=None, already_processed=False): """ Capture and preprocess a frame. Parameters ---------- frame : np.ndarray | None Optional externally supplied frame. When provided, the sensor skips reading from its own video source and uses this frame instead. already_processed : bool If True, treat `frame` as the final BGR tactile image and skip the configured frame processor. This is useful for pipelines that generate a derived image, such as an unrolled tactile view. Returns: np.ndarray | None: Processed image, or None on capture failure. """ # Simple direct read - buffer size of 1 should minimize frame accumulation # Timing measurement helps identify if capture or processing is the bottleneck t_capture_start = time.perf_counter() if frame is None: ret, frame = self._read_frame() else: ret = frame is not None t_capture_end = time.perf_counter() self._last_capture_time = t_capture_end - t_capture_start if not ret: print("Failed to get_img from video source.") return None # Start processing timing t_processing_start = time.perf_counter() self.frame = frame.copy() if isinstance(frame, np.ndarray) else frame if already_processed: img_bgr = frame if img_bgr.ndim == 2: img_bgr = cv2.cvtColor(img_bgr, cv2.COLOR_GRAY2BGR) elif img_bgr.ndim == 3 and img_bgr.shape[2] == 1: img_bgr = cv2.cvtColor(img_bgr[:, :, 0], cv2.COLOR_GRAY2BGR) else: img_bgr = self.frame_processor.process(frame) t_after_extract = time.perf_counter() t_after_stitch = time.perf_counter() self.img_size_raw = img_bgr.shape self._img_bgr_raw = img_bgr gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) res = tuple(self._config_resolution_stitch) if self._img_gray_buf is None or self._img_gray_buf.shape != (res[1], res[0]): self._img_gray_buf = np.empty((res[1], res[0]), dtype=np.uint8) cv2.resize(gray, res, dst=self._img_gray_buf, interpolation=cv2.INTER_LINEAR) t_after_resize = time.perf_counter() if self.mirror: cv2.flip(self._img_gray_buf, 1, dst=self._img_gray_buf) self.img = self._img_gray_buf t_processing_end = time.perf_counter() self._last_processing_time = t_processing_end - t_processing_start self._last_stage_times["extract_tiles"] = t_after_extract - t_processing_start self._last_stage_times["stitch"] = t_after_stitch - t_after_extract self._last_stage_times["resize"] = t_processing_end - t_after_stitch # FPS calculation self._fps = 1.0 / (time.perf_counter() - self.last_log_time) self.last_log_time = time.perf_counter() return self.img.copy() def _get_img_bgr(self): """Produce a resized+mirrored BGR image on demand (for display).""" if self._img_bgr_raw is None: return None res = tuple(self._config_resolution_stitch) bgr = cv2.resize(self._img_bgr_raw, res, interpolation=cv2.INTER_LINEAR) if self.mirror: cv2.flip(bgr, 1, dst=bgr) return bgr def detect_motion(self, threshold=6): # detect motion by optical flow magnitude mag = np.linalg.norm(self.flow, axis=-1) # print(f"mag: {mag.max()}, {mag.min()}, {mag.mean()}") return mag.mean() > threshold def compute_deformation(self, decompose=True, check_motion=True, threshold=6, tail_frames=10): """ Compute optical flow and optionally decompose it into force components. Parameters ---------- decompose : bool If True, run HHD decomposition (`compute_hhd`) after flow computation. check_motion : bool If True, gate decomposition with `detect_motion(threshold)`. When motion falls below threshold, decomposition may still run for up to `tail_frames` additional frames to avoid abrupt force drop-off. threshold : float Motion-detection threshold on mean flow magnitude. tail_frames : int Number of extra non-motion frames allowed to keep running HHD when `check_motion=True`. Returns ------- bool | None - `True` : motion is currently detected (`check_motion=True`) - `False` : motion is currently not detected (`check_motion=True`) - `None` : motion gating disabled (`check_motion=False`) """ if self.img is None: raise Exception("Warning: Get image first") self.compute_flow() if not check_motion: self.tail_frames_count = 0 if decompose: self.compute_hhd() return None is_motion = self.detect_motion(threshold=threshold) try: tail_frames = max(0, int(tail_frames)) except (TypeError, ValueError): tail_frames = 0 if is_motion: self.tail_frames_count = 0 if decompose: self.compute_hhd() return True # No motion this frame. Keep HHD alive briefly if requested, but do not # report this as "motion detected". if decompose and self.tail_frames_count < tail_frames: self.tail_frames_count += 1 self.compute_hhd() else: self.tail_frames_count = tail_frames return False def compute_flow(self): if self.img is None: raise Exception("Warning: Get image first") flow = self.__tracker.track(self.img) # flow = self.__tracker.track_raft(self.img) if hasattr(flow, 'download'): flow = flow.download() self.flow = flow def compute_hhd(self): flow = self.flow if flow is None: raise Exception("Warning: Compute flow first") orig_h, orig_w = flow.shape[:2] try: k = max(1, int(self._hhd_downsample)) except (TypeError, ValueError): k = 1 if k > 1: ds_w = max(1, orig_w // k) ds_h = max(1, orig_h // k) flow_ds = cv2.resize(flow, (ds_w, ds_h), interpolation=cv2.INTER_AREA) else: flow_ds = flow ds_shape = flow_ds.shape[:2] if self.__hhd is None or self._hhd_grid_shape != ds_shape: spacings = (18 * k, 18 * k) self.__hhd = nHHD(cuda=self.cuda, grid=ds_shape, spacings=spacings, poisson_backend=self.poisson_backend, debug=self.debug) self._hhd_grid_shape = ds_shape _, forces = self.__hhd.get_force(flow_ds) self.fnormal, self.fshearx, self.fsheary, self.ftor = forces vnormal_ds = self.__hhd.get_normal() vshear_ds = self.__hhd.get_shear() if k > 1: self.vnormal = cv2.resize(vnormal_ds, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR) self.vshear = cv2.resize(vshear_ds, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR) else: self.vnormal = vnormal_ds self.vshear = vshear_ds self.compute_deformation_valid = True return True def compute_contact(self, vnormal=None): """ Compute contact information from vnormal field. Args: vnormal: Optional vnormal field. If None, uses self.vnormal. Returns: dict: Contact result with centroid, contour, and depth_map """ multiple_foe = False # Use provided vnormal or fall back to internal state vnormal_to_use = vnormal if vnormal is not None else self.vnormal if vnormal_to_use is None: raise Exception("Warning: vnormal is required for contact computation") if multiple_foe: result = self.__contact3d.compute_multiple_foe(vnormal_to_use) else: if self.cuda: result = self.__contact3d.compute_foe_gpu(vnormal_to_use) else: result = self.__contact3d.compute_foe(vnormal_to_use) self.centroid = result['centroid'] self.centroid_found = result['centroid_found'] if 'contour' in result: self.contour = result['contour'] self.depth_map = result['divergence_map'] self.compute_contact_valid = True return result def read_info(self, *args) -> Tuple: """ Get sensor data by info type. Args: *args: info enum values (e.g., sensor.info.FPS, sensor.info.FNORMAL) Returns: Tuple: Requested data in order """ if any(arg in [self.info.IMG] for arg in args): if self.img is None: raise Exception("Warning: Get image first") # if any(arg in [self.info.VRAW, self.info.VNORMAL, self.info.VSHEAR, self.info.FNORMAL, self.info.FSHEARX, self.info.FSHEARY] for arg in args): # if not self.compute_deformation_valid: # print("Warning: Compute deformation first") # # raise Exception("Warning: Compute deformation first") # if any(arg in [self.info.CENTROID, self.info.CONTOUR, self.info.DEPTH] for arg in args): # if not self.compute_contact_valid: # print("Warning: Compute contact first") # raise Exception("Warning: Compute contact first") mapping = { # Return copies of arrays to prevent external modification of internal state self.info.VRAW: lambda: self.flow.copy() if self.flow is not None else np.array([]), self.info.VNORMAL: lambda: self.vnormal.copy() if self.vnormal is not None else np.array([]), self.info.VSHEAR: lambda: self.vshear.copy() if self.vshear is not None else np.array([]), # Scalar values are safe to return directly self.info.FNORMAL: lambda: self.fnormal, self.info.FSHEARX: lambda: self.fshearx, self.info.FSHEARY: lambda: self.fsheary, self.info.FPS: lambda: self._fps, self.info.TIMESTAMP: lambda: time.strftime('%H:%M:%S'), # Return copies of arrays to prevent external modification self.info.DEPTH: lambda: self.depth_map.copy() if self.depth_map is not None else np.array([]), # Centroid: can be tuple, list, or array - return a copy to prevent modification # Return None if centroid is not found or empty self.info.CENTROID: lambda: ( np.array(self.centroid).copy() if self.centroid is not None and len(self.centroid) > 0 else None ), self.info.CONTOUR: lambda: self.contour.copy() if self.contour is not None else np.array([]), self.info.IMG: lambda: self._get_img_bgr() if self._img_bgr_raw is not None else np.array([]), } return tuple(mapping.get(arg, lambda: None)() for arg in args if isinstance(arg, self.info)) def reset(self): """Reset optical flow tracker.""" self.__tracker.reset() if self.img is not None: self.initial_img = self.img.copy() def disconnect(self): """Release resources and stop async threads.""" self._capture_running = False if self._capture_thread is not None: self._capture_thread.join(timeout=1.0) self._capture_thread = None if self._use_pyav: if self._av_container is not None: self._av_container.close() self._av_container = None self._av_stream = None self._av_decoder = None if self.verbose: print("Sensor disconnected.") elif self.cap is not None and self.cap.isOpened(): self.cap.release() if self.verbose: print("Sensor disconnected.") cv2.destroyAllWindows() # ========================================================================= # Public API - Configuration & Calibration # ========================================================================= # def export_config(self, path="./config/exported_config.json"): # """Return current configuration dict.""" # configs.export_config(self._config, path) # if self.verbose: # print(f"Config exported to {path}") # return True # def stitch_cal(self): # """ # Load or create stitch calibration. # Returns: # np.ndarray: Stitched calibration image # """ # if self.cal_path is None: # self.cal_path = "./config/cal.npy" # if os.path.exists(self.cal_path): # if self.verbose: # print("Loading calibration from", self.cal_path) # self.stitch_calibration = np.load(self.cal_path, allow_pickle=True) # self.stitcher.stitch_calibration = self.stitch_calibration # return self.stitcher.stitch_rt_optimized() # else: # if self.verbose: # print("Creating new calibration...") # img, self.stitch_calibration = self.stitcher.stitch_cal() # os.makedirs("./config", exist_ok=True) # np.save(self.cal_path, self.stitch_calibration) # if self.verbose: # print("Calibration saved to", self.cal_path) # return img # ========================================================================= # Private - Initialization Helpers # ========================================================================= def _load_config(self): """Load sensor configuration from JSON file.""" self._config = configs.load_config(self.config_name) self.sensor_type = self._config.get( "type", "compound_eye" if "stitcher" in self._config else "single_camera", ) self.__config_cam = self._config["camera"] self._config_resolution_raw = self.__config_cam["resolution_raw"] self._config_fps = self.__config_cam["fps"] self._config_units = self.__config_cam.get("units", 1) self._config_resolution_stitch = self.__config_cam["resolution_stitch"] self._config_dis_preset = self._config["tracker"]["dis_preset"] self.__config_tracker = self._config["tracker"] self.__config_stitcher = self._config.get("stitcher") self.__config_contact = self._config["contact"] if self.verbose: print(f"Config: {self.config_name}") return True def _is_compound_eye(self): return str(self.sensor_type).lower() in ("compound_eye", "compound-eye", "compound") def _open_video_source(self, override, box): """Initialize video capture and calibration.""" # _is_camera_index = isinstance(self.vid_src, int) # _is_v4l2_path = isinstance(self.vid_src, str) and self.vid_src.startswith("/dev/video") # use_linux_camera_pyav = ( # platform.system() == "Linux" # and (_is_camera_index or _is_v4l2_path) # and av is not None # ) # if use_linux_camera_pyav: # device = self.vid_src if _is_v4l2_path else f"/dev/video{self.vid_src}" # if self.verbose: # print(f"Using PyAV Linux camera backend: {device}") # try: # self._av_container = av.open( # device, # format="v4l2", # options={ # "input_format": "mjpeg", # "video_size": f"{self._config_resolution_raw[0]}x{self._config_resolution_raw[1]}", # "framerate": str(self._config_fps), # }, # ) # self._av_stream = self._av_container.streams.video[0] # self._av_decoder = self._av_container.decode(self._av_stream) # self._use_pyav = True # except Exception as exc: # if self.verbose: # print(f"PyAV camera open failed ({exc}), fallback to OpenCV backend.") # self._use_pyav = False self.use_pyav = False if not self._use_pyav: self.cap = cv2.VideoCapture(self.vid_src) if self._is_live_camera: self.cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter.fourcc('M', 'J', 'P', 'G')) self.cap.set(cv2.CAP_PROP_FPS, self._config_fps) self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, self._config_resolution_raw[0]) self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, self._config_resolution_raw[1]) if self.verbose: print("Connecting to video source...") if not self._is_source_open(): raise Exception("Failed to open video source.") if self.verbose and not self._use_pyav and self.cap is not None: fourcc_int = int(self.cap.get(cv2.CAP_PROP_FOURCC)) fourcc_bytes = [(fourcc_int >> (8 * i)) & 0xFF for i in range(4)] # Show printable FOURCC only; some backends return 0 or non-printable bytes. if fourcc_int <= 0: fourcc = "N/A" else: fourcc = "".join(chr(b) if 32 <= b <= 126 else "." for b in fourcc_bytes) width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = self.cap.get(cv2.CAP_PROP_FPS) backend_id = int(self.cap.get(cv2.CAP_PROP_BACKEND)) print( f"OpenCV negotiated stream: FOURCC='{fourcc}' " f"(int={fourcc_int}, hex=0x{fourcc_int:08X}), " f"size={width}x{height}, fps={fps:.2f}, backend_id={backend_id}" ) # Wait for first valid frame frame = None while frame is None: ret, frame = self._read_frame_raw() if not ret: time.sleep(0.1) if not self._is_compound_eye(): self.stitch_valid = False self.initial_img = self._make_initial_gray(frame) elif override: if box is not None: self.stitcher.isstitch = self.isstitch = False self.stitcher.boxes = box else: print("Override=True but no box provided.") else: self._find_stitch_boxes(frame) img = self.frame_processor.process(frame) self.initial_img = self._make_initial_gray(img) # For live cameras, start a latest-frame capture thread to avoid latency buildup. if self._is_live_camera: self._start_capture_thread() def _make_initial_gray(self, img_bgr): gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY) self.resolution_stitch = gray.shape initial = cv2.resize(gray, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_CUBIC) if self.mirror: initial = cv2.flip(initial, 1) return initial.astype(np.uint8) def _find_stitch_boxes(self, img_src, show=False): """Detect vision tile boxes with camera stabilization.""" n_boxes, v_mean0, timeout = 0, 0, 5 start_time = cv2.getTickCount() while n_boxes < self._config_units: ret, frame = self._read_frame_raw() if not ret: continue v_mean = np.sum(cv2.mean(frame)) if abs(v_mean - v_mean0) < 5: n_boxes, boxes = self.stitcher.find_stitch_boxes(frame, show=show) if (cv2.getTickCount() - start_time) / cv2.getTickFrequency() > timeout: print("Timeout finding vision units.") self.stitcher.isstitch = self.isstitch = False break v_mean0 = v_mean self.n_boxes = n_boxes self.boxes = boxes self.stitch_valid = self.isstitch and n_boxes == self._config_units status = "valid for stitching" if self.stitch_valid else "non-stitch mode" print(f"Sensor: {status}") return n_boxes def _extract_tiles(self, img_src): """Extract tiles from source image.""" self.stitcher.get_stitch_unit(img_src) self.stitcher.extract_tiles() return True def _is_source_open(self): if self._use_pyav: return self._av_container is not None return self.cap is not None and self.cap.isOpened() def _read_frame(self): if self._capture_running: # Startup race guard: wait briefly until capture thread publishes first frame. if self._latest_frame is None: deadline = time.perf_counter() + 1.0 while time.perf_counter() < deadline: with self._latest_frame_lock: if self._latest_frame is not None: break time.sleep(0.001) # Return latest frame immediately (do not wait for next camera tick). # This keeps end-to-end latency minimal; caller may receive repeated # frames if processing runs faster than camera FPS. with self._latest_frame_lock: frame = self._latest_frame seq = self._latest_frame_seq if frame is not None: self._last_consumed_seq = seq return True, frame.copy() return False, None return self._read_frame_raw() def _read_frame_raw(self): if self._use_pyav: if self._av_decoder is None: return False, None try: frame = next(self._av_decoder) except StopIteration: return False, None except Exception: return False, None return True, frame.to_ndarray(format="bgr24") if self.cap is None: return False, None return self.cap.read() def _start_capture_thread(self): if self._capture_thread is not None: return self._capture_running = True self._capture_thread = threading.Thread(target=self._capture_loop, daemon=True) self._capture_thread.start() if self.verbose: print("Latest-frame capture thread started.") def _capture_loop(self): while self._capture_running: ret, frame = self._read_frame_raw() if not ret: time.sleep(0.001) continue with self._latest_frame_lock: self._latest_frame = frame self._latest_frame_seq += 1 # ========================================================================= # Private - Utilities # ========================================================================= def _print_timings(self, name, timings): """Print timing breakdown.""" print(f"\n[{name} timing]") for k, v in timings.items(): if k != 'total': print(f" {k}: {v:7.2f} ms") print(f" ─────────────────────────") print(f" TOTAL: {timings['total']:7.2f} ms ({1000/timings['total']:.1f} FPS)")