def __preprocess_and_save__(self, index):
obj_category = self.list_obj[index]
ins = self.list_instance[index]
obj = self.objnamelist.index(obj_category)
art_status = self.list_status[index]
frame_order = self.list_rank[index]
label = self.list_label[index]
h5_save_path = self.root_dir + '/hdf5/' + obj_category + '/' + ins + '/' + art_status
if (not os.path.exists(h5_save_path)):
os.makedirs(h5_save_path)
h5_save_name = h5_save_path + '/{}.h5'.format(frame_order)
num_parts = self.urdf_dict[obj_category][ins]['num_links']
model_offsets = self.urdf_dict[obj_category][ins]['link']
joint_offsets = self.urdf_dict[obj_category][ins]['joint']
parts_model_point = [None] * num_parts
parts_world_point = [None] * num_parts
parts_target_point = [None] * num_parts
parts_cloud_cam = [None] * num_parts
parts_cloud_world = [None] * num_parts
parts_cloud_canon = [None] * num_parts
parts_cloud_urdf = [None] * num_parts
parts_cloud_norm = [None] * num_parts
parts_world_pos = [None] * num_parts
parts_world_orn = [None] * num_parts
parts_urdf_pos = [None] * num_parts
parts_urdf_orn = [None] * num_parts
parts_urdf_box = [None] * num_parts
parts_model2world = [None] * num_parts
parts_canon2urdf = [None] * num_parts
parts_target_r = [None] * num_parts
parts_target_t = [None] * num_parts
parts_mask = [None] * num_parts
choose_x = [None] * num_parts
choose_y = [None] * num_parts
choose_to_whole = [None] * num_parts
# rgb/depth/label
print('current image: ', self.list_rgb[index])
img = Image.open(self.list_rgb[index])
img = np.array(img) #.astype(np.uint8)
depth = np.array(h5py.File(self.list_depth[index], 'r')['data'])
label = np.array(Image.open(self.list_label[index]))
# pose infos
pose_dict = self.meta_dict[obj_category][ins][art_status][
'frame_{}'.format(frame_order)]
urdf_dict = self.urdf_dict[obj_category][ins]
viewMat = np.array(pose_dict['viewMat']).reshape(4, 4).T
projMat = np.array(pose_dict['projMat']).reshape(4, 4).T
parts_world_pos[0] = np.array([0, 0, 0])
parts_world_orn[0] = np.array([0, 0, 0, 1])
for link in range(0, num_parts):
if link > 0:
parts_world_pos[link] = np.array(
pose_dict['obj'][link - 1][4]).astype(np.float32)
parts_world_orn[link] = np.array(
pose_dict['obj'][link - 1][5]).astype(np.float32)
for link in range(num_parts):
if link == 1 and num_parts == 2:
parts_urdf_pos[link] = np.array(urdf_dict['joint']['xyz'][
link -
1]) # todo, accumulate joints pffsets != link offsets
else:
parts_urdf_pos[link] = -np.array(
urdf_dict['link']['xyz'][link])
parts_urdf_orn[link] = np.array(urdf_dict['link']['rpy'][link])
for k in range(num_parts):
center_world_orn = parts_world_orn[k]
center_world_orn = np.array([
center_world_orn[3], center_world_orn[0], center_world_orn[1],
center_world_orn[2]
])
my_model2world_r = quaternion_matrix(
center_world_orn)[:4, :4] # [w, x, y, z]
my_model2world_t = parts_world_pos[k]
my_model2world_mat = np.copy(my_model2world_r)
for m in range(3):
my_model2world_mat[m, 3] = my_model2world_t[m]
my_world2camera_mat = viewMat
my_camera2clip_mat = projMat
my_model2camera_mat = np.dot(my_world2camera_mat,
my_model2world_mat)
parts_model2world[k] = my_model2world_mat
# depth to cloud data
mask = np.array((label[:, :] < num_parts) & (label[:, :] > -1)).astype(
np.uint8)
mask_whole = np.copy(mask)
for n in range(num_parts):
parts_mask[n] = np.array((label[:, :] == (n))).astype(np.uint8)
choose_to_whole[n] = np.where(parts_mask[n] > 0)
#>>>>>>>>>>------- rendering target pcloud from depth image --------<<<<<<<<<#
# first get projected map
ymap = self.ymap
xmap = self.xmap
h = self.height
w = self.width
u_map = ymap * 2 / w - 1
v_map = (512 - xmap) * 2 / h - 1
v1_map = xmap * 2 / h - 1
w_channel = -depth
projected_map = np.stack(
[u_map * w_channel, v_map * w_channel, depth,
w_channel]).transpose([1, 2, 0])
projected_map1 = np.stack(
[u_map * w_channel, v1_map * w_channel, depth,
w_channel]).transpose([1, 2, 0])
for s in range(num_parts):
x_set, y_set = choose_to_whole[s]
if len(x_set) < 10:
print('data is empty, skipping!!!')
return None
else:
choose_x[s] = x_set
choose_y[s] = y_set
# ---------------> from projected map into target part_cloud
# order: cam->world->canon)
projected_points = projected_map[
choose_x[s][:].astype(np.uint16),
choose_y[s][:].astype(np.uint16), :]
projected_points = np.reshape(projected_points, [-1, 4])
depth_channel = -projected_points[:, 3:4]
cloud_cam = np.dot(
projected_points[:, 0:2] -
np.dot(depth_channel, projMat[0:2, 2:3].T),
np.linalg.pinv(projMat[:2, :2].T))
projected_points1 = projected_map1[
choose_x[s][:].astype(np.uint16),
choose_y[s][:].astype(np.uint16), :]
projected_points1 = np.reshape(projected_points1, [-1, 4])
cloud_cam_real = np.dot(
projected_points1[:, 0:2] -
np.dot(depth_channel, projMat[0:2, 2:3].T),
np.linalg.pinv(projMat[:2, :2].T))
cloud_cam_real = np.concatenate((cloud_cam_real, depth_channel),
axis=1)
cloud_cam = np.concatenate((cloud_cam, depth_channel), axis=1)
cloud_cam_full = np.concatenate(
(cloud_cam, np.ones((cloud_cam.shape[0], 1))), axis=1)
# modify, todo
camera_pose_mat = np.linalg.pinv(viewMat.T)
camera_pose_mat[:3, :] = -camera_pose_mat[:3, :]
cloud_world = np.dot(cloud_cam_full, camera_pose_mat)
cloud_canon = np.dot(cloud_world,
np.linalg.pinv(parts_model2world[s].T))
# canon points should be points coordinates centered in the inertial frame
parts_cloud_cam[s] = cloud_cam_real[:, :3]
parts_cloud_world[s] = cloud_world[:, :3]
parts_cloud_canon[s] = cloud_canon[:, :3]
for k in range(num_parts):
center_joint_orn = parts_urdf_orn[k]
my_canon2urdf_r = euler_matrix(
center_joint_orn[0], center_joint_orn[1],
center_joint_orn[2])[:4, :4] # [w, x, y, z]
my_canon2urdf_t = parts_urdf_pos[k]
my_canon2urdf_mat = my_canon2urdf_r
for m in range(3):
my_canon2urdf_mat[m, 3] = my_canon2urdf_t[m]
part_points_space = np.concatenate(
(parts_cloud_canon[k],
np.ones((parts_cloud_canon[k].shape[0], 1))),
axis=1)
parts_cloud_urdf[k] = np.dot(part_points_space,
my_canon2urdf_mat.T)
#>>>>>>>>>>>>>>> go to PNCS space
for link in range(num_parts):
tight_w = max(parts_cloud_urdf[link][:, 0]) - min(
parts_cloud_urdf[link][:, 0])
tight_l = max(parts_cloud_urdf[link][:, 1]) - min(
parts_cloud_urdf[link][:, 1])
tight_h = max(parts_cloud_urdf[link][:, 2]) - min(
parts_cloud_urdf[link][:, 2])
norm_factor = np.sqrt(1) / np.sqrt(tight_w**2 + tight_l**2 +
tight_h**2)
base_p = np.array([
min(parts_cloud_urdf[link][:, 0]),
min(parts_cloud_urdf[link][:, 1]),
min(parts_cloud_urdf[link][:, 2])
]).reshape(1, 3)
extre_p = np.array([
max(parts_cloud_urdf[link][:, 0]),
max(parts_cloud_urdf[link][:, 1]),
max(parts_cloud_urdf[link][:, 2])
]).reshape(1, 3)
center_p = (extre_p - base_p) / 2 * norm_factor
parts_cloud_norm[link] = (parts_cloud_urdf[link][:, :3] -
base_p) * norm_factor + np.array(
[0.5, 0.5, 0.5]).reshape(
1, 3) - center_p.reshape(1, 3)
x_set_pcloud = np.concatenate(choose_x, axis=0)
y_set_pcloud = np.concatenate(choose_y, axis=0)
# save into h5 for rgb_img, input_pts, mask, correpsonding urdf_points
print('Writing to ', h5_save_name)
hf = h5py.File(h5_save_name, 'w')
hf.create_dataset('rgb', data=img)
hf.create_dataset('mask', data=mask)
cloud_cam = hf.create_group('gt_points')
for part_i, points in enumerate(parts_cloud_cam):
cloud_cam.create_dataset(str(part_i), data=points)
coord_gt = hf.create_group('gt_coords')
for part_i, points in enumerate(parts_cloud_urdf):
coord_gt.create_dataset(str(part_i), data=points)
hf.close()
################# for debug only, let me know if you have questions #################
if self.is_debug:
figure = plt.figure(dpi=200)
ax = plt.subplot(121)
plt.imshow(img)
plt.title('RGB image')
ax1 = plt.subplot(122)
plt.imshow(depth)
plt.title('depth image')
plt.show()
plot3d_pts(
[parts_cloud_cam],
[['part {}'.format(i) for i in range(len(parts_cloud_cam))]],
s=5,
title_name=['camera coords'])
plot3d_pts(
[parts_cloud_world],
[['part {}'.format(i) for i in range(len(parts_cloud_world))]],
s=5,
title_name=['world coords'])
plot3d_pts(
[parts_cloud_canon],
[['part {}'.format(i) for i in range(len(parts_cloud_canon))]],
s=5,
title_name=['canon coords'])
plot3d_pts(
[parts_cloud_urdf],
[['part {}'.format(i) for i in range(len(parts_cloud_urdf))]],
s=5,
title_name=['urdf coords'])
return None
def __preprocess_and_save__(self, index):
obj_category = self.list_obj[index]
ins = self.list_instance[index]
obj = self.objnamelist.index(obj_category)
art_status = self.list_status[index]
frame_order = self.list_rank[index]
label = self.list_label[index]
h5_save_path = (self.root_dir + "/hdf5/" + obj_category + "/" + ins +
"/" + art_status)
if not os.path.exists(h5_save_path):
os.makedirs(h5_save_path)
h5_save_name = h5_save_path + "/{}.h5".format(frame_order)
num_parts = self.urdf_dict[obj_category][ins]["num_links"]
new_num_parts = num_parts - 1
model_offsets = self.urdf_dict[obj_category][ins]["link"]
joint_offsets = self.urdf_dict[obj_category][ins]["joint"]
parts_model_point = [None] * new_num_parts
parts_world_point = [None] * new_num_parts
parts_target_point = [None] * new_num_parts
parts_cloud_cam = [None] * new_num_parts
parts_cloud_world = [None] * new_num_parts
parts_cloud_canon = [None] * new_num_parts
parts_cloud_urdf = [None] * new_num_parts
parts_cloud_norm = [None] * new_num_parts
parts_world_pos = [None] * new_num_parts
parts_world_orn = [None] * new_num_parts
parts_urdf_pos = [None] * new_num_parts
parts_urdf_orn = [None] * new_num_parts
parts_urdf_box = [None] * new_num_parts
parts_model2world = [None] * new_num_parts
parts_canon2urdf = [None] * new_num_parts
parts_target_r = [None] * new_num_parts
parts_target_t = [None] * new_num_parts
parts_mask = [None] * new_num_parts
choose_x = [None] * new_num_parts
choose_y = [None] * new_num_parts
choose_to_whole = [None] * new_num_parts
# rgb/depth/label
print("current image: ", self.list_rgb[index])
img = Image.open(self.list_rgb[index])
img = np.array(img) # .astype(np.uint8)
depth = np.array(h5py.File(self.list_depth[index], "r")["data"])
label = np.array(Image.open(self.list_label[index]))
# pose infos
pose_dict = self.meta_dict[obj_category][ins][art_status][
"frame_{}".format(frame_order)]
urdf_dict = self.urdf_dict[obj_category][ins]
viewMat = np.array(pose_dict["viewMat"]).reshape(4, 4).T
projMat = np.array(pose_dict["projMat"]).reshape(4, 4).T
parts_world_pos[0] = np.array([0, 0, 0])
parts_world_orn[0] = np.array([0, 0, 0, 1])
# SAPIEN: skip the virtual base part
for link in range(1, num_parts):
if link > 0:
parts_world_pos[link - 1] = np.array(
pose_dict["obj"][link - 1][4]).astype(np.float32)
parts_world_orn[link - 1] = np.array(
pose_dict["obj"][link - 1][5]).astype(np.float32)
for link in range(1, num_parts):
if link == 1 and num_parts == 2:
parts_urdf_pos[link] = np.array(urdf_dict["joint"]["xyz"][
link -
1]) # todo, accumulate joints offsets != link offsets
else:
# Only change this branch for SAPIEN data
parts_urdf_pos[link - 1] = -np.array(
urdf_dict["link"]["xyz"][link][0])
parts_urdf_orn[link - 1] = np.array(
urdf_dict["link"]["rpy"][link][0])
for k in range(1, num_parts):
center_world_orn = parts_world_orn[k - 1]
center_world_orn = np.array([
center_world_orn[3],
center_world_orn[0],
center_world_orn[1],
center_world_orn[2],
])
my_model2world_r = quaternion_matrix(
center_world_orn)[:4, :4] # [w, x, y, z]
my_model2world_t = parts_world_pos[k - 1]
my_model2world_mat = np.copy(my_model2world_r)
for m in range(3):
my_model2world_mat[m, 3] = my_model2world_t[m]
my_world2camera_mat = viewMat
my_camera2clip_mat = projMat
my_model2camera_mat = np.dot(my_world2camera_mat,
my_model2world_mat)
parts_model2world[k - 1] = my_model2world_mat
# depth to cloud data
# SAPIEN, throw away the label 0 for virtual base link
mask = np.array((label[:, :] < num_parts) & (label[:, :] > 0)).astype(
np.uint8)
mask_whole = np.copy(mask)
for n in range(1, num_parts):
parts_mask[n - 1] = np.array((label[:, :] == (n))).astype(np.uint8)
choose_to_whole[n - 1] = np.where(parts_mask[n - 1] > 0)
# >>>>>>>>>>------- rendering target pcloud from depth image --------<<<<<<<<<#
# first get projected map
# ymap = self.ymap
# xmap = self.xmap
# h = self.height
# w = self.width
# u_map = ymap * 2 / w - 1
# v_map = (512 - xmap) * 2 / h - 1
# v1_map = xmap * 2 / h - 1
# w_channel = -depth
# projected_map = np.stack(
# [u_map * w_channel, v_map * w_channel, depth, w_channel]
# ).transpose([1, 2, 0])
# projected_map1 = np.stack(
# [u_map * w_channel, v1_map * w_channel, depth, w_channel]
# ).transpose([1, 2, 0])
for s in range(1, num_parts):
x_set, y_set = choose_to_whole[s - 1]
if len(x_set) < 10:
print(ins, art_status, frame_order)
print(
"data is empty, skipping!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
)
return None
else:
choose_x[s - 1] = x_set
choose_y[s - 1] = y_set
# ---------------> from projected map into target part_cloud
# order: cam->world->canon)
# MotionNet Modification
point_num = np.shape(choose_x[s - 1])[0]
cloud_cam = []
cloud_cam_real = []
for i in range(point_num):
x = choose_x[s - 1][i]
y = choose_y[s - 1][i]
y_real = 512 - y
cloud_cam.append([
(y + projMat[0, 2]) * depth[x, y] / projMat[0, 0],
-(x + projMat[1, 2]) * depth[x, y] / (-projMat[1, 1]),
-depth[x, y],
])
cloud_cam_real.append([
(y + projMat[0, 2]) * depth[x, y] / projMat[0, 0],
-(x + projMat[1, 2]) * depth[x, y] / (-projMat[1, 1]),
-depth[x, y],
])
# cloud_cam.append(
# [
# (x + projMat[0, 2]) * depth[x, y_real] / projMat[0, 0],
# -(y_real + projMat[1, 2]) * depth[x, y_real] / (-projMat[1, 1]),
# -depth[x, y_real],
# ]
# )
# pdb.set_trace()
# print(cloud_cam)
# break
cloud_cam = np.array(cloud_cam)
cloud_cam_real = np.array(cloud_cam_real)
# pdb.set_trace()
# MotionNet Modification
# projected_points = projected_map[choose_x[s-1][:].astype(np.uint16), choose_y[s-1][:].astype(np.uint16), :]
# projected_points = np.reshape(projected_points, [-1, 4])
# depth_channel = - projected_points[:, 3:4]
# cloud_cam = np.dot(projected_points[:, 0:2] - np.dot(depth_channel, projMat[0:2, 2:3].T), np.linalg.pinv(projMat[:2, :2].T))
# projected_points1 = projected_map1[choose_x[s-1][:].astype(np.uint16), choose_y[s-1][:].astype(np.uint16), :]
# projected_points1 = np.reshape(projected_points1, [-1, 4])
# cloud_cam_real = np.dot(projected_points1[:, 0:2] - np.dot(depth_channel, projMat[0:2, 2:3].T), np.linalg.pinv(projMat[:2, :2].T))
# cloud_cam_real = np.concatenate((cloud_cam_real, depth_channel), axis=1)
# cloud_cam = np.concatenate((cloud_cam, depth_channel), axis=1)
cloud_cam_full = np.concatenate(
(cloud_cam, np.ones((cloud_cam.shape[0], 1))), axis=1)
# modify, todo
camera_pose_mat = np.linalg.pinv(viewMat.T)
# MotionNet modification
# camera_pose_mat[:3, :] = -camera_pose_mat[:3, :]
cloud_world = np.dot(cloud_cam_full, camera_pose_mat)
cloud_canon = np.dot(cloud_world,
np.linalg.pinv(parts_model2world[s - 1].T))
# canon points should be points coordinates centered in the inertial frame
parts_cloud_cam[s - 1] = cloud_cam_real[:, :3]
parts_cloud_world[s - 1] = cloud_world[:, :3]
parts_cloud_canon[s - 1] = cloud_canon[:, :3]
for k in range(1, num_parts):
center_joint_orn = parts_urdf_orn[k - 1]
my_canon2urdf_r = euler_matrix(
center_joint_orn[0], center_joint_orn[1],
center_joint_orn[2])[:4, :4] # [w, x, y, z]
my_canon2urdf_t = parts_urdf_pos[k - 1]
my_canon2urdf_mat = my_canon2urdf_r
for m in range(3):
my_canon2urdf_mat[m, 3] = my_canon2urdf_t[m]
part_points_space = np.concatenate(
(
parts_cloud_canon[k - 1],
np.ones((parts_cloud_canon[k - 1].shape[0], 1)),
),
axis=1,
)
parts_cloud_urdf[k - 1] = np.dot(part_points_space,
my_canon2urdf_mat.T)
# >>>>>>>>>>>>>>> go to NPCS space
# NPCS here is not stored into hdf5
for link in range(1, num_parts):
tight_w = max(parts_cloud_urdf[link - 1][:, 0]) - min(
parts_cloud_urdf[link - 1][:, 0])
tight_l = max(parts_cloud_urdf[link - 1][:, 1]) - min(
parts_cloud_urdf[link - 1][:, 1])
tight_h = max(parts_cloud_urdf[link - 1][:, 2]) - min(
parts_cloud_urdf[link - 1][:, 2])
norm_factor = np.sqrt(1) / np.sqrt(tight_w**2 + tight_l**2 +
tight_h**2)
base_p = np.array([
min(parts_cloud_urdf[link - 1][:, 0]),
min(parts_cloud_urdf[link - 1][:, 1]),
min(parts_cloud_urdf[link - 1][:, 2]),
]).reshape(1, 3)
extre_p = np.array([
max(parts_cloud_urdf[link - 1][:, 0]),
max(parts_cloud_urdf[link - 1][:, 1]),
max(parts_cloud_urdf[link - 1][:, 2]),
]).reshape(1, 3)
center_p = (extre_p - base_p) / 2 * norm_factor
parts_cloud_norm[link - 1] = (
(parts_cloud_urdf[link - 1][:, :3] - base_p) * norm_factor +
np.array([0.5, 0.5, 0.5]).reshape(1, 3) -
center_p.reshape(1, 3))
x_set_pcloud = np.concatenate(choose_x, axis=0)
y_set_pcloud = np.concatenate(choose_y, axis=0)
# save into h5 for rgb_img, input_pts, mask, correpsonding urdf_points
print("Writing to ", h5_save_name)
hf = h5py.File(h5_save_name, "w")
hf.create_dataset("rgb", data=img)
hf.create_dataset("mask", data=mask)
cloud_cam = hf.create_group("gt_points")
for part_i, points in enumerate(parts_cloud_cam):
cloud_cam.create_dataset(str(part_i), data=points)
coord_gt = hf.create_group("gt_coords")
for part_i, points in enumerate(parts_cloud_urdf):
coord_gt.create_dataset(str(part_i), data=points)
hf.close()
################# for debug only, let me know if you have questions #################
if self.is_debug:
figure = plt.figure(dpi=200)
ax = plt.subplot(121)
plt.imshow(img)
plt.title('RGB image')
ax1 = plt.subplot(122)
plt.imshow(depth)
plt.title('depth image')
plt.show()
plot3d_pts(
[parts_cloud_cam],
[['part {}'.format(i) for i in range(len(parts_cloud_cam))]],
s=5,
title_name=['camera coords'])
plot3d_pts(
[parts_cloud_world],
[['part {}'.format(i) for i in range(len(parts_cloud_world))]],
s=5,
title_name=['world coords'])
plot3d_pts(
[parts_cloud_canon],
[['part {}'.format(i) for i in range(len(parts_cloud_canon))]],
s=5,
title_name=['canon coords'])
plot3d_pts(
[parts_cloud_urdf],
[['part {}'.format(i) for i in range(len(parts_cloud_urdf))]],
s=5,
title_name=['urdf coords'])
return None
def get_part_bounding_box(my_dir, test_ins, args, viz=False):
pts_m = {}
bbox3d_all = {}
start = time.time()
m_file = my_dir + '/shape2motion/pickle/{}_pts.pkl'.format(args.item)
c_file = my_dir + '/shape2motion/pickle/{}_corners.pkl'.format(args.item)
n_file = my_dir + '/shape2motion/pickle/{}.pkl'.format(args.item)
if args.process:
for item in os.listdir(test_ins):
print('now fetching for item {}'.format(item))
pts, nf, cpts = get_model_pts(root_dset, args.item, item)
pt_ii = []
for p, pt in enumerate(pts):
pt_s = np.concatenate(pt, axis=0)
np.random.shuffle(pt_s)
# pt_s = pt_s[::20, :]
pt_ii.append(pt_s)
print('We have {} pts'.format(pt_ii[p].shape[0]))
if pt_ii is not []:
pts_m[item] = pt_ii
else:
print('!!!!! {} model loading is wrong'.format(item))
end_t = time.time()
with open(m_file, 'wb') as f:
pickle.dump(pts_m, f)
else:
with open(m_file, 'rb') as f:
pts_m = pickle.load(f)
with open(c_file, 'rb') as f:
pts_c = pickle.load(f)
with open(n_file, 'rb') as f:
pts_n = pickle.load(f)
# for item in list(pts_m.keys()):
for item in test_ins:
pts = pts_m[item]
norm_factors = pts_n[item]
norm_corners = pts_c[item]
pt_ii = []
bbox3d_per_part = []
for p, pt in enumerate(
pts
): # todo: assume we are dealing part-nocs, so model pts are processed
norm_factor = norm_factors[p + 1]
norm_corner = norm_corners[p + 1]
nocs_corner = np.copy(
norm_corner) # copy is very important, as they are
print('norm_corner:\n', norm_corner)
pt_nocs = (pt - norm_corner[0]) * norm_factor + np.array(
[0.5, 0.5, 0.5]).reshape(1, 3) - 0.5 * (
norm_corner[1] - norm_corner[0]) * norm_factor
nocs_corner[0] = np.array([0.5, 0.5, 0.5]).reshape(
1,
3) - 0.5 * (norm_corner[1] - norm_corner[0]) * norm_factor
nocs_corner[1] = np.array([0.5, 0.5, 0.5]).reshape(
1,
3) + 0.5 * (norm_corner[1] - norm_corner[0]) * norm_factor
bbox3d_per_part.append(nocs_corner)
np.random.shuffle(pt_nocs)
pt_ii.append(pt_nocs[0:2000, :]) # sub-sampling
print('We have {} pts'.format(pt_ii[p].shape[0]))
if pt_ii is not []:
pts_m[item] = pt_ii
else:
print('!!!!! {} model loading is wrong'.format(item))
assert bbox3d_per_part != []
bbox3d_all[item] = bbox3d_per_part
end_t = time.time()
print('It takes {} seconds to get: \n'.format(end_t - start),
list(pts_m.keys()))
if viz:
plot3d_pts([pts_m['0001']], [['part 0', 'part 1']],
s=5,
title_name=['sampled model pts'],
dpi=200)
return bbox3d_all, pts_m
def get_all_objs(root_dset,
obj_category,
item,
obj_file_list=None,
offsets=None,
is_debug=False,
verbose=False):
"""
offsets is usually 0, but sometimes could be [x, y, z] in array 1*3, it could be made to a K*3 array if necessary
"""
norm_factors = []
pts_list = []
name_list = []
target_dir = root_dset + '/objects/' + obj_category + '/' + item
offset = 0
if obj_file_list is None:
for k, obj_file in enumerate(glob.glob(target_dir +
'/part_objs/*.obj')):
if offsets is not None:
offset = offsets[k:k + 1, :]
if is_debug:
print('obj_file is: ', obj_file)
try:
tm = trimesh.load(obj_file)
vertices_obj = np.array(tm.vertices)
except:
dict_mesh, _, _, _ = load_model_split(obj_file)
vertices_obj = np.concatenate(dict_mesh['v'], axis=0)
pts_list.append(vertices_obj + offset)
name_obj = obj_file.split('.')[0].split('/')[-1]
name_list.append(name_obj)
else:
for k, obj_files in enumerate(obj_file_list):
if offsets is not None:
offset = offsets[k:k + 1, :]
if obj_files is not None and not isinstance(obj_files, list):
try:
tm = trimesh.load(obj_files)
vertices_obj = np.array(tm.vertices)
except:
dict_mesh, _, _, _ = load_model_split(obj_files)
vertices_obj = np.concatenate(dict_mesh['v'], axis=0)
pts_list.append(vertices_obj + offset)
name_obj = obj_files.split('.')[0].split('/')[-1]
name_list.append(
name_obj) # which should follow the right order
elif isinstance(obj_files, list):
if verbose:
print('{} part has {} obj files'.format(k, len(obj_files)))
part_pts = []
name_objs = []
for obj_file in obj_files:
if obj_file is not None and not isinstance(obj_file, list):
try:
tm = trimesh.load(obj_file)
vertices_obj = np.array(tm.vertices)
except:
dict_mesh, _, _, _ = load_model_split(obj_file)
vertices_obj = np.concatenate(dict_mesh['v'],
axis=0)
name_obj = obj_file.split('.')[0].split('/')[-1]
name_objs.append(name_obj)
part_pts.append(vertices_obj)
part_pts_whole = np.concatenate(part_pts, axis=0)
pts_list.append(part_pts_whole + offset)
name_list.append(name_objs) # which should follow the right
if is_debug:
print('name_list is: ', name_list)
parts_a = []
parts_a = pts_list
parts_b = [None] * len(obj_file_list)
# dof_rootd_Aa001_r.obj dof_rootd_Aa002_r.obj none_motion.obj
# bike: part2: 'dof_Aa001_Ca001_r', 'dof_rootd_Aa001_r'
if obj_category == 'bike':
part0 = []
part1 = []
part2 = []
part0 = pts_list
for i, name_obj in enumerate(name_list):
if name_obj in ['dof_Aa001_Ca001_r', 'dof_rootd_Aa001_r']:
print('part 2 adding ', name_obj)
part2.append(pts_list[i])
else:
print('part 1 adding ', name_obj)
part1.append(pts_list[i])
parts = [part0, part1, part2]
elif obj_category == 'eyeglasses':
for i, name_obj in enumerate(name_list):
if name_obj in ['none_motion']:
parts_b[0] = []
parts_b[0].append(pts_list[i])
if name_obj in ['dof_rootd_Aa001_r']:
parts_b[1] = []
parts_b[1].append(pts_list[i])
elif name_obj in ['dof_rootd_Aa002_r']:
parts_b[2] = []
parts_b[2].append(pts_list[i])
parts = [parts_a] + parts_b
else:
parts_a = []
parts_a = pts_list
parts_b = [None] * len(name_list)
for i, name_obj in enumerate(name_list):
parts_b[i] = []
parts_b[i].append(pts_list[i])
parts = [parts_a] + parts_b
corner_pts = [None] * len(parts)
for j in range(len(parts)):
if is_debug:
print('Now checking ', j)
part_gts = np.concatenate(parts[j], axis=0)
print('part_gts: ', part_gts.shape)
tight_w = max(part_gts[:, 0]) - min(part_gts[:, 0])
tight_l = max(part_gts[:, 1]) - min(part_gts[:, 1])
tight_h = max(part_gts[:, 2]) - min(part_gts[:, 2])
corner_pts[j] = np.amin(part_gts, axis=1)
norm_factor = np.sqrt(1) / np.sqrt(tight_w**2 + tight_l**2 +
tight_h**2)
norm_factors.append(norm_factor)
corner_pt_left = np.amin(part_gts, axis=0, keepdims=True)
corner_pt_right = np.amax(part_gts, axis=0, keepdims=True)
corner_pts[j] = [corner_pt_left, corner_pt_right
] # [index][left/right][x, y, z], numpy array
if is_debug:
print('Group {} has {} points with shape {}'.format(
j, len(corner_pts[j]), corner_pts[j][0].shape))
if verbose:
plot3d_pts([[part_gts[::2]]], ['model pts'],
s=15,
title_name=['GT model pts {}'.format(j)],
sub_name=str(j))
# for k in range(len(parts[j])):
# plot3d_pts([[parts[j][k][::2]]], ['model pts of part {}'.format(k)], s=15, title_name=['GT model pts'], sub_name=str(k))
return parts[1:], norm_factors, corner_pts
def evaluateModelRotation(Rotation,
SourceHom,
TargetHom,
PassThreshold,
rt_ref=None,
viz_normal=False):
SourceCentroid = np.mean(SourceHom[:3, :], axis=1)
TargetCentroid = np.mean(TargetHom[:3, :], axis=1)
nPoints = SourceHom.shape[1]
CenteredSource = SourceHom[:3, :] - np.tile(SourceCentroid,
(nPoints, 1)).transpose()
CenteredTarget = TargetHom[:3, :] - np.tile(TargetCentroid,
(nPoints, 1)).transpose()
stdSource = np.sqrt(np.var(CenteredSource[:3, :], axis=1).sum())
stdTarget = np.sqrt(np.var(CenteredTarget[:3, :], axis=1).sum())
CenteredSource = CenteredSource / stdSource
CenteredTarget = CenteredTarget / stdTarget
origin = np.zeros((1, 3))
if viz_normal:
if rt_ref is not None:
RotatedSource = np.matmul(rt_ref[:3, :3], CenteredSource)
plot3d_pts([[
RotatedSource[:3].transpose(), CenteredTarget[:3].transpose(),
origin
]], [['GT rotated source', 'target', 'origin']],
s=100,
title_name=['normalized source and target pts'],
color_channel=None,
save_fig=False,
sub_name='default')
plot3d_pts([[
CenteredSource[:3].transpose(), CenteredTarget[:3].transpose(),
origin
]], [['source', 'target', 'origin']],
s=100,
title_name=['normalized source and target pts'],
color_channel=None,
save_fig=False,
sub_name='default')
RotatedSource = np.matmul(Rotation, CenteredSource)
if viz_normal:
plot3d_pts([[
RotatedSource[:3].transpose(), CenteredTarget[:3].transpose(),
origin
]], [['Pred rotated source', 'target', 'origin']],
s=100,
title_name=['normalized source and target pts'],
color_channel=None,
save_fig=False,
sub_name='default')
Diff = CenteredTarget - RotatedSource
ResidualVec = np.linalg.norm(Diff[:3, :], axis=0)
Residual = np.linalg.norm(ResidualVec)
InlierIdx = np.where(ResidualVec < PassThreshold)
nInliers = np.count_nonzero(InlierIdx)
# Residual = np.linalg.norm(ResidualVec[InlierIdx[0]]) # todo
InlierRatio = nInliers / SourceHom.shape[1]
return Residual, InlierRatio, InlierIdx[0]
def svd_pts(SourceHom,
TargetHom,
raw_axis=None,
rotated_axis=None,
viz_sample=False,
index=0):
# Copy of original paper is at: http://web.stanford.edu/class/cs273/refs/umeyama.pdf
SourceCentroid = np.mean(SourceHom[:3, :], axis=1)
TargetCentroid = np.mean(TargetHom[:3, :], axis=1)
nPoints = SourceHom.shape[1]
# pre-centering
CenteredSource = SourceHom[:3, :] - np.tile(SourceCentroid,
(nPoints, 1)).transpose()
CenteredTarget = TargetHom[:3, :] - np.tile(TargetCentroid,
(nPoints, 1)).transpose()
# pre-scaling
stdSource = np.sqrt(np.var(CenteredSource[:3, :], axis=1).sum())
stdTarget = np.sqrt(np.var(CenteredTarget[:3, :], axis=1).sum())
try:
CenteredSource = CenteredSource / stdSource
CenteredTarget = CenteredTarget / stdTarget
except:
CenteredSource = CenteredSource
CenteredTarget = CenteredTarget
origin = np.zeros((1, 3))
if viz_sample:
if raw_axis is not None:
# raw_axis[:3, 0:1].transpose(), rotated_axis[:3, 0:1].transpose(), 'axis', 'rotated axis',
plot3d_pts(
[[
CenteredSource[:3].transpose(),
CenteredTarget[:3].transpose(),
raw_axis.transpose(),
rotated_axis.transpose(), origin
]], [[
'source', 'target', 'raw_axis_point', 'rotated_axis_point',
'origin'
]],
s=100,
title_name=[
'{}th iteration points for coords descent'.format(index)
],
color_channel=None,
save_fig=False,
sub_name='default')
if raw_axis is not None:
CenteredSource = np.concatenate([CenteredSource, raw_axis[:3]], axis=1)
CenteredTarget = np.concatenate([CenteredTarget, rotated_axis[:3]],
axis=1)
nPoints = nPoints + raw_axis.shape[1]
CovMatrix = np.matmul(CenteredTarget,
np.transpose(CenteredSource)) / nPoints
if np.isnan(CovMatrix).any():
print('nPoints:', nPoints)
print(SourceHom.shape)
print(TargetHom.shape)
return None, None, None
U, D, Vh = np.linalg.svd(CovMatrix, full_matrices=True)
d = (np.linalg.det(U) * np.linalg.det(Vh)) < 0.0
if d:
D[-1] = -D[-1]
U[:, -1] = -U[:, -1]
return U, D, Vh
def estimateSimilarityUmeyamaCoords(SourceHom0, TargetHom0, SourceHom1, TargetHom1, joint_axis, joint_pts=None, rt_ref=[None, None], rt_pre=[None, None], \
viz=False, viz_ransac=False, viz_sample=False, use_jt_pts=False, use_ext_rot=False, verbose=False, index=0):
"""
SourceHom0: [4, 5]
joint_pts : [4, 5]
joint_axis: [4, 1]
"""
U, D0, Vh = svd_pts(SourceHom0, TargetHom0) #
R0 = np.matmul(U, Vh).T # Transpose is the one that works
U, D1, Vh = svd_pts(SourceHom1, TargetHom1) #
R1 = np.matmul(U, Vh).T #
# begin EM
max_iter = 100
# max_iter = 1 # todo
StopThreshold = 2 * np.cos(0.5 / 180 * np.pi)
if viz_sample:
plot3d_pts([[
SourceHom0[:3].transpose(), SourceHom1[:3].transpose(),
TargetHom0[:3].transpose(), TargetHom1[:3].transpose(),
joint_pts[:3].transpose()
]], [['source0', 'source1', 'target0', 'target1', 'joint_points']],
s=100,
title_name=['sampled points'],
color_channel=None,
save_fig=False,
sub_name='default')
joint_axis_tiled0 = np.tile(joint_axis, (1, int(SourceHom0.shape[1] / 5)))
joint_axis_tiled1 = np.tile(joint_axis, (1, int(SourceHom1.shape[1] / 5)))
# joint_axis_tiled0 = np.tile(joint_axis, (1, int(SourceHom0.shape[1])))
# joint_axis_tiled1 = np.tile(joint_axis, (1, int(SourceHom1.shape[1])))
if use_ext_rot and rt_pre[0] is not None:
# print('using external rotation')
R0 = rt_pre[0][:3, :3].T
R1 = rt_pre[1][:3, :3].T
else:
r_list = [[R0], [R1]]
for i in range(max_iter):
rotated_axis = np.matmul(R0.T, joint_axis_tiled1[:3]) # [3, 1]
U, D1, Vh = svd_pts(SourceHom1,
TargetHom1,
joint_axis_tiled1,
rotated_axis,
viz_sample=viz_sample,
index=2 * i)
R1_new = np.matmul(U, Vh).T
rotated_axis = np.matmul(R1_new.T, joint_axis_tiled0[:3])
U, D0, Vh = svd_pts(SourceHom0,
TargetHom0,
joint_axis_tiled0,
rotated_axis,
viz_sample=viz_sample,
index=2 * i + 1)
R0_new = np.matmul(U, Vh).T
eigen_sum0 = np.trace(np.matmul(R0_new.T, R0)) - 1
eigen_sum1 = np.trace(np.matmul(R1_new.T, R1)) - 1
R0 = R0_new
R1 = R1_new
r_list[0].append(R0)
r_list[1].append(R1)
if eigen_sum0 > StopThreshold and eigen_sum1 > StopThreshold:
# if verbose:
# print('Algorithm converges at {}th iteration for Coordinate Descent'.format(i))
break
if viz_ransac and index < 10: # and SourceHom0.shape[1]>5:
ang_dis_list = [[], []]
for j in range(2):
q_gt = quaternion_from_matrix(rt_ref[j][:3, :3])
for rot_iter in r_list[j]:
q_iter = quaternion_from_matrix(rot_iter.T)
ang_dis = 2 * np.arccos(sum(q_iter * q_gt)) * 180 / np.pi
if ang_dis > 180:
ang_dis = 360 - ang_dis
ang_dis_list[j].append(ang_dis)
fig = plt.figure(dpi=200)
for j in range(2):
ax = plt.subplot(1, 2, j + 1)
plt.plot(range(len(ang_dis_list[j])), ang_dis_list[j])
plt.xlabel('iteration')
plt.ylabel('rotation error')
plt.title('{}th sampling part {}'.format(index, j))
plt.show()
Rs = [R0, R1]
if use_jt_pts:
if viz_sample:
plot3d_pts([[
SourceHom0[:3].transpose(), SourceHom1[:3].transpose(),
TargetHom0[:3].transpose(), TargetHom1[:3].transpose(),
joint_pts[:3].transpose()
]], [['source0', 'source1', 'target0', 'target1', 'joint_points']],
s=100,
title_name=['sampled points'],
color_channel=None,
save_fig=False,
sub_name='default')
final_scale, Ts, OutTrans = compute_scale_translation(
[SourceHom0, SourceHom1], [TargetHom0, TargetHom1], Rs, joint_pts)
if verbose:
print("scale by adding joints are \n: {}".format(final_scale))
else:
if viz_sample:
plot3d_pts([[
SourceHom0[:3].transpose(), SourceHom1[:3].transpose(),
TargetHom0[:3].transpose(), TargetHom1[:3].transpose()
]], [['source0', 'source1', 'target0', 'target1']],
s=100,
title_name=['points after sampling'],
color_channel=None,
save_fig=False,
sub_name='default')
final_scale0, T0, OutTrans0 = est_ST(SourceHom0, TargetHom0, D0, Rs[0])
final_scale1, T1, OutTrans1 = est_ST(SourceHom1, TargetHom1, D1, Rs[1])
final_scale = [final_scale0, final_scale1]
Ts = [T0, T1]
OutTrans = [OutTrans0, OutTrans1]
if verbose:
print("scale by direct solving per part are \n: {}".format(
final_scale))
return final_scale, Rs, Ts, OutTrans