prepare for calibration data collection

2026-01-19 10:53:51 +08:00
parent 781fe5aa9b
commit 2109a959f3
7 changed files with 473 additions and 8 deletions
@@ -0,0 +1,141 @@
+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='s1.mp4', 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", isstitch=True, cuda=True, 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
+    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
+            writer.write(sensor.frame)
+
+        sensor.compute_deformation()
+        sensor.compute_contact()
+
+        fps, flow, vnormal, img, contour, centroid, depth_map = sensor.read_info(
+            sensor.info.FPS,
+            sensor.info.VRAW,
+            sensor.info.VNORMAL,
+            sensor.info.IMG,
+            sensor.info.CONTOUR,
+            sensor.info.CENTROID,
+            sensor.info.DEPTH
+        )
+        
+        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
+            if sensor.frame is not None:
+                frame_height, frame_width = sensor.frame.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("i"):
+            print("Up arrow pressed")
+        if key == ord("k"):
+            print("Down arrow pressed")
+        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退出程序...")
+            break
+        if key == ord("r"):
+            # 重置光流追踪器（当追踪漂移时使用）
+            sensor.reset()
+            print("传感器已重置")
+
+        
+    # =========================================================================
+    # 4. 清理资源
+    # =========================================================================
+    if writer is not None:
+        writer.release()
+    sensor.disconnect()  # 释放视频源和相关资源
+
+
+if __name__ == '__main__':
+    main()