orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

create util.py

Browse Source
This commit is contained in:
bluepulsar
2025-12-25 20:46:13 +08:00
parent f68893e16e
commit 231f070143
6 changed files with 212 additions and 243 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/orisys/sensor.py
+50 -219
View File
@@ -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)