ykcheungab
253 lines
9.6 KiB
Python
253 lines
9.6 KiB
Python
import os
|
|
|
|
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
|
import argparse
|
|
import time
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import pyqtgraph as pg
|
|
from PyQt5.QtGui import QFont, QImage
|
|
from pyqtgraph.Qt import QtGui
|
|
|
|
import orisys
|
|
|
|
|
|
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', '#0B1016') # 深蓝灰背景
|
|
pg.setConfigOption('foreground', '#A0A0A0') # 浅灰文字
|
|
pg.setConfigOptions(antialias=True) # 抗锯齿
|
|
theme_color = '#00E5FF'
|
|
line_pen = pg.mkPen(color=theme_color, width=2)
|
|
grad_brush = make_gradient_brush(theme_color, max_alpha=120)
|
|
|
|
axis_font = QFont("Segoe UI", 11)
|
|
p1 = win.addPlot()
|
|
p1.setLabel('left', 'N')
|
|
p1.getAxis('left').label.setFont(axis_font)
|
|
p1.getAxis('left').label.setRotation(-90)
|
|
p1.showGrid(x=True, y=True, alpha=0.3)
|
|
win.nextRow()
|
|
p2 = win.addPlot()
|
|
p2.setLabel('left', 'X')
|
|
p2.getAxis('left').label.setFont(axis_font)
|
|
p2.getAxis('left').label.setRotation(-90)
|
|
p2.showGrid(x=True, y=True, alpha=0.3)
|
|
win.nextRow()
|
|
p3 = win.addPlot()
|
|
p3.setLabel('left', 'Y')
|
|
p3.setLabel('bottom', 'Time', units='s')
|
|
p3.getAxis('left').label.setFont(axis_font)
|
|
p3.getAxis('left').label.setRotation(-90)
|
|
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 # 图表数据长度
|
|
pg.setConfigOption('background', 'k')
|
|
pg.setConfigOption('foreground', 'w')
|
|
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
|
|
fps = 30.0
|
|
for i in range(3):
|
|
data[:,0,i] = np.linspace(0.0, npoints / fps, npoints)
|
|
|
|
|
|
|
|
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, (365, 365))
|
|
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, (365, 365))
|
|
|
|
sensor1 = orisys.Sensor(args.video, config_name="./config/ddjx01.json", cal_path= "./config/dd01a_0306.npy", isstitch=True, cuda=False, verbose=True)
|
|
|
|
start_time = time.time()
|
|
is_save = False
|
|
frame_count = 0
|
|
writer = None
|
|
while True:
|
|
isimg = sensor1.get_img() # 提取并拼接复眼图像
|
|
if isimg is None:
|
|
break
|
|
|
|
|
|
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)
|
|
# 显示向量场
|
|
_bg_bgr = (22, 16, 11)
|
|
black_img = np.empty_like(sensor1.img)
|
|
black_img[:, :, 0] = _bg_bgr[0]
|
|
black_img[:, :, 1] = _bg_bgr[1]
|
|
black_img[:, :, 2] = _bg_bgr[2]
|
|
|
|
arrows = orisys.util.draw_arrows(black_img, flow, threshold=2, grid_spacing=20, arrow_scale=0.5,below_threshold_color=(160, 160, 160),edge_margin=15,colormap=cv2.COLORMAP_PARULA)
|
|
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_JET)
|
|
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)
|
|
update(plot_af, plot_tfx, plot_tfy, data, 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, (365, 365), 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()
|