update test script

calibration_data_collection.py

@@ -7,7 +7,7 @@ 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('-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()
@@ -27,6 +27,8 @@ def main():
     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
@@ -36,7 +38,24 @@ def main():
 
         if is_save and writer is not None and sensor.frame is not None:
             frame_count += 1
-            writer.write(sensor.frame)
+            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()
@@ -54,7 +73,7 @@ def main():
         arrows = orisys.util.draw_arrows(
             img,
             flow,
-            threshold=4,
+            threshold=2,
             grid_spacing=10,
             arrow_scale=1.0
         )
@@ -81,8 +100,27 @@ def main():
             # 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]
+            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
@@ -105,9 +143,16 @@ def main():
             frame_count = 0
 
         if key == ord("i"):
-            print("Up arrow pressed")
+            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"):
-            print("Down arrow pressed")
+            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:
@@ -122,6 +167,9 @@ def main():
 
         if key == ord("q"):
             print("\n退出程序...")
+            if writer is not None:
+                writer.release()
+                writer = None
             break
         if key == ord("r"):
             # 重置光流追踪器（当追踪漂移时使用）