create util.py
This commit is contained in:
bluepulsar
+3
-6
@@ -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()
|
||||
|
||||
+5
-7
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+50
-219
@@ -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,10 +33,16 @@ 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,12 +304,9 @@ 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")
|
||||
else:
|
||||
@@ -353,172 +345,11 @@ class Sensor:
|
||||
# print("Image stitched with shape:", img_stitched_height.shape)
|
||||
return img_stitched_height
|
||||
|
||||
# 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)
|
||||
|
||||
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
|
||||
|
||||
|
||||
|
||||
# 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)
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user