diff --git a/calibration_data_collection.py b/calibration_data_collection.py index 716ad1a..cc29f8e 100644 --- a/calibration_data_collection.py +++ b/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"): # 重置光流追踪器(当追踪漂移时使用)