diff --git a/benchmark.py b/benchmark.py index 9396414..5eadc03 100644 --- a/benchmark.py +++ b/benchmark.py @@ -12,7 +12,7 @@ def main(): # cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件 # 方式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: 从摄像头读取（取消注释使用） # sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) diff --git a/config/compound_eye.png b/config/compound_eye.png index 3aea01a..eed73b6 100644 Binary files a/config/compound_eye.png and b/config/compound_eye.png differ diff --git a/config/prototype_of3.json b/config/prototype_of3.json index 3b1245d..e1d8c06 100644 --- a/config/prototype_of3.json +++ b/config/prototype_of3.json @@ -6,7 +6,7 @@ "fps": 30, "units": 15, "units_used": 9, - "resolution_stitch": [128, 128], + "resolution_stitch": [384, 384], "stitch_res_list": [384,256,128,64] }, "stitcher": { diff --git a/config/prototype_raft.json b/config/prototype_raft.json index e0b65f7..097e2bb 100644 --- a/config/prototype_raft.json +++ b/config/prototype_raft.json @@ -6,7 +6,7 @@ "fps": 30, "units": 15, "units_used": 9, - "resolution_stitch": [128, 128] + "resolution_stitch": [384, 384] }, "stitcher": { "min_area_frac": 0.01, diff --git a/example1_basic_pipeline.py b/example1_basic_pipeline.py index 23de99d..3d7bb1a 100644 --- a/example1_basic_pipeline.py +++ b/example1_basic_pipeline.py @@ -57,8 +57,8 @@ def main(): arrows = orisys.util.draw_arrows( img, flow, - threshold=4, - grid_spacing=40, + threshold=2, + grid_spacing=20, arrow_scale=1.0 ) cv2.imshow("Deformation", arrows) diff --git a/promote_video.py b/promote_video.py new file mode 100644 index 0000000..091fae1 --- /dev/null +++ b/promote_video.py @@ -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() diff --git a/src/orisys/tracker.py b/src/orisys/tracker.py index 273c46a..9e31628 100644 --- a/src/orisys/tracker.py +++ b/src/orisys/tracker.py @@ -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() diff --git a/src/orisys/util.py b/src/orisys/util.py index ef1a1a1..55a78a1 100644 --- a/src/orisys/util.py +++ b/src/orisys/util.py @@ -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