""" Stch Calibration OK """ import os import numpy as np from scipy import ndimage import cv2 from .Stitcher import Stitcher from .tracker import Tracker from .nHHD import * import time from scipy.interpolate import griddata from enum import Enum, auto from typing import Optional, Tuple, Union, Any from .getDensity import getFilteredDensity from . import configs class Sensor: """ Sensor helper that can decompose optical flow via Fourier methods. Stores the decomposed components and global norms as attributes after compute(). """ def __init__(self, vid_src, isstitch=False, verbose=True, cuda=True, config_name="prototype"): """ Initialize Sensor with video source and configuration. Args: vid_src: Video source (int for camera ID, str for file path, None for default camera) isstitch: Whether to use stitching mode for multi-tile sensor (default: False) verbose: Whether to print initialization logs (default: True) cuda: Whether to use CUDA acceleration for optical flow & decomposition (default: True) """ try: import cupy as cp import cupyx.scipy.signal as cp_signal self.CUPY_AVAILABLE = True except ImportError: self.CUPY_AVAILABLE = False # Initialize settings if vid_src is not None: vid_src = vid_src else: print("User must specify a video source.") self.vid_src, self.isstitch, self.cuda, self.verbose = vid_src, isstitch, cuda, verbose self.config_name = config_name # Initialize Sensor Configuration self.get_config() # Initialize components self.__stitcher = Stitcher(isstitch=isstitch, index=self.config_index) self.tracker = Tracker(adaptive=True, cuda=cuda) self.hhd = None if self.verbose: print("Components initialized") # Setup video capture self.cap = cv2.VideoCapture(vid_src) 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]) # Wait for first valid frame if self.verbose: print("Sensor connecting to video source...") frame = None while frame is None: ret, frame = self.cap.read() if not ret: time.sleep(1) if self.verbose: print("Video source connected.") # Initialize and calibrate self.__find_stitch_boxes(frame, show=False) self.__get_tiles(frame) self.stitch_cal() if self.verbose: print("Sensor initialized and calibrated.") # Initialize state variables self.stitch_valid = False self.n_boxes = None self.stitch_calibration = [] self.img = self.flow = self.normal = self.shear = self.force = self.depth = None self.img_filtered = None # For temporal noise reduction self.fnormal=self.fshearx=self.fsheary=self.ftor=0.0 self.imgupdate = self.vupdate = False # Logging state self.last_log_time = time.time() self.fps = 0 if self.verbose: print("Settings:") print(f" Video source: {vid_src}") print(f" Stitch mode: {isstitch}") print(f" CUDA acceleration: {cuda}") print(f" Cupy available: {self.CUPY_AVAILABLE}") print(f" OpenCV CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}") class info(Enum): VRAW = auto() VNORMAL = auto() VSHEAR = auto() FPS = auto() TIMESTAMP = auto() FNORMAL = auto() FSHEARX = auto() FSHEARY = auto() DEPTH = auto() def get_config(self): """ Load and return sensor configuration from JSON file. Returns: dict: Dictionary containing sensor configuration """ self.config = configs.load_config(self.config_name) self.config_resolution_raw = self.config["camera"]["resolution_raw"] self.config_fps = self.config["camera"]["fps"] self.config_units = self.config["camera"]["units"] self.config_resolution_stitch = self.config["camera"]["resolution_stitch"] self.config_index = self.config["camera"]["index"] self.config_pos_default = self.config["camera"]["pos_default"] def get_img(self): """ Capture frame from video source, convert to grayscale, and perform stitching. Returns: np.ndarray: Stitched grayscale image (contiguous array) or None if capture fails """ ret, frame = self.cap.read() if not ret: print("Failed to get_img from video source.") self.cap.release() exit(1) return None gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY) self.__stitcher.get_stitch_unit(gray_frame) self.__get_tiles(gray_frame) img = np.ascontiguousarray(self.__stitcher.stitch_rt_optimized()) img_resized = cv2.resize(img, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC) # Apply temporal noise reduction using exponential moving average # Blend current frame with previous filtered frame to reduce temporal noise alpha = 0.7 # Weight for current frame (0.7 = 70% current, 30% previous) if self.img_filtered is None: # First frame: no previous frame to blend with self.img_filtered = img_resized.astype(np.float32) else: # Temporal filtering: blend current frame with previous filtered frame self.img_filtered = alpha * img_resized.astype(np.float32) + (1 - alpha) * self.img_filtered # Convert back to uint8 self.img = self.img_filtered.astype(np.uint8) self.imgupdate = True return self.img def compute_deformation(self): """ Track optical flow and decompose into force field. Returns: tuple: (flow, force) where: - flow: Optical flow field (2D vector field) - force: Decomposed force components from nHHD， ([normal, shear_x, shear_y, torque]). Note: torque is experimental. """ flow = self.tracker.track(self.img) if hasattr(flow, 'download'): flow = flow.download() if self.hhd is None: self.hhd = nHHD(grid=(flow.shape[0], flow.shape[1]), spacings=(18, 18)) self.hhd.decompose(flow) self.imgupdate = False self.vupdate = True self.flow = flow self.normal = self.hhd.get_normal() self.shear = self.hhd.get_shear() flag, [self.fnormal,self.fshearx,self.fsheary,self.ftor] = self.hhd.get_force(flow) self.depth = getFilteredDensity(self.flow, use_cuda=self.cuda) self.fps = 1/(time.time() - self.last_log_time) self.last_log_time = time.time() return self.flow def read_info(self, *args: info) -> Tuple: """ Get sensor data according to specified info types. Args: *args: Variable number of info enum values specifying desired data Returns: Tuple: Data in the order specified by args. Each element corresponds to: - VRAW: Raw optical flow field (np.ndarray) - VNORMAL: Normal component of flow (np.ndarray) - VSHEAR: Shear component of flow (np.ndarray) - FNORMAL: Normal force magnitude (float) - FSHEAR: Shear force magnitude (float) - FPS: Current frames per second (float) - TIMESTAMP: Current timestamp string (str) - DEPTH: Depth image (np.ndarray) Example: normal, shear, fps = sensor.read_info(sensor.info.VNORMAL, sensor.info.VSHEAR, sensor.info.FPS) Note: Must call compute_deformation() before reading flow/force data """ if not self.vupdate: print("Warning: Update field first via compute_deformation() before reading flow/force info.") results = [] for arg in args: if not isinstance(arg, self.info): print(f"Warning: Invalid argument type {type(arg)}, skipping") continue if arg == self.info.VRAW: results.append(self.flow if self.flow is not None else np.array([])) elif arg == self.info.VNORMAL: results.append(self.normal if self.normal is not None else np.array([])) elif arg == self.info.VSHEAR: results.append(self.shear if self.shear is not None else np.array([])) elif arg == self.info.FNORMAL: results.append(self.fnormal) elif arg == self.info.FSHEARX: results.append(self.fshearx) elif arg == self.info.FSHEARY: results.append(self.fsheary) elif arg == self.info.FPS: results.append(self.fps) elif arg == self.info.TIMESTAMP: results.append(time.strftime('%H:%M:%S')) elif arg == self.info.DEPTH: results.append(self.depth if self.depth is not None else np.array([])) return tuple(results) if results else () def reset(self): self.tracker.reset() def disconnect(self): """ Release video capture resources. """ if self.cap.isOpened(): self.cap.release() if self.verbose: print("Sensor disconnected.") cv2.destroyAllWindows() def __get_tiles(self, img_src): """ Extract individual tiles from source image for stitching. Args: img_src: Source image (grayscale) containing multiple vision tiles Returns: bool: True (always succeeds) """ self.__stitcher.get_stitch_unit(img_src) self.__stitcher.extract_tiles() return True def __find_stitch_boxes(self, img_src, show=False): """ Detect vision tile boxes in the frame, waiting for camera stabilization. Args: img_src: Initial source image (not actively used, boxes detected from live capture) show: Whether to display detected boxes visually (default: False) Returns: int: Number of boxes detected (should be 15 for valid stitching) """ n_boxes, v_mean0, timeout = 0, 0, 10 start_time = cv2.getTickCount() # print("Waiting for camera stabilization...") 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 = self.__stitcher.find_stitch_boxes(frame, show=show) if (cv2.getTickCount() - start_time) / cv2.getTickFrequency() > timeout: print("Timeout for finding valid vision units.") self.isstitch = False self.__stitcher.isstitch = False break v_mean0 = v_mean # print("Camera stabilized.") # print("Checking validabity for stitching...") self.n_boxes = n_boxes self.stitch_valid = self.isstitch and n_boxes == self.config_units if self.isstitch: print("Sensor is valid for stithcing." if self.stitch_valid else "Sensor not valid to stitching, switching to non-stitch mode") else: print("Using non-stitch mode.") return n_boxes def stitch_cal(self): """ Load or create stitch calibration for aligning multiple tiles. Returns: np.ndarray: Stitched image after calibration Note: - Loads calibration from ./config/cal.npy if exists - Creates new calibration and saves to ./config/cal.npy if not found - Calibration contains transformation parameters for tile alignment """ cal_path = "./config/cal.npy" if os.path.exists(cal_path): print("Loading existing calibration from", cal_path) self.stitch_calibration = np.load(cal_path, allow_pickle=True) self.__stitcher.stitch_calibration = self.stitch_calibration img_stitched_height = self.__stitcher.stitch_rt_optimized() else: print("Calibration file not found, creating new calibration...") img_stitched_height, self.stitch_calibration = self.__stitcher.stitch_cal() os.makedirs("./config", exist_ok=True) np.save(cal_path, self.stitch_calibration) print("Calibration saved to", cal_path) # print("Image stitched with shape:", img_stitched_height.shape) return img_stitched_height def draw_arrows(self, info_type: info,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0) -> np.ndarray: """ Draw colored optical flow arrows on sensor image based on flow magnitude. Args: threshold: Minimum flow magnitude to display arrow (default: 10.0) grid_spacing: Spacing between arrows in pixels (default: 20) arrow_scale: Scale factor for arrow length (default: 3.0) Returns: np.ndarray: 400x400 BGR image with colored arrows (JET colormap) """ # Create working copies to avoid modifying originals image = self.img.copy() if info_type is None: print("Flow field type not specified, defaulting to show raw vector field.") optical_flow = self.flow.copy() elif info_type == self.info.VRAW: optical_flow = self.flow.copy() elif info_type == self.info.VNORMAL: optical_flow = self.normal.copy() elif info_type == self.info.VSHEAR: optical_flow = self.shear.copy() # Convert grayscale to BGR to enable colored arrow drawing if len(image.shape) == 2 or image.shape[2] == 1: image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) # Calculate scaling factors between image and flow dimensions img_height, img_width = image.shape[:2] flow_height, flow_width = optical_flow.shape[:2] scale_x, scale_y = img_width / flow_width, img_height / flow_height # Create sampling grid for flow field (used for interpolation source points) sample_step = 19 # Spacing between sample points flow_start = np.stack(np.meshgrid(range(0, flow_width, sample_step), range(0, flow_height, sample_step)), 2) flow_end = (optical_flow[flow_start[:, :, 1], flow_start[:, :, 0], :] * 2.0 + flow_start).astype(np.int32) # Create regular grid for arrow placement in image coordinates grid_x, grid_y = np.mgrid[grid_spacing//2:flow_height:grid_spacing, grid_spacing//2:flow_width:grid_spacing] x_axis = (grid_y.reshape(-1) * scale_x).astype(np.int64) # Scale to image width y_axis = (grid_x.reshape(-1) * scale_y).astype(np.int64) # Scale to image height # Interpolate flow vectors from sample grid to regular arrow grid using cubic interpolation grid_vector = griddata( flow_start.reshape(-1, 2), (flow_end - flow_start).reshape(-1, 2), # Flow from start to end (original to new position) (grid_y, grid_x), method='cubic', fill_value=0 ) grid_vector[np.isnan(grid_vector)] = 0 # Replace NaN with zero for stability grid_flat_scaled = grid_vector.reshape(-1, 2) * np.array([scale_x, scale_y]) # Compute arrow endpoints and ensure they stay within image bounds end_x = (grid_flat_scaled[:, 0] / arrow_scale + x_axis).astype(np.int64) end_y = (grid_flat_scaled[:, 1] / arrow_scale + y_axis).astype(np.int64) valid_mask = (y_axis >= 0) & (y_axis < img_height) & (x_axis >= 0) & (x_axis < img_width) valid_mask &= (end_x >= 0) & (end_x < img_width) & (end_y >= 0) & (end_y < img_height) # Filter arrows by magnitude threshold and compute colors using JET colormap magnitudes = np.linalg.norm(grid_flat_scaled, axis=1) valid_mask &= (magnitudes >= 0) # Only show arrows above threshold valid_indices = np.where(valid_mask)[0] mag_range = magnitudes.max() - magnitudes.min() # Normalize magnitudes to 0-255 range for colormap application, but set below-threshold values to 0 (blue) norm_magnitudes = ((magnitudes - magnitudes.min()) / mag_range * 255.0).astype(np.uint8) if mag_range > 0 else np.zeros_like(magnitudes, dtype=np.uint8) norm_magnitudes[magnitudes < threshold] = 0 # Keep lowest color (blue) for below-threshold magnitudes colors_array = cv2.applyColorMap(norm_magnitudes.reshape(-1, 1), cv2.COLORMAP_JET) # JET: blue (low) -> red (high) # Draw arrows at each valid grid point with magnitude-based coloring for i in valid_indices: color = tuple(map(int, colors_array[i, 0])) # Scale arrow length and thickness based on threshold if magnitudes[i] < threshold: # Below threshold: draw as small dot (zero-length arrow with larger tip) cv2.circle(image, (x_axis[i], y_axis[i]), radius=1, color=color, thickness=-1) else: # Above threshold: draw full arrow with thickness proportional to magnitude # Arrow rooted at original position (x_axis, y_axis), head points to new position (end_x, end_y) thickness = max(1, int(2 * magnitudes[i] / magnitudes.max())) if magnitudes.max() > 0 else 2 cv2.arrowedLine(image, (x_axis[i], y_axis[i]), (end_x[i], end_y[i]), color, thickness=thickness, tipLength=0.3) out_img = cv2.resize(image, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC) # if fps: # out_img = self.draw_fps(image=out_img) return out_img def draw_magnitude_map(self, info_type: info = None, threshold=3.0, colormap=cv2.COLORMAP_JET) -> np.ndarray: """ Draw color map of pointwise magnitude of velocity field. Args: info_type: Type of flow field to visualize (VRAW, VNORMAL, VSHEAR, or None for total) threshold: Minimum magnitude to display (below threshold shows as zero/blue) colormap: OpenCV colormap to use (default: COLORMAP_JET) Returns: np.ndarray: 400x400 BGR image with colored magnitude map """ # Get the appropriate flow field if info_type is None or info_type == self.info.VRAW: flow_field = self.flow.copy() if self.flow is not None else None elif info_type == self.info.VNORMAL: flow_field = self.normal.copy() if self.normal is not None else None elif info_type == self.info.VSHEAR: flow_field = self.shear.copy() if self.shear is not None else None else: # For combined normal + shear magnitude if self.normal is not None and self.shear is not None: flow_field = self.normal + self.shear else: flow_field = None if flow_field is None: # Return blank image if no flow data return np.zeros((self.config_resolution_stitch[0], self.config_resolution_stitch[1], 3), dtype=np.uint8) # Compute magnitude at each point magnitude = np.sqrt(flow_field[:, :, 0]**2 + flow_field[:, :, 1]**2) # Smooth by downsampling and interpolating back downsample_factor = 10 h, w = magnitude.shape magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor] magnitude_smoothed = cv2.resize(magnitude_downsampled, (w, h), interpolation=cv2.INTER_CUBIC) # Apply threshold - set below-threshold values to zero magnitude_thresholded = magnitude_smoothed.copy() magnitude_thresholded[magnitude_smoothed < threshold] = 0 # Normalize to 0-255 range mag_min, mag_max = magnitude_thresholded.min(), magnitude_thresholded.max() if mag_max > mag_min: magnitude_norm = ((magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255).astype(np.uint8) else: magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8) # Apply colormap magnitude_colored = cv2.applyColorMap(magnitude_norm, colormap) # Resize to output size out_img = cv2.resize(magnitude_colored, (self.config_resolution_stitch[0], self.config_resolution_stitch[1]), interpolation=cv2.INTER_CUBIC) return out_img def draw_fps(self, image=None): """ Draw FPS counter on the top left corner of the image. """ font_scale=0.5 thickness=2 color=(0, 255, 0) image = image.copy() fps = self.fps # Draw FPS text fps_text = f"FPS: {fps:.1f}" position = (10, 30) # Top left corner with margin font = cv2.FONT_HERSHEY_SIMPLEX # Draw FPS text cv2.putText(image, fps_text, position, font, font_scale, color, thickness, cv2.LINE_AA) return image