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

update promotion video gui script

Browse Source
This commit is contained in:
ykcheungab
2026-02-03 15:58:10 +08:00
parent 22de58986a
commit b8b180d848
8 changed files with 313 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/orisys/tracker.py
+6 -3
View File
@@ -1,9 +1,7 @@
import cv2
import numpy as np
import time
import torch
import torchvision.transforms as T
from torchvision.models.optical_flow import raft_large
class Tracker:
def __init__(self,
@@ -27,6 +25,9 @@ class Tracker:
dis_preset = config_tracker["dis_preset"]
loop_closure_interval = config_tracker["loop_closure_interval"]
self.__algorithm = config_tracker["algorithm"]
if self.__algorithm == "raft":
import torch
if self.__algorithm == "raft" and (not self.cuda or not torch.cuda.is_available()):
self.__algorithm = "dis"
@@ -199,6 +200,8 @@ class Tracker:
"""
Setup a simpler version using available models with performance optimizations
"""
import torchvision.transforms as T
from torchvision.models.optical_flow import raft_large
# Load model and move to GPU
model = raft_large(pretrained=True, progress=False).to(self.__raft_device)
model = model.eval()
src/orisys/util.py
+38 -65
View File
@@ -1,78 +1,51 @@
import cv2
import numpy as np
from scipy.interpolate import griddata
def draw_arrows(frame, optical_flow,fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0):
def draw_arrows(frame, optical_flow, fps=True, threshold=5, grid_spacing=40, arrow_scale=1.0, margin = 0.8, colormap=cv2.COLORMAP_JET, below_threshold_color=(0, 0, 255), edge_margin=15):
"""
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)
Draw colored optical flow arrows on the image by magnitude (JET colormap).
Frame and optical_flow must have the same height/width; arrows are drawn in pixel coordinates.
"""
# Create working copies to avoid modifying originals
image = frame.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
h, w = image.shape[:2]
# Arrow grid in pixel coordinates (assume frame and flow same size)
rr = np.arange(grid_spacing // 2, h, grid_spacing)
cc = np.arange(grid_spacing // 2, w, grid_spacing)
rr, cc = np.meshgrid(rr, cc, indexing="ij")
rr, cc = rr.reshape(-1), cc.reshape(-1)
# Displacement at each grid point (flow stored as half-displacement)
disp = 2.0 * optical_flow[rr, cc, :]
mag = np.linalg.norm(disp, axis=1)
dx = disp[:, 0] / arrow_scale
dy = disp[:, 1] / arrow_scale
length = np.sqrt(dx * dx + dy * dy)
# Cap arrow length at grid_spacing
scale = np.where(length > grid_spacing*margin, grid_spacing / np.maximum(length, 1e-9), 1.0)*margin
dx, dy = dx * scale, dy * scale
end_x = (cc + dx).astype(np.int32)
end_y = (rr + dy).astype(np.int32)
valid = (
(cc >= edge_margin) & (cc < w - edge_margin) & (rr >= edge_margin) & (rr < h - edge_margin)
& (end_x >= 0) & (end_x < w) & (end_y >= 0) & (end_y < h)
)
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)
mag_range = mag.max() - mag.min()
norm_mag = ((mag - mag.min()) / mag_range * 255).astype(np.uint8) if mag_range > 0 else np.zeros_like(mag, dtype=np.uint8)
norm_mag[mag < threshold] = 0
colors = cv2.applyColorMap(norm_mag.reshape(-1, 1), colormap)
for i in np.where(valid)[0]:
pt1, pt2 = (int(cc[i]), int(rr[i])), (int(end_x[i]), int(end_y[i]))
color = tuple(int(c) for c in colors[i, 0])
if mag[i] < threshold:
cv2.circle(image, pt1, radius=1, color=below_threshold_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)
thickness = max(1, int(2 * mag[i] / mag.max())) if mag.max() > 0 else 2
cv2.arrowedLine(image, pt1, pt2, color, thickness=thickness, tipLength=0.2)
return image