TPP-G6
196 lines
7.3 KiB
Python
196 lines
7.3 KiB
Python
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()
|