update promotion video gui script
This commit is contained in:
+1
-1
@@ -12,7 +12,7 @@ def main():
|
|||||||
# cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件
|
# cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件
|
||||||
|
|
||||||
# 方式1: 从视频文件读取
|
# 方式1: 从视频文件读取
|
||||||
sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_0.npy", isstitch=True, cuda=True, verbose=True)
|
sensor = orisys.Sensor("s1.mp4", config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_0.npy", isstitch=True, cuda=True, verbose=True)
|
||||||
|
|
||||||
# 方式2: 从摄像头读取(取消注释使用)
|
# 方式2: 从摄像头读取(取消注释使用)
|
||||||
# sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
|
# sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
|
||||||
|
|||||||
Binary file not shown.
|
Before Width: | Height: | Size: 1.9 MiB After Width: | Height: | Size: 2.0 MiB |
@@ -6,7 +6,7 @@
|
|||||||
"fps": 30,
|
"fps": 30,
|
||||||
"units": 15,
|
"units": 15,
|
||||||
"units_used": 9,
|
"units_used": 9,
|
||||||
"resolution_stitch": [128, 128],
|
"resolution_stitch": [384, 384],
|
||||||
"stitch_res_list": [384,256,128,64]
|
"stitch_res_list": [384,256,128,64]
|
||||||
},
|
},
|
||||||
"stitcher": {
|
"stitcher": {
|
||||||
|
|||||||
@@ -6,7 +6,7 @@
|
|||||||
"fps": 30,
|
"fps": 30,
|
||||||
"units": 15,
|
"units": 15,
|
||||||
"units_used": 9,
|
"units_used": 9,
|
||||||
"resolution_stitch": [128, 128]
|
"resolution_stitch": [384, 384]
|
||||||
},
|
},
|
||||||
"stitcher": {
|
"stitcher": {
|
||||||
"min_area_frac": 0.01,
|
"min_area_frac": 0.01,
|
||||||
|
|||||||
@@ -57,8 +57,8 @@ def main():
|
|||||||
arrows = orisys.util.draw_arrows(
|
arrows = orisys.util.draw_arrows(
|
||||||
img,
|
img,
|
||||||
flow,
|
flow,
|
||||||
threshold=4,
|
threshold=2,
|
||||||
grid_spacing=40,
|
grid_spacing=20,
|
||||||
arrow_scale=1.0
|
arrow_scale=1.0
|
||||||
)
|
)
|
||||||
cv2.imshow("Deformation", arrows)
|
cv2.imshow("Deformation", arrows)
|
||||||
|
|||||||
@@ -0,0 +1,264 @@
|
|||||||
|
import os
|
||||||
|
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
||||||
|
import cv2
|
||||||
|
import orisys
|
||||||
|
import numpy as np
|
||||||
|
import time
|
||||||
|
import argparse
|
||||||
|
import pyqtgraph as pg
|
||||||
|
from PyQt5.QtGui import QImage, QFont
|
||||||
|
from pyqtgraph.Qt import QtGui
|
||||||
|
|
||||||
|
# 实时图表
|
||||||
|
bg_color = (46, 41, 36)
|
||||||
|
bg_color_hex = "#24292E"
|
||||||
|
text_color_hex = "#D1D5DA"
|
||||||
|
text_color = (209, 213, 218)
|
||||||
|
|
||||||
|
def make_gradient_brush(color_hex, max_alpha=150):
|
||||||
|
"""
|
||||||
|
创建一个从上到下渐变的笔刷 (兼容 PyQt5 和 PyQt6)
|
||||||
|
"""
|
||||||
|
c_top = QtGui.QColor(color_hex)
|
||||||
|
c_top.setAlpha(max_alpha) # 顶部半透明
|
||||||
|
|
||||||
|
c_bot = QtGui.QColor(color_hex)
|
||||||
|
c_bot.setAlpha(0) # 底部全透明
|
||||||
|
|
||||||
|
# 创建线性渐变:(0,0)到(0,1) 代表从控件顶部到底部
|
||||||
|
grad = QtGui.QLinearGradient(0, 0, 0, 1)
|
||||||
|
|
||||||
|
if hasattr(QtGui.QGradient, 'CoordinateMode'):
|
||||||
|
grad.setCoordinateMode(QtGui.QGradient.CoordinateMode.ObjectBoundingMode)
|
||||||
|
else:
|
||||||
|
grad.setCoordinateMode(QtGui.QGradient.ObjectBoundingMode)
|
||||||
|
# ==============================================
|
||||||
|
|
||||||
|
grad.setColorAt(0, c_top) # 0.0 (顶端) 颜色
|
||||||
|
grad.setColorAt(1, c_bot) # 1.0 (底端) 颜色
|
||||||
|
|
||||||
|
return QtGui.QBrush(grad)
|
||||||
|
|
||||||
|
def create_plot():
|
||||||
|
|
||||||
|
pg.setConfigOption('background', bg_color_hex) # 深蓝灰背景
|
||||||
|
pg.setConfigOption('foreground', text_color_hex) # 浅灰文字
|
||||||
|
pg.setConfigOptions(antialias=True) # 抗锯齿
|
||||||
|
# theme_color = '#00E5FF'
|
||||||
|
theme_color = '#50abff'
|
||||||
|
line_pen = pg.mkPen(color=theme_color, width=2)
|
||||||
|
grad_brush = make_gradient_brush(theme_color, max_alpha=120)
|
||||||
|
show_values = False
|
||||||
|
|
||||||
|
title_font = QFont("Segoe UI", 14, QFont.Bold)
|
||||||
|
axis_font = QFont("Segoe UI", 10, QFont.Bold)
|
||||||
|
p1 = win.addPlot(title="Normal")
|
||||||
|
p1.titleLabel.item.setFont(title_font)
|
||||||
|
p1.getAxis('left').setStyle(showValues=show_values)
|
||||||
|
p1.showGrid(x=True, y=True, alpha=0.3)
|
||||||
|
win.nextRow()
|
||||||
|
p2 = win.addPlot(title="Shear X")
|
||||||
|
p2.titleLabel.item.setFont(title_font)
|
||||||
|
p2.getAxis('left').setStyle(showValues=show_values)
|
||||||
|
p2.showGrid(x=True, y=True, alpha=0.3)
|
||||||
|
win.nextRow()
|
||||||
|
p3 = win.addPlot(title="Shear Y")
|
||||||
|
p3.titleLabel.item.setFont(title_font)
|
||||||
|
p3.getAxis('left').setStyle(showValues=show_values)
|
||||||
|
p3.setLabel('bottom', 'Time', units='s')
|
||||||
|
p3.getAxis('bottom').label.setFont(axis_font)
|
||||||
|
p3.showGrid(x=True, y=True, alpha=0.3)
|
||||||
|
plot_af = p1.plot(pen=line_pen, fillLevel=0, fillBrush=grad_brush)
|
||||||
|
plot_tfx = p2.plot(pen=line_pen, fillLevel=0, fillBrush=grad_brush)
|
||||||
|
plot_tfy = p3.plot(pen=line_pen, fillLevel=0, fillBrush=grad_brush)
|
||||||
|
return p1, p2, p3, plot_af, plot_tfx, plot_tfy
|
||||||
|
|
||||||
|
def update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3):
|
||||||
|
# global plot_af, plot_tfx, plot_tfy, data, p1, p2, p3
|
||||||
|
plot_af.setData(data[:,:,0])
|
||||||
|
p1.setYRange(min(data[:,1,0]),max(data[:,1,0]))
|
||||||
|
plot_tfx.setData(data[:,:,1])
|
||||||
|
p2.setYRange(min(data[:,1,1]),max(data[:,1,1]))
|
||||||
|
plot_tfy.setData(data[:,:,2])
|
||||||
|
p3.setYRange(min(data[:,1,2]),max(data[:,1,2]))
|
||||||
|
|
||||||
|
# 深度图
|
||||||
|
def draw_magnitude_map(vfield, threshold=3.0, colormap=cv2.COLORMAP_JET) -> np.ndarray:
|
||||||
|
"""
|
||||||
|
Draw color map of pointwise magnitude of velocity field.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
vfield: Velocity field to visualize (2D vector field)
|
||||||
|
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
|
||||||
|
"""
|
||||||
|
flow_field = vfield
|
||||||
|
|
||||||
|
# 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, (400, 400), interpolation=cv2.INTER_CUBIC)
|
||||||
|
|
||||||
|
return out_img
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
npoints = 100 # 图表数据长度
|
||||||
|
plot_threshold = 0 # only plot force values above this (below → 0)
|
||||||
|
pg.setConfigOption('background', bg_color_hex)
|
||||||
|
pg.setConfigOption('foreground', bg_color_hex)
|
||||||
|
app = pg.mkQApp("Plotting Example")
|
||||||
|
win = pg.GraphicsLayoutWidget(show=True, title="Basic plotting examples")
|
||||||
|
win.setWindowTitle('pyqtgraph example: Plotting')
|
||||||
|
p1, p2, p3, plot_af, plot_tfx, plot_tfy = create_plot()
|
||||||
|
data = np.zeros([npoints,2,3])
|
||||||
|
start_time = None # Will be set when sensor starts
|
||||||
|
for i in range(3):
|
||||||
|
data[:,0,i] = np.linspace(-npoints,0,npoints)/30
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
global start_time
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
parser.add_argument("--video","-v", type=str, default="s1.mp4", help="视频文件路径")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
arrow_writer = cv2.VideoWriter(args.video.replace(".mp4", "_arrows.mp4"), cv2.VideoWriter_fourcc(*'mp4v'), 30, (384, 384))
|
||||||
|
pyqt_writer = cv2.VideoWriter(args.video.replace(".mp4", "_pyqt.mp4"), cv2.VideoWriter_fourcc(*'mp4v'), 30, (1600, 1200))
|
||||||
|
depth_writer = cv2.VideoWriter(args.video.replace(".mp4", "_depth.mp4"), cv2.VideoWriter_fourcc(*'mp4v'), 30, (384, 384))
|
||||||
|
|
||||||
|
sensor1 = orisys.Sensor(args.video, config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_98.npy", isstitch=True, cuda=False, verbose=True)
|
||||||
|
|
||||||
|
start_time = time.time()
|
||||||
|
is_save = False
|
||||||
|
frame_count = 0
|
||||||
|
writer = None
|
||||||
|
while True:
|
||||||
|
sensor1.get_img() # 提取并拼接复眼图像
|
||||||
|
|
||||||
|
|
||||||
|
sensor1.compute_deformation() # 计算形变和分解
|
||||||
|
sensor1.compute_contact()
|
||||||
|
# 提取信息:
|
||||||
|
# 向量场: VRAW, VNORMAL, VSHEAR: 原始,法向,切向
|
||||||
|
# 帧率,时间戳: FPS, TIMESTAMP
|
||||||
|
# 变形量总和: FNORMAL , FSHEARX, FSHEARY : 法向,切向(X),切向(Y)
|
||||||
|
fps, fn, fx, fy, flow, depth_map = 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)
|
||||||
|
# 显示向量场
|
||||||
|
# Background color #24292E (BGR)
|
||||||
|
black_img = np.full_like(sensor1.img, bg_color)
|
||||||
|
arrows = orisys.util.draw_arrows(black_img, flow, threshold=2, grid_spacing=20, arrow_scale=0.5, colormap=cv2.COLORMAP_PARULA, below_threshold_color=text_color)
|
||||||
|
cv2.imshow("Deformation", arrows)
|
||||||
|
|
||||||
|
# 显示深度图
|
||||||
|
div_abs = depth_map
|
||||||
|
if div_abs.max() > div_abs.min():
|
||||||
|
div_normalized = ((div_abs - div_abs.min()) / (div_abs.max() - div_abs.min()) * 255).astype(np.uint8)
|
||||||
|
else:
|
||||||
|
div_normalized = np.zeros_like(div_abs, dtype=np.uint8)
|
||||||
|
|
||||||
|
# 应用颜色映射(使用JET:蓝色=低散度,红色=高散度)
|
||||||
|
depth_colored = cv2.applyColorMap(div_normalized, cv2.COLORMAP_PARULA)
|
||||||
|
# Set background color for zero/low values (BGR color #24292E)
|
||||||
|
depth_colored[div_normalized < 5] = bg_color
|
||||||
|
cv2.imshow("Depth Map", depth_colored)
|
||||||
|
# cv2.imshow("Depth Map", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET))
|
||||||
|
|
||||||
|
# 显示图表 - 更新时间轴
|
||||||
|
current_time = time.time() - start_time
|
||||||
|
for i in range(3):
|
||||||
|
data[:-1, 0, i] = data[1:, 0, i] # shift time data left
|
||||||
|
data[:-1, 1, i] = data[1:, 1, i] # shift force data left
|
||||||
|
data[-1, 0, :] = frame_count/30.0 # set current time for all plots
|
||||||
|
data[-1, 1, 0] = fn
|
||||||
|
data[-1, 1, 1] = abs(fx)
|
||||||
|
data[-1, 1, 2] = abs(fy)
|
||||||
|
# Moving average on force data (column 1): simple loop, no singularities at ends
|
||||||
|
ma_window = 4
|
||||||
|
data_plot = data.copy()
|
||||||
|
for i in range(3):
|
||||||
|
x = data_plot[:, 1, i]
|
||||||
|
n = len(x)
|
||||||
|
out = np.empty_like(x, dtype=np.float64)
|
||||||
|
for j in range(n):
|
||||||
|
start = max(0, j - ma_window + 1)
|
||||||
|
out[j] = np.mean(x[start : j + 1])
|
||||||
|
data_plot[:, 1, i] = out
|
||||||
|
below = data_plot[:, 1, :] < plot_threshold
|
||||||
|
data_plot[:, 1, :] = np.where(below, 0, data_plot[:, 1, :])
|
||||||
|
update(plot_af, plot_tfx, plot_tfy, data_plot, p1, p2, p3) # Update plot
|
||||||
|
frame_count += 1
|
||||||
|
|
||||||
|
# Convert pyqtgraph plot to image
|
||||||
|
# Grab the GraphicsLayoutWidget as QPixmap
|
||||||
|
pixmap = win.grab()
|
||||||
|
# Convert QPixmap to QImage
|
||||||
|
qimg = pixmap.toImage()
|
||||||
|
# Convert QImage to numpy array (RGB format)
|
||||||
|
qimg = qimg.convertToFormat(QImage.Format_RGB888)
|
||||||
|
width = qimg.width()
|
||||||
|
height = qimg.height()
|
||||||
|
# Convert QImage to numpy array
|
||||||
|
ptr = qimg.bits()
|
||||||
|
ptr.setsize(height * width * 3)
|
||||||
|
arr = np.array(ptr, dtype=np.uint8).reshape((height, width, 3))
|
||||||
|
# Convert RGB to BGR for OpenCV
|
||||||
|
pyqt_img_bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
|
||||||
|
# Resize to match video writer dimensions (1920, 1080)
|
||||||
|
pyqt_img_resized = cv2.resize(pyqt_img_bgr, (1600, 1200), interpolation=cv2.INTER_LINEAR)
|
||||||
|
|
||||||
|
arrow_writer.write(arrows)
|
||||||
|
pyqt_writer.write(pyqt_img_resized)
|
||||||
|
depth_frame = cv2.resize(depth_colored, (384, 384), interpolation=cv2.INTER_LINEAR)
|
||||||
|
depth_writer.write(depth_frame)
|
||||||
|
# depth_writer.write(depth_map)
|
||||||
|
|
||||||
|
key = cv2.waitKey(1)
|
||||||
|
if key & 0xFF == ord("q"):
|
||||||
|
if writer is not None:
|
||||||
|
writer.release()
|
||||||
|
writer = None
|
||||||
|
print("\n 停止程序")
|
||||||
|
break
|
||||||
|
if key == ord("r"):
|
||||||
|
# 重置光流追踪器(当追踪漂移时使用)
|
||||||
|
sensor1.reset()
|
||||||
|
print("传感器已重置")
|
||||||
|
|
||||||
|
# Release video writers
|
||||||
|
arrow_writer.release()
|
||||||
|
pyqt_writer.release()
|
||||||
|
depth_writer.release()
|
||||||
|
sensor1.disconnect() # 断开传感器连接
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
||||||
@@ -1,9 +1,7 @@
|
|||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import time
|
import time
|
||||||
import torch
|
|
||||||
import torchvision.transforms as T
|
|
||||||
from torchvision.models.optical_flow import raft_large
|
|
||||||
|
|
||||||
class Tracker:
|
class Tracker:
|
||||||
def __init__(self,
|
def __init__(self,
|
||||||
@@ -27,6 +25,9 @@ class Tracker:
|
|||||||
dis_preset = config_tracker["dis_preset"]
|
dis_preset = config_tracker["dis_preset"]
|
||||||
loop_closure_interval = config_tracker["loop_closure_interval"]
|
loop_closure_interval = config_tracker["loop_closure_interval"]
|
||||||
self.__algorithm = config_tracker["algorithm"]
|
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()):
|
if self.__algorithm == "raft" and (not self.cuda or not torch.cuda.is_available()):
|
||||||
self.__algorithm = "dis"
|
self.__algorithm = "dis"
|
||||||
@@ -199,6 +200,8 @@ class Tracker:
|
|||||||
"""
|
"""
|
||||||
Setup a simpler version using available models with performance optimizations
|
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
|
# Load model and move to GPU
|
||||||
model = raft_large(pretrained=True, progress=False).to(self.__raft_device)
|
model = raft_large(pretrained=True, progress=False).to(self.__raft_device)
|
||||||
model = model.eval()
|
model = model.eval()
|
||||||
|
|||||||
+34
-61
@@ -1,78 +1,51 @@
|
|||||||
import cv2
|
import cv2
|
||||||
import numpy as np
|
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.
|
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.
|
||||||
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 = frame.copy()
|
image = frame.copy()
|
||||||
# Convert grayscale to BGR to enable colored arrow drawing
|
|
||||||
if len(image.shape) == 2 or image.shape[2] == 1:
|
if len(image.shape) == 2 or image.shape[2] == 1:
|
||||||
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
|
||||||
|
|
||||||
# Calculate scaling factors between image and flow dimensions
|
h, w = image.shape[:2]
|
||||||
img_height, img_width = image.shape[:2]
|
# Arrow grid in pixel coordinates (assume frame and flow same size)
|
||||||
flow_height, flow_width = optical_flow.shape[:2]
|
rr = np.arange(grid_spacing // 2, h, grid_spacing)
|
||||||
scale_x, scale_y = img_width / flow_width, img_height / flow_height
|
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)
|
||||||
|
|
||||||
# Create sampling grid for flow field (used for interpolation source points)
|
# Displacement at each grid point (flow stored as half-displacement)
|
||||||
sample_step = 19 # Spacing between sample points
|
disp = 2.0 * optical_flow[rr, cc, :]
|
||||||
flow_start = np.stack(np.meshgrid(range(0, flow_width, sample_step), range(0, flow_height, sample_step)), 2)
|
mag = np.linalg.norm(disp, axis=1)
|
||||||
flow_end = (optical_flow[flow_start[:, :, 1], flow_start[:, :, 0], :] * 2.0 + flow_start).astype(np.int32)
|
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)
|
||||||
|
|
||||||
# Create regular grid for arrow placement in image coordinates
|
valid = (
|
||||||
grid_x, grid_y = np.mgrid[grid_spacing//2:flow_height:grid_spacing, grid_spacing//2:flow_width:grid_spacing]
|
(cc >= edge_margin) & (cc < w - edge_margin) & (rr >= edge_margin) & (rr < h - edge_margin)
|
||||||
x_axis = (grid_y.reshape(-1) * scale_x).astype(np.int64) # Scale to image width
|
& (end_x >= 0) & (end_x < w) & (end_y >= 0) & (end_y < h)
|
||||||
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
|
mag_range = mag.max() - mag.min()
|
||||||
grid_flat_scaled = grid_vector.reshape(-1, 2) * np.array([scale_x, scale_y])
|
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)
|
||||||
|
|
||||||
# Compute arrow endpoints and ensure they stay within image bounds
|
for i in np.where(valid)[0]:
|
||||||
end_x = (grid_flat_scaled[:, 0] / arrow_scale + x_axis).astype(np.int64)
|
pt1, pt2 = (int(cc[i]), int(rr[i])), (int(end_x[i]), int(end_y[i]))
|
||||||
end_y = (grid_flat_scaled[:, 1] / arrow_scale + y_axis).astype(np.int64)
|
color = tuple(int(c) for c in colors[i, 0])
|
||||||
valid_mask = (y_axis >= 0) & (y_axis < img_height) & (x_axis >= 0) & (x_axis < img_width)
|
if mag[i] < threshold:
|
||||||
valid_mask &= (end_x >= 0) & (end_x < img_width) & (end_y >= 0) & (end_y < img_height)
|
cv2.circle(image, pt1, radius=1, color=below_threshold_color, thickness=-1)
|
||||||
|
|
||||||
# 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:
|
else:
|
||||||
# Above threshold: draw full arrow with thickness proportional to magnitude
|
thickness = max(1, int(2 * mag[i] / mag.max())) if mag.max() > 0 else 2
|
||||||
# Arrow rooted at original position (x_axis, y_axis), head points to new position (end_x, end_y)
|
cv2.arrowedLine(image, pt1, pt2, color, thickness=thickness, tipLength=0.2)
|
||||||
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
|
return image
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user