""" 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 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) 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, config_name="prototype", cal_path=None, override=False, box=None, mirror=True, poisson_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) """ 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.") # Store settings self.vid_src = vid_src self.isstitch = isstitch self.cuda = cuda self.verbose = verbose 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 shape = (self.__config_cam["resolution_stitch"][1], self.__config_cam["resolution_stitch"][0]) self.img_size_raw = None self.img_filtered = None # Optical Flow Data Variables (public) self.flow = np.zeros(shape + (2,), dtype=np.float32) self.vnormal = np.zeros_like(self.flow) self.vshear = np.zeros_like(self.flow) 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 # Initialize components self.stitcher = Stitcher(isstitch=isstitch, config_stitcher=self.__config_stitcher, verbose=self.verbose) self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, dis_preset=self._config_dis_preset, config_tracker=self.__config_tracker) self.__contact3d = Contact3D(config_contact=self.__config_contact, verbose=self.verbose) self.__hhd = None self.poisson_backend = poisson_backend # Setup video capture 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): """ Capture and preprocess a frame. Returns: np.ndarray: Processed grayscale image """ # 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() ret, frame = self.cap.read() 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() self.frame = frame # Extract and stitch tiles # self.stitcher.get_stitch_unit(frame) self._extract_tiles(frame) img = np.ascontiguousarray(self.stitcher.stitch_rt_optimized()) self.img_size_raw = img.shape # Resize img = cv2.resize(img, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_LINEAR) if self.mirror: img = cv2.flip(img, 1) self.img = img.astype(np.uint8) t_processing_end = time.perf_counter() self._last_processing_time = t_processing_end - t_processing_start # FPS calculation self._fps = 1.0 / (time.perf_counter() - self.last_log_time) self.last_log_time = time.perf_counter() return None def detect_motion(self, threshold=4): if self.initial_img is None: return False diff = cv2.absdiff(self.img, self.initial_img) return diff.mean() > threshold def compute_deformation(self, decompose=True): """ Compute optical flow and decompose into force components. Returns: np.ndarray: Optical flow field """ if self.img is None: raise Exception("Warning: Get image first") self.compute_flow() if decompose: self.compute_hhd() return True 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") # Initialize HHD solver if self.__hhd is None: self.__hhd = nHHD(cuda=self.cuda, grid=flow.shape[:2], spacings=(18, 18), poisson_backend=self.poisson_backend) # Force computation # When OpenCL backend is selected, use host force path to ensure # Poisson solve goes through PoissonSolver.solve_opencl(). _, forces = self.__hhd.get_force(flow) # if self.poisson_backend == 'opencl': # _, forces = self.__hhd.get_force_host(flow) # elif self.cuda and self.CUPY_AVAILABLE: # _, forces = self.__hhd.get_force_gpu_v2(flow) # else: # _, forces = self.__hhd.get_force_host(flow) self.fnormal, self.fshearx, self.fsheary, self.ftor = forces # Get components self.vnormal = self.__hhd.get_normal() self.vshear = self.__hhd.get_shear() # Update state 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.img.copy() if self.img 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() self.initial_img = self.img.copy() def disconnect(self): """Release resources and stop async threads.""" if 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.__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["units"] 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["stitcher"] self.__config_contact = self._config["contact"] if self.verbose: print(f"Config: {self.config_name}") return True def _open_video_source(self, override, box): """Initialize video capture and calibration.""" self.cap = cv2.VideoCapture(self.vid_src) # Set buffer size to 1 to reduce latency and avoid frame accumulation # Note: Not all backends support this, but it helps when available try: self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1) except: pass # Some backends don't support this property 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.cap.isOpened(): raise Exception("Failed to open video source.") # Wait for first valid frame frame = None while frame is None: ret, frame = self.cap.read() if not ret: time.sleep(0.1) if 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) self._extract_tiles(frame) img = self.stitch_cal() img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) self.resolution_stitch = img.shape self.initial_img = cv2.resize(img, tuple(self._config_resolution_stitch), interpolation=cv2.INTER_CUBIC) if self.mirror: self.initial_img = cv2.flip(self.initial_img, 1) self.initial_img = self.initial_img.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.cap.read() 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 # ========================================================================= # 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)")