import os os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" import cv2 import orisys import numpy as np import argparse def main(): parser = argparse.ArgumentParser(description='Calibration Data Collection') parser.add_argument('-i','--vid_source', type=str, default=0, help='Video Source') parser.add_argument('-id','--device_id', type=int, default=0, help='Device ID') args = parser.parse_args() device_id = args.device_id vid_source = args.vid_source sample_name = f"dd01a_{device_id}" #dd01 (平面款), dd01a (平面款 手搓), dd02 (手指款) test_type = "normal" # shearX, shearY, normal print("\n") print(f"Collecting data for device ID: {device_id} on test type: {test_type}, video source: {vid_source}") print("\n") sensor = orisys.Sensor(vid_source, config_name="./config/prototype_of3.json", cal_path= f"./config/{sample_name}.npy", cuda=False, verbose=True) print(f"\n按 'q' 退出程序，按 'r' 重置光流追踪器，按 's' 开始保存，按 'e' 停止保存，按 'j' 减小力，按 'l' 增大力") force = 0 # kgf force_step = 0.1 # kgf force_min = 0 # kgf force_max = 3 # kgf is_save = False save_mode = ["raw", "stitched", "arrows", "contact", "depth"] save_mode_index = 0 frame_count = 0 writer = None save_path = None while True: sensor.get_img() if is_save and writer is not None and sensor.frame is not None: frame_count += 1 if save_mode_index == 0: save_img = sensor.frame elif save_mode_index == 1: save_img = sensor.img elif save_mode_index == 2: save_img = arrows elif save_mode_index == 3: save_img = image_with_foe elif save_mode_index == 4: save_img = div_normalized # 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) sensor.compute_deformation() sensor.compute_contact() fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = sensor.read_info( sensor.info.FPS, sensor.info.VRAW, sensor.info.VNORMAL, sensor.info.IMG, sensor.info.CONTOUR, sensor.info.CENTROID, sensor.info.DEPTH, sensor.info.FNORMAL, sensor.info.FSHEARX, sensor.info.FSHEARY ) arrows = orisys.util.draw_arrows( img, flow, threshold=4, grid_spacing=10, arrow_scale=1.0 ) cv2.putText(arrows, f"FPS: {fps:.0f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2) cv2.imshow("Deformation", arrows) image_with_foe = orisys.util.draw_contact(img, contour, centroid) cv2.imshow("Contact Centroid", image_with_foe) 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：蓝色=低散度，红色=高散度） cv2.imshow("Depth Map", cv2.applyColorMap(div_normalized, cv2.COLORMAP_JET)) key = cv2.waitKey(1) & 0xFF if key == ord("s"): save_name = f"{sample_name}_{test_type}_{force:.1f}kgf" save_path = f"./data/{sample_name}/{test_type}/{save_name}.mp4" # Ensure directory exists os.makedirs(os.path.dirname(save_path), exist_ok=True) # Get frame dimensions from the actual frame save_img = None if save_mode_index == 0: save_img = sensor.frame elif save_mode_index == 1: save_img = sensor.img elif save_mode_index == 2: save_img = arrows elif save_mode_index == 3: save_img = image_with_foe elif save_mode_index == 4: save_img = div_normalized if save_img is not None: # 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] 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 else: print("No frame available to determine video dimensions") 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 if key == ord("p"): print(f"Normal Force: {fnormal:.1f} kgf, Shear ForceX: {fshearx:.1f} kgf, Shear ForceY: {fsheary:.1f} kgf") if key == ord("i"): save_mode_index += 1 if save_mode_index >= len(save_mode): save_mode_index = 0 print(f"Save mode: {save_mode[save_mode_index]}") if key == ord("k"): save_mode_index -= 1 if save_mode_index < 0: save_mode_index = len(save_mode) - 1 print(f"Save mode: {save_mode[save_mode_index]}") if key == ord("j"): force -= force_step if force < force_min: force = force_min print(f"Force: {force:.1f} kgf") if key == ord("l"): force += force_step if force > force_max: force = force_max print(f"Force: {force:.1f} kgf") if key == ord("q"): print("\n退出程序...") if writer is not None: writer.release() writer = None break if key == ord("r"): # 重置光流追踪器（当追踪漂移时使用） sensor.reset() print("传感器已重置") # ========================================================================= # 4. 清理资源 # ========================================================================= if writer is not None: writer.release() sensor.disconnect() # 释放视频源和相关资源 if __name__ == '__main__': main()