ykcheung
309 lines
11 KiB
Python
309 lines
11 KiB
Python
import os
|
|
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
|
|
import cv2
|
|
import orisys
|
|
import numpy as np
|
|
import time
|
|
|
|
import pyqtgraph as pg
|
|
from PyQt5.QtGui import QImage
|
|
|
|
def create_plot():
|
|
p1 = win.addPlot(title="法相")
|
|
p1.setLabel('bottom', '时间', units='s')
|
|
win.nextRow()
|
|
p2 = win.addPlot(title="切向 X")
|
|
p2.setLabel('bottom', '时间', units='s')
|
|
win.nextRow()
|
|
p3 = win.addPlot(title="切向 Y")
|
|
p3.setLabel('bottom', '时间', units='s')
|
|
plot_af = p1.plot(pen='w')
|
|
plot_tfx = p2.plot(pen='w')
|
|
plot_tfy = p3.plot(pen='w')
|
|
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 plot_to_frame(plot_widget, width=None, height=None):
|
|
"""
|
|
Convert pyqtgraph GraphicsLayoutWidget to numpy array (BGR format for OpenCV).
|
|
|
|
Args:
|
|
plot_widget: pyqtgraph GraphicsLayoutWidget to capture
|
|
width: Optional width for output frame (default: widget width)
|
|
height: Optional height for output frame (default: widget height)
|
|
|
|
Returns:
|
|
np.ndarray: BGR image array (height, width, 3) ready for cv2.VideoWriter
|
|
"""
|
|
# Get the widget size
|
|
if width is None:
|
|
width = plot_widget.width()
|
|
if height is None:
|
|
height = plot_widget.height()
|
|
|
|
# Grab the widget as a QPixmap
|
|
pixmap = plot_widget.grab()
|
|
|
|
# Resize if needed
|
|
if pixmap.width() != width or pixmap.height() != height:
|
|
pixmap = pixmap.scaled(width, height)
|
|
|
|
# Convert QPixmap to QImage
|
|
qimage = pixmap.toImage()
|
|
|
|
# Convert QImage to RGB format
|
|
qimage = qimage.convertToFormat(QImage.Format_RGB888)
|
|
|
|
# Get image dimensions
|
|
img_width = qimage.width()
|
|
img_height = qimage.height()
|
|
|
|
# Convert QImage to numpy array
|
|
# Use constBits() to get raw bytes, then reshape
|
|
buffer = qimage.constBits()
|
|
if buffer is not None:
|
|
arr = np.frombuffer(buffer, dtype=np.uint8).reshape((img_height, img_width, 3))
|
|
else:
|
|
# Fallback: create empty array
|
|
arr = np.zeros((img_height, img_width, 3), dtype=np.uint8)
|
|
|
|
# Convert RGB to BGR for OpenCV
|
|
bgr = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR)
|
|
|
|
return bgr
|
|
|
|
# 深度图
|
|
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
|
|
for i in range(3):
|
|
data[:,0,i] = np.linspace(-npoints,0,npoints)
|
|
|
|
|
|
|
|
def main():
|
|
global start_time
|
|
|
|
# 打开视频
|
|
# sensor1 = Sensor("./s1.mp4",isstitch=True,cuda = True,verbose = True)
|
|
# 方式1: 从视频文件读取
|
|
sensor1 = orisys.Sensor("./data/pulse_test_1.mp4", config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_98.npy", isstitch=True, cuda=False, verbose=True)
|
|
|
|
start_time = time.time()
|
|
is_save = True
|
|
frame_count = 0
|
|
writer = None
|
|
save_path = None
|
|
|
|
# Video writer for plot
|
|
is_save_plot = True
|
|
plot_writer = None
|
|
plot_frame_count = 0
|
|
plot_save_path = None
|
|
|
|
if True:
|
|
save_path = f"./data/pulse_test_arrows.mp4"
|
|
# Ensure directory exists
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
# Get frame dimensions from the actual frame
|
|
save_img = sensor1.frame
|
|
|
|
# if save_img is not None:
|
|
if True:
|
|
# Convert grayscale (1 channel) to BGR (3 channels) if needed for dimension check
|
|
# temp_img = save_img.copy()
|
|
# if len(temp_img.shape) == 2 or (len(temp_img.shape) == 3 and temp_img.shape[2] == 1):
|
|
# if len(temp_img.shape) == 3:
|
|
# temp_img = temp_img.squeeze(2)
|
|
# temp_img = cv2.cvtColor(temp_img, cv2.COLOR_GRAY2BGR)
|
|
|
|
# frame_height, frame_width = temp_img.shape[:2]
|
|
frame_height = 400
|
|
frame_width = 400
|
|
writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (frame_width, frame_height))
|
|
if writer.isOpened():
|
|
is_save = True
|
|
print(f"Start writing to {save_path}")
|
|
else:
|
|
print(f"Failed to open video writer for {save_path}")
|
|
writer = None
|
|
|
|
# Setup video writer for plot
|
|
plot_save_path = f"./data/pulse_test_plot.mp4"
|
|
os.makedirs(os.path.dirname(plot_save_path), exist_ok=True)
|
|
# Get plot widget dimensions
|
|
plot_width = win.width()
|
|
plot_height = win.height()
|
|
# Use fixed dimensions if widget not yet rendered
|
|
if plot_width <= 0 or plot_height <= 0:
|
|
plot_width = 800
|
|
plot_height = 600
|
|
|
|
plot_writer = cv2.VideoWriter(plot_save_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (plot_width, plot_height))
|
|
if plot_writer.isOpened():
|
|
is_save_plot = True
|
|
print(f"Start writing plot to {plot_save_path}")
|
|
else:
|
|
print(f"Failed to open plot video writer for {plot_save_path}")
|
|
plot_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, contour, centroid, depth_map = sensor1.read_info(sensor1.info.FPS, sensor1.info.FNORMAL, sensor1.info.FSHEARX, sensor1.info.FSHEARY, sensor1.info.VRAW,sensor1.info.CONTOUR,
|
|
sensor1.info.CENTROID,
|
|
sensor1.info.DEPTH)
|
|
print(fps,fn,fx,fy)
|
|
|
|
|
|
# 显示触碰点和区域 # 步骤 3.2: 在图像上触碰点和区域
|
|
white_img = np.full_like(sensor1.img, 0)
|
|
|
|
image_with_foe = orisys.util.draw_contact(white_img, contour, centroid,region_color=(255, 255, 255))
|
|
cv2.imshow("Contact Centroid", image_with_foe)
|
|
# 显示向量场
|
|
|
|
arrows = orisys.util.draw_arrows(white_img, flow, threshold=5, grid_spacing=10, arrow_scale=1.0)
|
|
cv2.imshow("Deformation", arrows)
|
|
# 显示深度图
|
|
mag = draw_magnitude_map(flow)
|
|
cv2.imshow("Magnitude", mag)
|
|
|
|
if is_save and writer is not None and sensor1.frame is not None:
|
|
frame_count += 1
|
|
save_img = arrows
|
|
|
|
# Convert grayscale (1 channel) to BGR (3 channels) if needed
|
|
if len(save_img.shape) == 2 or (len(save_img.shape) == 3 and save_img.shape[2] == 1):
|
|
if len(save_img.shape) == 3:
|
|
save_img = save_img.squeeze(2)
|
|
save_img = cv2.cvtColor(save_img, cv2.COLOR_GRAY2BGR)
|
|
|
|
writer.write(save_img)
|
|
|
|
# 显示图表 - 更新时间轴
|
|
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, :] = current_time # set current time for all plots
|
|
data[-1, 1, 0] = fn
|
|
data[-1, 1, 1] = fx
|
|
data[-1, 1, 2] = fy
|
|
update(plot_af, plot_tfx, plot_tfy, data, p1, p2, p3) # Update plot
|
|
|
|
# Capture and save plot frame
|
|
if is_save_plot and plot_writer is not None:
|
|
plot_frame_count += 1
|
|
# Process Qt events to ensure plot is updated
|
|
app.processEvents()
|
|
# Capture plot as frame
|
|
plot_frame = plot_to_frame(win)
|
|
plot_writer.write(plot_frame)
|
|
|
|
key = cv2.waitKey(1)
|
|
if key & 0xFF == ord("q"):
|
|
print("\n 停止程序")
|
|
break
|
|
|
|
if key == ord("e"):
|
|
is_save = False
|
|
if writer is not None:
|
|
writer.release()
|
|
writer = None
|
|
if save_path is not None:
|
|
print(f"Stop writing to {save_path}, total frames: {frame_count}")
|
|
else:
|
|
print(f"Stop writing, total frames: {frame_count}")
|
|
frame_count = 0
|
|
|
|
# Stop plot recording
|
|
is_save_plot = False
|
|
if plot_writer is not None:
|
|
plot_writer.release()
|
|
plot_writer = None
|
|
if plot_save_path is not None:
|
|
print(f"Stop writing plot to {plot_save_path}, total frames: {plot_frame_count}")
|
|
else:
|
|
print(f"Stop writing plot, total frames: {plot_frame_count}")
|
|
plot_frame_count = 0
|
|
|
|
sensor1.disconnect() # 断开传感器连接
|
|
|
|
# Clean up video writers
|
|
if writer is not None:
|
|
writer.release()
|
|
if plot_writer is not None:
|
|
plot_writer.release()
|
|
print(f"Plot video saved: {plot_save_path}, total frames: {plot_frame_count}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|