Diff of /src/compute_target.py [000000] .. [c1b1c5]

Switch to side-by-side view

--- a
+++ b/src/compute_target.py
@@ -0,0 +1,351 @@
+# Given HPE results and ratios, compute the final target point
+
+import cv2
+import matplotlib.pyplot as plt
+from utils.trajectory_io import *
+import pickle
+from utils.triangulation import *
+import open3d as o3d
+from scipy.linalg import null_space
+
+from subject_info import SUBJECT_NAME, SCAN_POSE
+import argparse
+import time
+
+parser = argparse.ArgumentParser(description='Compute target')
+parser.add_argument('--pose_model', type=str, default='ViTPose_large', help='pose model')
+parser.add_argument('--subject_name', type=str, default='John Doe', help='subject name')
+parser.add_argument('--scan_pose', type=str, default='none', help='scan pose')
+parser.add_argument('--target1_r1', type=float, default=0.3, help='target1_r1')
+parser.add_argument('--target1_r2', type=float, default=0.1, help='target1_r2')
+parser.add_argument('--target2_r1', type=float, default=0.3, help='target2_r1')
+parser.add_argument('--target2_r2', type=float, default=0.55, help='target2_r2')
+parser.add_argument('--target4_r1', type=float, default=0.35, help='target4_r1')
+parser.add_argument('--target4_r2', type=float, default=0.1, help='target4_r2')
+
+args = parser.parse_args()
+POSE_MODEL = args.pose_model
+if args.subject_name != 'John Doe':
+    SUBJECT_NAME = args.subject_name
+if args.scan_pose != 'none':
+    SCAN_POSE = args.scan_pose
+
+
+def target12_3D(l_shoulder_cam1, l_shoulder_cam2, r_shoulder_cam1, r_shoulder_cam2,
+                cam1_intr, cam2_intr, T_base_cam1, T_base_cam2, ratio_1=0.3, ratio_2=0.1):
+    '''
+    A 3D method to determine the 1st and 2nd target 3D coordinate. Didn't use nipples.
+    :X1, X2 --3D coordinates of the right, left shoulder
+    :X3 -- 3D coordinate of the point determined by ratio_1
+    :t1 -- the direction vector of the line connecting X1 and X2
+    :t2 -- the direction vector of the line connecting X3 and target
+    :param ratio_1: ratio on the line connecting X1 and X2
+    :param ratio_2: ratio between X3 to target and X1 to X2
+    '''
+    X1 = reconstruct(r_shoulder_cam1, r_shoulder_cam2, cam1_intr, cam2_intr, T_base_cam1, T_base_cam2)
+    X1 = triangulation_post_process(X1)
+    X2 = reconstruct(l_shoulder_cam1, l_shoulder_cam2, cam1_intr, cam2_intr, T_base_cam1, T_base_cam2)
+    X2 = triangulation_post_process(X2)
+
+    X3 = X1 + ratio_1 * (X2 - X1)
+
+    t1 = ((X2 - X1) / np.linalg.norm(X2 - X1)).squeeze()
+    n = np.array([0, 0, 1])  # z-axis of the base frame / normal vector of the base frame's x-y plane
+    A = np.vstack([t1, n])
+    t2 = null_space(A)  # t2 is perpendicular to both t1 and n
+    t2 = t2 / np.linalg.norm(t2)  # normalize t2
+
+    if t2[0] < 0:
+        t2 = -t2
+
+    target = X3 + ratio_2 * np.linalg.norm(X2 - X1) * t2
+    target = triangulation_post_process(target, verbose=False)
+    position_data['front'] = [X1, X2, t2]
+
+    return target
+
+
+def triangulation_post_process(original_3d, verbose=False):
+    '''
+    Use the original x&y values to find a closest point in the point cloud.
+    Use the new point's z value as the new z, combined with the original x&y.
+    :param original_3d: The original 3D coordinate from triangulation
+    '''
+    if verbose:
+        print("original_3d: ", original_3d)
+    idx = np.argmin(np.square(coordinates[:, 0:2] - original_3d.squeeze()[0:2]).sum(axis=1))
+    if verbose:
+        new_3d = coordinates[idx]
+        print("new_3d: ", new_3d)
+    new_z = coordinates[idx, 2]
+    new_3d = np.vstack((original_3d[0:2], [new_z]))
+    return new_3d
+
+
+def target4_3D(r_shoulder_cam1, r_shoulder_cam2, r_hip_cam1, r_hip_cam2,
+               cam1_intr, cam2_intr, T_base_cam1, T_base_cam2, ratio_1=0.35, ratio_2=0.1):
+    '''
+    Similar to 1st&2nd target but ratio 2 operated on X1 to X3, instead of X1 to X2
+    :X1, X2 --3D coordinates of the right shoulder and the right hip
+    :X3 -- 3D coordinate of the point determined by ratio_1
+    :t1 -- the direction vector of the line connecting X1 and X2
+    :t2 -- the direction vector of the line connecting X3 and target
+    :param ratio_1: ratio on the line connecting X1 and X2
+    :param ratio_2: ratio between X3 to target and X1 to X3
+    '''
+    X1 = reconstruct(r_shoulder_cam1, r_shoulder_cam2, cam1_intr, cam2_intr, T_base_cam1, T_base_cam2)
+    X1 = triangulation_post_process(X1)
+    X2 = reconstruct(r_hip_cam1, r_hip_cam2, cam1_intr, cam2_intr, T_base_cam1, T_base_cam2)
+    X2 = triangulation_post_process(X2)
+
+    X3 = X1 + ratio_1 * (X2 - X1)
+
+    t1 = ((X2 - X1) / np.linalg.norm(X2 - X1)).squeeze()
+    n = np.array([0, 0, 1])
+    A = np.vstack([t1, n])
+    t2 = null_space(A)
+    t2 = t2 / np.linalg.norm(t2)
+
+    if t2[1] < 0:
+        t2 = -t2
+
+    target = X3 + ratio_2 * np.linalg.norm(X1 - X3) * t2
+    target = triangulation_post_process(target)
+    position_data['side'] = [X1, X2, t2]
+
+    return target
+
+
+def draw_front_target_point(target_point1, target_point2, l_shoulder1, r_shoulder1, l_shoulder2, r_shoulder2, target):
+    image1 = plt.imread(folder_path + 'color_images/cam_1.jpg')
+    cv2.circle(image1, (int(target_point1[0]), int(target_point1[1])), 2, (36, 255, 12), 2, -1)
+    cv2.circle(image1, (int(l_shoulder1[0]), int(l_shoulder1[1])), 2, (36, 255, 12), 2, -1)
+    cv2.circle(image1, (int(r_shoulder1[0]), int(r_shoulder1[1])), 2, (36, 255, 12), 2, -1)
+    cv2.line(image1, (int(l_shoulder1[0]), int(l_shoulder1[1])), (int(r_shoulder1[0]), int(r_shoulder1[1])),
+             color=(255, 0, 0), thickness=2)
+
+    plt.subplot(121)
+    plt.imshow(image1)
+
+    image2 = plt.imread(folder_path + 'color_images/cam_2.jpg')
+    cv2.circle(image2, (int(target_point2[0]), int(target_point2[1])), 2, (36, 255, 12), 2, -1)
+    cv2.circle(image2, (int(l_shoulder2[0]), int(l_shoulder2[1])), 2, (36, 255, 12), 2, -1)
+    cv2.circle(image2, (int(r_shoulder2[0]), int(r_shoulder2[1])), 2, (36, 255, 12), 2, -1)
+    cv2.line(image2, (int(l_shoulder2[0]), int(l_shoulder2[1])), (int(r_shoulder2[0]), int(r_shoulder2[1])),
+             color=(255, 0, 0), thickness=2)
+
+    plt.subplot(122)
+    plt.imshow(image2)
+
+    plt.show()
+
+    if target == 1:
+        plt.imsave(folder_path + 'target1_highlight/cam_1.jpg', image1)
+        plt.imsave(folder_path + 'target1_highlight/cam_2.jpg', image2)
+    else:
+        plt.imsave(folder_path + 'target2_highlight/cam_1.jpg', image1)
+        plt.imsave(folder_path + 'target2_highlight/cam_2.jpg', image2)
+
+
+def draw_side_target_point(target_point1, target_point2, r_hip1, r_shoulder1, r_hip2, r_shoulder2):
+    image1 = plt.imread(folder_path + 'color_images/cam_1.jpg')
+    cv2.circle(image1, (int(target_point1[0]), int(target_point1[1])), 3, (36, 255, 12), 3, -1)
+    cv2.circle(image1, (int(r_hip1[0]), int(r_hip1[1])), 2, (36, 255, 12), 2, -1)
+    cv2.circle(image1, (int(r_shoulder1[0]), int(r_shoulder1[1])), 2, (36, 255, 12), 2, -1)
+    cv2.line(image1, (int(r_hip1[0]), int(r_hip1[1])), (int(r_shoulder1[0]), int(r_shoulder1[1])),
+             color=(255, 0, 0), thickness=2)
+
+    plt.subplot(121)
+    plt.imshow(image1)
+
+    image2 = plt.imread(folder_path + 'color_images/cam_2.jpg')
+    cv2.circle(image2, (int(target_point2[0]), int(target_point2[1])), 3, (36, 255, 12), 3, -1)
+    cv2.circle(image2, (int(r_hip2[0]), int(r_hip2[1])), 2, (36, 255, 12), 2, -1)
+    cv2.circle(image2, (int(r_shoulder2[0]), int(r_shoulder2[1])), 2, (36, 255, 12), 2, -1)
+    cv2.line(image2, (int(r_hip2[0]), int(r_hip2[1])), (int(r_shoulder2[0]), int(r_shoulder2[1])),
+             color=(255, 0, 0), thickness=2)
+
+    plt.subplot(122)
+    plt.imshow(image2)
+
+    plt.show()
+    plt.imsave(folder_path + 'target4_highlight/cam_1.jpg', image1)
+    plt.imsave(folder_path + 'target4_highlight/cam_2.jpg', image2)
+
+folder_path = 'src/data/' + SUBJECT_NAME + '/' + SCAN_POSE + '/'  # when run from src directory
+# folder_path = '../data/' + SUBJECT_NAME + '/' + SCAN_POSE + '/'  # when run from evaluation directory
+
+# read RGBD intrinsics
+with open(folder_path + 'intrinsics/cam_1_intrinsics.pickle', 'rb') as f:
+    cam1_intr = pickle.load(f)
+
+with open(folder_path + 'intrinsics/cam_2_intrinsics.pickle', 'rb') as f:
+    cam2_intr = pickle.load(f)
+
+with open(folder_path + 'intrinsics/cam_1_depth_intr.pickle', 'rb') as f:
+    depth_cam1_intr = pickle.load(f)
+
+with open(folder_path + 'intrinsics/cam_2_depth_intr.pickle', 'rb') as f:
+    depth_cam2_intr = pickle.load(f)
+
+# read extrinsics
+camera_poses = read_trajectory(folder_path + "odometry.log")
+T_cam1_base = camera_poses[0].pose
+T_base_cam1 = np.linalg.inv(T_cam1_base)
+# print("T_base_cam1/extr: \n", T_base_cam1)
+T_cam2_base = camera_poses[1].pose
+T_base_cam2 = np.linalg.inv(T_cam2_base)
+# print("T_base_cam2/extr: \n", T_base_cam2)
+
+# read pose keypoints
+with open(folder_path + POSE_MODEL + '/keypoints/cam_1_keypoints.pickle', 'rb') as f:
+    cam1_keypoints = pickle.load(f)
+    if POSE_MODEL == "OpenPose":
+        l_shoulder1 = np.array(cam1_keypoints['people'][0]['pose_keypoints_2d'][15:17])
+        r_shoulder1 = np.array(cam1_keypoints['people'][0]['pose_keypoints_2d'][6:8])
+        l_hip1 = np.array(cam1_keypoints['people'][0]['pose_keypoints_2d'][36:38])
+        r_hip1 = np.array(cam1_keypoints['people'][0]['pose_keypoints_2d'][27:29])
+    else:
+        l_shoulder1 = cam1_keypoints[0]['keypoints'][5][:2]
+        r_shoulder1 = cam1_keypoints[0]['keypoints'][6][:2]
+        l_hip1 = cam1_keypoints[0]['keypoints'][11][:2]
+        r_hip1 = cam1_keypoints[0]['keypoints'][12][:2]
+
+with open(folder_path + POSE_MODEL + '/keypoints/cam_2_keypoints.pickle', 'rb') as f:
+    cam2_keypoints = pickle.load(f)
+    if POSE_MODEL == "OpenPose":
+        l_shoulder2 = np.array(cam2_keypoints['people'][0]['pose_keypoints_2d'][15:17])
+        r_shoulder2 = np.array(cam2_keypoints['people'][0]['pose_keypoints_2d'][6:8])
+        l_hip2 = np.array(cam2_keypoints['people'][0]['pose_keypoints_2d'][36:38])
+        r_hip2 = np.array(cam2_keypoints['people'][0]['pose_keypoints_2d'][27:29])
+    else:
+        l_shoulder2 = cam2_keypoints[0]['keypoints'][5][:2]
+        r_shoulder2 = cam2_keypoints[0]['keypoints'][6][:2]
+        l_hip2 = cam2_keypoints[0]['keypoints'][11][:2]
+        r_hip2 = cam2_keypoints[0]['keypoints'][12][:2]
+
+############ TSDF volume ############
+start_time = time.time()
+volume = o3d.pipelines.integration.ScalableTSDFVolume(
+    voxel_length=1 / 512.0,
+    sdf_trunc=0.04,
+    color_type=o3d.pipelines.integration.TSDFVolumeColorType.RGB8)
+
+# show RGBD images from all the views
+for i in range(len(camera_poses)):
+    # print("Integrate {:d}-th image into the volume.".format(i))
+    color = o3d.io.read_image(folder_path + "color_images/cam_{}.jpg".format(i+1))
+    depth = o3d.io.read_image(folder_path + "depth_images/cam_{}.png".format(i+1))
+
+    depth_image = np.asanyarray(depth)
+    color_image = np.asanyarray(color)
+
+    depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(depth_image, alpha=0.09), cv2.COLORMAP_JET)
+    depth_colormap_dim = depth_colormap.shape
+    color_colormap_dim = color_image.shape
+
+    # If depth and color resolutions are different, resize color image to match depth image for display
+    if depth_colormap_dim != color_colormap_dim:
+        resized_color_image = cv2.resize(color_image, dsize=(depth_colormap_dim[1], depth_colormap_dim[0]),
+                                         interpolation=cv2.INTER_AREA)
+        images = np.hstack((resized_color_image, depth_colormap))
+    else:
+        images = np.hstack((color_image, depth_colormap))
+
+    plt.imsave(folder_path + "depth_colormap/cam_{}.png".format(i+1), depth_colormap)
+
+    # Show images
+    cv2.namedWindow('RealSense', cv2.WINDOW_AUTOSIZE)
+    plt.imshow(images)
+    plt.show()
+
+    rgbd = o3d.geometry.RGBDImage.create_from_color_and_depth(
+        color, depth, depth_trunc=4.0, convert_rgb_to_intensity=False)
+
+    cam_intr = depth_cam1_intr if i == 0 else depth_cam2_intr  # use depth camera's intrinsics
+
+    intr = o3d.camera.PinholeCameraIntrinsic(
+        width=640,
+        height=480,
+        fx=cam_intr[0,0],
+        fy=cam_intr[1,1],
+        cx=cam_intr[0,2],
+        cy=cam_intr[1,2]
+    )
+
+    volume.integrate(rgbd, intr, np.linalg.inv(camera_poses[i].pose))
+print("Volume Integration: {:.3f} s".format(time.time() - start_time))
+
+# point cloud generation
+pcd = volume.extract_point_cloud()
+downpcd = pcd.voxel_down_sample(voxel_size=0.01)
+downpcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
+o3d.visualization.draw_geometries([downpcd])
+
+coordinates = np.asarray(downpcd.points)
+normals = np.asarray(downpcd.normals)
+
+# save data, write pcd
+o3d.io.write_point_cloud(folder_path + 'downpcd.ply', downpcd)
+with open(folder_path + 'pcd_coordinates.pickle', 'wb') as f:
+    pickle.dump(coordinates, f)
+
+with open(folder_path + 'pcd_normals.pickle', 'wb') as f:
+    pickle.dump(normals, f)
+
+# Read pcd coordianates
+with open(folder_path + 'pcd_coordinates.pickle', 'rb') as f:
+    coordinates = pickle.load(f)
+
+final_target_list = []  # the target coordinates used to navigate the robotic arm
+position_data = {}  # store
+
+if SCAN_POSE == 'front':
+    target1 = target12_3D(l_shoulder1, l_shoulder2, r_shoulder1, r_shoulder2,
+                          cam1_intr, cam2_intr, T_base_cam1, T_base_cam2, ratio_1=args.target1_r1,
+                          ratio_2=args.target1_r2)
+    target2 = target12_3D(l_shoulder1, l_shoulder2, r_shoulder1, r_shoulder2,
+                          cam1_intr, cam2_intr, T_base_cam1, T_base_cam2, ratio_1=args.target2_r1,
+                          ratio_2=args.target2_r2)
+
+    target1_2d_cam1 = from_homog(cam1_intr @ T_base_cam1[0:3, :] @ to_homog(target1)).squeeze()
+    target1_2d_cam2 = from_homog(cam2_intr @ T_base_cam2[0:3, :] @ to_homog(target1)).squeeze()
+
+    target2_2d_cam1 = from_homog(cam1_intr @ T_base_cam1[0:3, :] @ to_homog(target2)).squeeze()
+    target2_2d_cam2 = from_homog(cam2_intr @ T_base_cam2[0:3, :] @ to_homog(target2)).squeeze()
+
+    draw_front_target_point(target1_2d_cam1, target1_2d_cam2, l_shoulder1, r_shoulder1, l_shoulder2, r_shoulder2, target=1)
+    draw_front_target_point(target2_2d_cam1, target2_2d_cam2, l_shoulder1, r_shoulder1, l_shoulder2, r_shoulder2, target=2)
+
+    final_target_list.append(target1)
+    final_target_list.append(target2)
+
+elif SCAN_POSE == 'side':
+    target4 = target4_3D(r_shoulder1, r_shoulder2, r_hip1, r_hip2,
+                         cam1_intr, cam2_intr, T_base_cam1, T_base_cam2, ratio_1=args.target4_r1,
+                         ratio_2=args.target4_r2)
+
+    target4_2d_cam1 = from_homog(cam1_intr @ T_base_cam1[0:3, :] @ to_homog(target4)).squeeze()
+    target4_2d_cam2 = from_homog(cam2_intr @ T_base_cam2[0:3, :] @ to_homog(target4)).squeeze()
+
+    draw_side_target_point(target4_2d_cam1, target4_2d_cam2, r_hip1, r_shoulder1, r_hip2, r_shoulder2)
+
+    final_target_list.append(target4)
+
+# # save results with initial ratio parameters
+# with open(folder_path + POSE_MODEL + '/final_target.pickle', 'wb') as f:
+#     pickle.dump(final_target_list, f)
+#
+# with open(folder_path + POSE_MODEL + '/position_data.pickle', 'wb') as f:
+#     pickle.dump(position_data, f)
+
+# # save results with optimized ratio parameters based on planar euclidean distance
+# with open(folder_path + POSE_MODEL + '/final_target_opt.pickle', 'wb') as f:
+#     pickle.dump(final_target_list, f)
+#
+# with open(folder_path + POSE_MODEL + '/position_data_opt.pickle', 'wb') as f:
+#     pickle.dump(position_data, f)
+
+# # save as temporary file for evaluation
+# with open(folder_path + POSE_MODEL + '/tmp_target_test.pickle', 'wb') as f:
+#     pickle.dump(final_target_list, f)