From 231f070143fc70e3f62610ae3e196c64895a11a4 Mon Sep 17 00:00:00 2001 From: bluepulsar Date: Thu, 25 Dec 2025 20:46:13 +0800 Subject: [PATCH] create util.py --- example1_basic.py | 9 +- example2_config.py | 12 +- src/orisys/__init__.py | 2 +- src/orisys/configs/prototype.json | 10 -- src/orisys/sensor.py | 269 ++++++------------------------ src/orisys/util.py | 153 +++++++++++++++++ 6 files changed, 212 insertions(+), 243 deletions(-) create mode 100644 src/orisys/util.py diff --git a/example1_basic.py b/example1_basic.py index 7ae83d9..4a44c3f 100644 --- a/example1_basic.py +++ b/example1_basic.py @@ -1,15 +1,13 @@ import os os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" import cv2 -from orisys import Sensor -import numpy as np - +import orisys def main(): global start_time # 打开视频 - sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True) + sensor1 = orisys.Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True) # 打开摄像头 # sensor1 = Sensor(0,isstitch=True,cuda = True,verbose = True) @@ -24,7 +22,7 @@ def main(): fps, fn, fx, fy, flow, depth = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW, sensor1.info.DEPTH) print(fps,fn,fx,fy) # 显示向量场 - arrows = sensor1.draw_arrows(sensor1.info.VRAW, threshold=5, grid_spacing=40, arrow_scale=1.0) + arrows = orisys.util.draw_arrows(sensor1.img, flow, threshold=5, grid_spacing=40, arrow_scale=1.0) cv2.imshow("Deformation", arrows) depth_img = cv2.applyColorMap((depth*1* 255.0).astype('uint8'), cv2.COLORMAP_HOT) @@ -41,6 +39,5 @@ def main(): sensor1.disconnect() # 断开传感器连接 - if __name__ == '__main__': main() diff --git a/example2_config.py b/example2_config.py index 6b24558..bf40cab 100644 --- a/example2_config.py +++ b/example2_config.py @@ -1,15 +1,14 @@ import os os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" import cv2 -from orisys import Sensor -import numpy as np - +import orisys def main(): global start_time # 打开视频 - sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True, config_name="prototype") + sensor1 = orisys.Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True, config_name="prototype") + # sensor1 = Sensor(0,isstitch=True,cuda = True,verbose = True, config_name="prototype") # 打开摄像头 # sensor1 = Sensor(0,isstitch=True,cuda = True,verbose = True) @@ -22,9 +21,9 @@ def main(): # 变形量总和: FNORMAL , FSHEARX, FSHEARY : 法向，切向（X)，切向（Y) # 深度图: DEPTH fps, fn, fx, fy, flow, depth = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW, sensor1.info.DEPTH) - print(fps,fn,fx,fy) + # 显示向量场 - arrows = sensor1.draw_arrows(sensor1.info.VRAW, threshold=5, grid_spacing=40, arrow_scale=1.0) + arrows = orisys.util.draw_arrows(sensor1.img, flow, threshold=5, grid_spacing=40, arrow_scale=1.0) cv2.imshow("Deformation", arrows) depth_img = cv2.applyColorMap((depth*1* 255.0).astype('uint8'), cv2.COLORMAP_HOT) @@ -41,6 +40,5 @@ def main(): sensor1.disconnect() # 断开传感器连接 - if __name__ == '__main__': main() diff --git a/src/orisys/__init__.py b/src/orisys/__init__.py index 4f305f3..eadd478 100644 --- a/src/orisys/__init__.py +++ b/src/orisys/__init__.py @@ -5,7 +5,7 @@ A Python library for processing and analyzing optical tactile sensor data, including image stitching, optical flow tracking, and deformation decomposition. """ -__version__ = "0.1.0" +__version__ = "0.2.0" __author__ = "Orisys" from .sensor import Sensor diff --git a/src/orisys/configs/prototype.json b/src/orisys/configs/prototype.json index 0094f99..1d6c253 100644 --- a/src/orisys/configs/prototype.json +++ b/src/orisys/configs/prototype.json @@ -9,16 +9,6 @@ "resolution_stitch": [400, 400], "index": [1,2,3,6,7,8,11,12,13], "pos_default": [[260, 56], [577, 62], [897, 69], [1214, 75], [1530, 81], [256, 391], [570, 397], [889, 403], [1208, 409], [1523, 416], [248, 724], [563, 730], [882, 738], [1201, 743], [1516, 751]] - }, - "elastomer": { - "thickness_mm": 3.0, - "marker_pattern": "random_color", - "marker_version": "1.0" - }, - "mechanical": { - "width_mm": 10, - "height_mm": 10, - "depth_mm": 10 } } diff --git a/src/orisys/sensor.py b/src/orisys/sensor.py index 1554c76..ae9ccbf 100644 --- a/src/orisys/sensor.py +++ b/src/orisys/sensor.py @@ -1,6 +1,3 @@ -""" - Stch Calibration OK -""" import os import numpy as np import cv2 @@ -13,12 +10,9 @@ from enum import Enum, auto from typing import Optional, Tuple, Union, Any from .getDensity import getFilteredDensity from . import configs +from .util import draw_arrows, draw_magnitude_map, draw_fps 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. @@ -28,6 +22,7 @@ class Sensor: 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) + config_name: Name of the mdoel (default: "prototype") """ try: import cupy as cp @@ -38,9 +33,15 @@ class Sensor: # Initialize settings if vid_src is None: - raise ValueError("User must specify a video source.") + raise Exception("User must specify a video source.") self.vid_src, self.isstitch, self.cuda, self.verbose, self.config_name = vid_src, isstitch, cuda, verbose, config_name + + # Initialize state variables + self.stitch_valid = self.imgupdate = self.vupdate = False + self.stitch_calibration = [] + self.img = self.flow = self.normal = self.shear = self.force = self.depth = self.n_boxes = self.img_filtered = None + self.fnormal=self.fshearx=self.fsheary=self.ftor=0.0 # Initialize Sensor Configuration self.get_config() @@ -50,46 +51,16 @@ class Sensor: 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) - - # 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 = self.imgupdate = self.vupdate = False - self.stitch_calibration = [] - self.img = self.flow = self.normal = self.shear = self.force = self.depth = self.n_boxes = self.img_filtered = None - self.fnormal=self.fshearx=self.fsheary=self.ftor=0.0 + self.__open_video_source() # Logging state self.last_log_time = time.time() self.fps = 0 if self.verbose: - print("Settings:" {self.config_name}) print(f" Video source: {vid_src}") - print(f" Stitch mode: {isstitch}") + print(f" Stitch mode: {self.stitch_valid}") print(f" CUDA acceleration: {self.cuda}") print(f" Cupy available: {self.CUPY_AVAILABLE}") print(f" OpenCV CUDA devices: {cv2.cuda.getCudaEnabledDeviceCount()}") @@ -105,6 +76,34 @@ class Sensor: FSHEARY = auto() DEPTH = auto() + def __open_video_source(self): + + # Setup video capture + self.cap = cv2.VideoCapture(self.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 + if self.cap.isOpened(): + while frame is None: + ret, frame = self.cap.read() + if not ret: + time.sleep(1) + else: + raise Exception("Failed to open video source.") + + # Initialize stitching calibration + self.__find_stitch_boxes(frame, show=False) + self.__get_tiles(frame) + self.stitch_cal() + if self.verbose: + print("Sensor initialized and calibrated.") + def get_config(self): """ Load and return sensor configuration from JSON file. @@ -123,11 +122,10 @@ class Sensor: if self.verbose: print("Config loaded: ", str(self.config_name)) print(" Resolution raw: ", str(self.config_resolution_raw)) - print(" FPS: ", str(self.config_fps)) - print(" Units: ", str(self.config_units)) - print(" Resolution stitch: ", str(self.config_resolution_stitch)) - print(" Index: ", str(self.config_index)) - print(" Pos default: ", str(self.config_pos_default)) + print(" Camera FPS: ", str(self.config_fps)) + print(" Number of VisionUnits: ", str(self.config_units)) + print(" Resolution after stitch: ", str(self.config_resolution_stitch)) + print(" Index for stitching: ", str(self.config_index)) def get_img(self): """ @@ -165,8 +163,6 @@ class Sensor: self.imgupdate = True return self.img - - def compute_deformation(self): """ Track optical flow and decompose into force field. @@ -197,7 +193,6 @@ class Sensor: 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: @@ -291,7 +286,7 @@ class Sensor: Returns: int: Number of boxes detected (should be 15 for valid stitching) """ - n_boxes, v_mean0, timeout = 0, 0, 10 + n_boxes, v_mean0, timeout = 0, 0, 5 start_time = cv2.getTickCount() # print("Waiting for camera stabilization...") @@ -309,11 +304,8 @@ class Sensor: 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") @@ -353,172 +345,11 @@ class Sensor: # 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) + # def draw_arrows(self, optical_flow, fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0): + # return draw_arrows(self.img, optical_flow, fps, threshold, grid_spacing, arrow_scale) - 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 - + # def draw_magnitude_map(self, flow_field, threshold=3.0, colormap=cv2.COLORMAP_JET): + # return draw_magnitude_map(self.img, flow_field, threshold, colormap) - - + # def draw_fps(self, fps): + # return draw_fps(self.img, fps) diff --git a/src/orisys/util.py b/src/orisys/util.py new file mode 100644 index 0000000..4cb0ff8 --- /dev/null +++ b/src/orisys/util.py @@ -0,0 +1,153 @@ +import cv2 +import numpy as np +from scipy.interpolate import griddata + +def draw_arrows(image, optical_flow,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0): + """ + 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 + + # 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) + return image + +def draw_magnitude_map(image, flow_field, threshold=3.0, colormap=cv2.COLORMAP_JET): + """ + 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 magnitude_colored + +def draw_fps(image, fps): + """ + Draw FPS counter on the top left corner of the image. + """ + font_scale=0.5 + thickness=2 + color=(0, 255, 0) + image = image.copy() + + # 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 \ No newline at end of file