def get_object_pose_in_world(object_pose, camera_pose, fat_world_pose=None, type='quat'):
# Returns in cm
object_pose_matrix = np.zeros((4,4))
object_pose_matrix[:3,:3] = RT_transform.quat2mat(get_wxyz_quaternion(object_pose['quaternion_xyzw']))
object_pose_matrix[:,3] = object_pose['location'] + [1]
# camera_pose_matrix = np.zeros((4,4))
# camera_pose_matrix[:3, :3] = RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe']))
# camera_pose_matrix[:, 3] = camera_pose['location_worldframe'] + [1]
camera_pose_matrix = get_camera_pose_in_world(camera_pose, fat_world_pose, type='rot', cam_to_body=None)
object_pose_world = np.matmul(camera_pose_matrix, object_pose_matrix)
# object_pose_world = np.matmul(np.linalg.inv(camera_pose_matrix), object_pose_matrix)
# scale = np.array([[0.01,0,0,0],[0,0.01,0,0],[0,0,0.01,0],[0,0,0,1]])
# object_pose_world = np.matmul(scale, object_pose_world)
if fat_world_pose is not None:
fat_world_matrix = np.zeros((4,4))
fat_world_matrix[:3,:3] = RT_transform.quat2mat(get_wxyz_quaternion(fat_world_pose['quaternion_xyzw']))
fat_world_matrix[:,3] = fat_world_pose['location'] + [1]
object_pose_world = np.matmul(fat_world_matrix, object_pose_world)
# print(object_pose_world)
if type == 'quat':
quat = RT_transform.mat2quat(object_pose_world[:3, :3]).tolist()
return object_pose_world[:3,3].tolist(), get_xyzw_quaternion(quat)
def get_camera_pose_in_world(camera_pose, fat_world_pose=None, type='quat', cam_to_body=None):
# print(camera_pose)
# this matrix gives world to camera transform
camera_pose_matrix = np.zeros((4,4))
# camera_rotation = np.linalg.inv(RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe'])))
camera_rotation = RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe']))
camera_pose_matrix[:3, :3] = camera_rotation
camera_location = [i for i in camera_pose['location_worldframe']]
# camera_pose_matrix[:, 3] = np.matmul(-1 * camera_rotation, camera_location).tolist() + [1]
camera_pose_matrix[:, 3] = camera_location + [1]
# # make it camera to world
# camera_pose_matrix = np.linalg.inv(camera_pose_matrix)
if cam_to_body is not None:
camera_pose_matrix = np.matmul(camera_pose_matrix, np.linalg.inv(cam_to_body))
if fat_world_pose is not None:
fat_world_matrix = np.zeros((4,4))
fat_world_matrix[:3,:3] = RT_transform.quat2mat(get_wxyz_quaternion(fat_world_pose['quaternion_xyzw']))
fat_world_matrix[:,3] = fat_world_pose['location'] + [1]
camera_pose_world = np.matmul(fat_world_matrix, camera_pose_matrix)
else:
camera_pose_world = camera_pose_matrix
# make it camera to world
# camera_pose_world = np.linalg.inv(camera_pose_world)
if type == 'quat':
quat = RT_transform.mat2quat(camera_pose_world[:3, :3]).tolist()
return camera_pose_world[:3,3], get_xyzw_quaternion(quat)
elif type == 'rot':
return camera_pose_world
def render(self,
r,
t,
light_position,
light_intensity,
class_name,
brightness_k=0,
r_type='quat'):
'''
:param r:
:param t:
:param light_position:
:param light_intensity:
:param brightness_k: choose which brightness in __init__
:param r_type:
:return:
'''
if r_type == 'quat':
R = RT_transform.quat2mat(r)
elif r_type == 'mat':
R = r
self.class_name = class_name
self.brightness_k = brightness_k
self.render_kernels[
self.class_name][brightness_k]['u_view'] = self._get_view_mtx(
R, t)
self.render_kernels[
self.class_name][brightness_k]['u_light_position'] = light_position
self.render_kernels[self.class_name][brightness_k]['u_normal'] = \
np.array(np.matrix(np.dot(
self.render_kernels[self.class_name][brightness_k]['u_view'].reshape(4,4),
self.render_kernels[self.class_name][brightness_k]['u_model'].reshape(4,4))).I.T)
self.render_kernels[self.class_name][brightness_k][
'u_light_intensity'] = light_intensity
app.run(framecount=0, framerate=0)
rgb_buffer = np.zeros((self.height, self.width, 4), dtype=np.float32)
gl.glReadPixels(0, 0, self.width, self.height, gl.GL_RGBA, gl.GL_FLOAT,
rgb_buffer)
rgb_gl = np.copy(rgb_buffer)
rgb_gl.shape = 480, 640, 4
rgb_gl = rgb_gl[::-1, :]
rgb_gl = np.round(rgb_gl[:, :, :3] * 255).astype(
np.uint8) # Convert to [0, 255]
rgb_gl = rgb_gl[:, :, [2, 1, 0]]
depth_buffer = np.zeros((self.height, self.width), dtype=np.float32)
gl.glReadPixels(0, 0, self.width, self.height, gl.GL_DEPTH_COMPONENT,
gl.GL_FLOAT, depth_buffer)
depth_gl = np.copy(depth_buffer)
depth_gl.shape = 480, 640
depth_gl = depth_gl[::-1, :]
depth_bg = depth_gl == 1
depth_gl = 2 * self.zFar * self.zNear / (self.zFar + self.zNear -
(self.zFar - self.zNear) *
(2 * depth_gl - 1))
depth_gl[depth_bg] = 0
return rgb_gl, depth_gl
def gen_poses():
pz = np.array([0, 0, 1])
pose_dict, quat_stat, trans_stat, new_points = stat_poses()
observed_prefix_list = ['{:06d}'.format(i + 1) for i in range(NUM_IMAGES)]
sel_classes = classes
observed_pose_dict = {
cls_name: np.zeros((NUM_IMAGES, 7))
for cls_name in sel_classes
}
for cls_i, cls_name in enumerate(sel_classes):
# uncomment here to only generate data for ape
# if cls_name not in ['ape']:
# continue
src_quat_mean = quat_stat[cls_name]['quat_mean']
src_quat_std = quat_stat[cls_name]['quat_std']
src_trans_mean = trans_stat[cls_name]['trans_mean']
src_trans_std = trans_stat[cls_name]['trans_std']
deg_max = new_points[cls_name]['angle_max']
for i in tqdm(range(NUM_IMAGES)):
observed_prefix = observed_prefix_list[i]
# randomly generate a pose
tgt_quat = np.random.normal(0, 1, 4)
tgt_quat = tgt_quat / np.linalg.norm(tgt_quat)
if tgt_quat[0] < 0:
tgt_quat *= -1
tgt_trans = np.random.normal(src_trans_mean, src_trans_std)
tgt_rot_m = se3.quat2mat(tgt_quat)
new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3, ))
pz_mean = new_points[cls_name]['pz_mean']
deg = angle(new_pz, pz_mean)
# r_dist, t_dist = calc_rt_dist_q(tgt_quat, src_quat, tgt_trans, src_trans)
transform = np.matmul(K, tgt_trans.reshape(3, 1))
center_x = float(transform[0] / transform[2])
center_y = float(transform[1] / transform[2])
count = 0
while (deg > deg_max
or not (48 < center_x < (640 - 48) and 48 < center_y <
(480 - 48))):
# randomly generate a pose
tgt_quat = np.random.normal(0, 1, 4)
tgt_quat = tgt_quat / np.linalg.norm(tgt_quat)
if tgt_quat[0] < 0:
tgt_quat *= -1
tgt_trans = np.random.normal(src_trans_mean, src_trans_std)
tgt_rot_m = se3.quat2mat(tgt_quat)
new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3, ))
pz_mean = new_points[cls_name]['pz_mean']
deg = angle(new_pz, pz_mean)
transform = np.matmul(K, tgt_trans.reshape(3, 1))
center_x = float(transform[0] / transform[2])
center_y = float(transform[1] / transform[2])
count += 1
if count % 100 == 0:
print(
observed_prefix, cls_name, count,
"deg < deg_max={}: {}, 48 < center_x < (640-48): {}, 48 < center_y < (480-48): {}"
.format(deg_max, deg <= deg_max, 48 < center_x <
(640 - 48), 48 < center_y < (480 - 48)))
tgt_pose_q = np.zeros((7, ), dtype='float32')
tgt_pose_q[:4] = tgt_quat
tgt_pose_q[4:] = tgt_trans
observed_pose_dict[cls_name][i, :] = tgt_pose_q
# write pose
poses_file = os.path.join(pose_dir, 'LM6d_ds_train_observed_pose_all.pkl')
with open(poses_file, 'wb') as f:
cPickle.dump(observed_pose_dict, f, 2)
def main():
gen_images = True
for class_idx, class_name in idx2class.items():
train_pair = []
val_pair = []
print("start ", class_idx, class_name)
if class_name in ["__back_ground__"]:
continue
# uncomment here to only generate data for ape
# if class_name not in ['ape']:
# continue
if gen_images:
# init renderer
model_dir = os.path.join(LM6d_root, "models", class_name)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
for set_type in ["all"]:
with open(
os.path.join(observed_set_root, "{}_{}.txt".format(class_name, "all")),
"r",
) as f:
all_observed_list = [x.strip() for x in f.readlines()]
# with open(
# os.path.join(observed_set_root, '{}_{}.txt'.format(
# class_name, 'train')), 'r') as f:
# train_observed_list = [x.strip() for x in f.readlines()]
with open(
os.path.join(observed_set_root, "{}_{}.txt".format(class_name, "test")),
"r",
) as f:
test_observed_list = [x.strip() for x in f.readlines()]
with open(rendered_pose_path.format(set_type, class_name)) as f:
str_rendered_pose_list = [x.strip().split(" ") for x in f.readlines()]
rendered_pose_list = np.array(
[[float(x) for x in each_pose] for each_pose in str_rendered_pose_list]
)
rendered_per_observed = 10
assert len(rendered_pose_list) == 10 * len(
all_observed_list
), "{} vs {}".format(len(rendered_pose_list), len(all_observed_list))
for idx, observed_index in enumerate(tqdm(all_observed_list)):
video_name, observed_prefix = observed_index.split("/")
rendered_dir = os.path.join(rendered_root_dir, class_name)
mkdir_if_missing(rendered_dir)
for inner_idx in range(rendered_per_observed):
if gen_images:
# if gen_images and observed_index in test_observed_list and inner_idx == 0: # only generate my_val_v{} # noqa:E501
image_file = os.path.join(
rendered_dir,
"{}_{}-color.png".format(observed_prefix, inner_idx),
)
depth_file = os.path.join(
rendered_dir,
"{}_{}-depth.png".format(observed_prefix, inner_idx),
)
rendered_idx = idx * rendered_per_observed + inner_idx
pose_rendered_q = rendered_pose_list[rendered_idx]
rgb_gl, depth_gl = render_machine.render(
pose_rendered_q[:4], pose_rendered_q[4:]
)
rgb_gl = rgb_gl.astype("uint8")
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(image_file, rgb_gl)
cv2.imwrite(depth_file, depth_gl)
pose_rendered_file = os.path.join(
rendered_dir,
"{}_{}-pose.txt".format(observed_prefix, inner_idx),
)
text_file = open(pose_rendered_file, "w")
text_file.write("{}\n".format(class_idx))
pose_rendered_m = np.zeros((3, 4))
pose_rendered_m[:, :3] = RT_transform.quat2mat(
pose_rendered_q[:4]
)
pose_rendered_m[:, 3] = pose_rendered_q[4:]
pose_ori_m = pose_rendered_m
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
pose_ori_m[0, 0],
pose_ori_m[0, 1],
pose_ori_m[0, 2],
pose_ori_m[0, 3],
pose_ori_m[1, 0],
pose_ori_m[1, 1],
pose_ori_m[1, 2],
pose_ori_m[1, 3],
pose_ori_m[2, 0],
pose_ori_m[2, 1],
pose_ori_m[2, 2],
pose_ori_m[2, 3],
)
text_file.write(pose_str)
if observed_index in test_observed_list:
if inner_idx == 0:
val_pair.append(
"{} {}/{}_{}".format(
observed_index,
class_name,
observed_prefix,
inner_idx,
)
)
else:
train_pair.append(
"{} {}/{}_{}".format(
observed_index, class_name, observed_prefix, inner_idx
)
)
train_pair_set_file = os.path.join(
pair_set_dir, "train_{}.txt".format(class_name)
)
train_pair = sorted(train_pair)
with open(train_pair_set_file, "w") as text_file:
for x in train_pair:
text_file.write("{}\n".format(x))
test_pair_set_file = os.path.join(
pair_set_dir, "my_val_{}.txt".format(class_name)
)
val_pair = sorted(val_pair)
with open(test_pair_set_file, "w") as text_file:
for x in val_pair:
text_file.write("{}\n".format(x))
print(class_name, " done")
def main():
gen_images = (
True
) ##################################################################### gen_images
for class_idx, class_name in enumerate(class_list):
# train_pair = []
val_pair = []
if class_name in ["__back_ground__"]:
continue
if not class_name in sel_classes:
continue
print("start ", class_name)
model_path = os.path.join(
TLESS_root, "models/{}/obj_{}_scaled.ply".format(class_name, class_name)
)
model = inout.load_ply(model_path)
for video_name in sel_videos: # 02
for set_type in ["test"]:
with open(
os.path.join(
real_set_root,
"{}_{}_{}.txt".format(class_name, video_name, "test"),
),
"r",
) as f:
test_real_list = [x.strip() for x in f.readlines()]
with open(
rendered_pose_path.format(set_type, class_name, video_name)
) as f:
str_rendered_pose_list = [
x.strip().split(" ") for x in f.readlines()
]
rendered_pose_list = np.array(
[
[float(x) for x in each_pose]
for each_pose in str_rendered_pose_list
]
)
rendered_per_real = 1
assert len(rendered_pose_list) == 1 * len(
test_real_list
), "{} vs {}".format(len(rendered_pose_list), len(test_real_list))
for idx, real_index in enumerate(tqdm(test_real_list)):
_, cls_name, vid_name, real_prefix = real_index.split("/")
rendered_dir = os.path.join(
rendered_root_dir, class_name, video_name
)
mkdir_if_missing(rendered_dir)
for inner_idx in range(rendered_per_real):
if gen_images:
# if gen_images and real_index in test_real_list and inner_idx == 0: # only generate my_val_v{}
image_file = os.path.join(
rendered_dir, "{}-color.png".format(real_prefix)
)
depth_file = os.path.join(
rendered_dir, "{}-depth.png".format(real_prefix)
)
rendered_idx = idx * rendered_per_real + inner_idx
pose_rendered_q = rendered_pose_list[rendered_idx]
R = se3.quat2mat(pose_rendered_q[:4])
t = pose_rendered_q[4:]
surf_color = None # (1, 0, 0) # ?????
im_size = (width, height) # (w, h)
ren_rgb, ren_depth = renderer.render(
model,
im_size,
K_0,
R,
t,
clip_near=ZNEAR,
clip_far=ZFAR,
surf_color=surf_color,
mode="rgb+depth",
)
ren_depth = (ren_depth * depth_factor).astype(np.uint16)
cv2.imwrite(image_file, ren_rgb)
cv2.imwrite(depth_file, ren_depth)
pose_rendered_file = os.path.join(
rendered_dir, "{}-pose.txt".format(real_prefix)
)
text_file = open(pose_rendered_file, "w")
text_file.write("{}\n".format(class_idx))
pose_rendered_m = np.zeros((3, 4))
pose_rendered_m[:, :3] = se3.quat2mat(pose_rendered_q[:4])
pose_rendered_m[:, 3] = pose_rendered_q[4:]
pose_ori_m = pose_rendered_m
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
pose_ori_m[0, 0],
pose_ori_m[0, 1],
pose_ori_m[0, 2],
pose_ori_m[0, 3],
pose_ori_m[1, 0],
pose_ori_m[1, 1],
pose_ori_m[1, 2],
pose_ori_m[1, 3],
pose_ori_m[2, 0],
pose_ori_m[2, 1],
pose_ori_m[2, 2],
pose_ori_m[2, 3],
)
text_file.write(pose_str)
val_pair.append(
"{} test/{}/{}/{}".format(
real_index, class_name, video_name, real_prefix
)
)
test_pair_set_file = os.path.join(
pair_set_dir,
"my_val_{}_video{}_{}.txt".format(version, video_name, class_name),
)
val_pair = sorted(val_pair)
with open(test_pair_set_file, "w") as text_file:
for x in val_pair:
text_file.write("{}\n".format(x))
print(class_name, video_name, " done")
pose_rendered_q[:4], pose_rendered_q[4:])
rgb_gl = rgb_gl.astype('uint8')
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(image_file, rgb_gl)
cv2.imwrite(depth_file, depth_gl)
pose_file = os.path.join(
rendered_dir,
'{}_{}_{}-pose.txt'.format(class_name, real_prefix,
inner_idx))
text_file = open(pose_file, 'w')
text_file.write("{}\n".format(class_idx))
pose_rendered_m = np.zeros((3, 4))
pose_rendered_m[:, :3] = se3.quat2mat(
pose_rendered_q[:4])
pose_rendered_m[:, 3] = pose_rendered_q[4:]
pose_ori_m = pose_rendered_m
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}"\
.format(pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3],
pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3],
pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3])
text_file.write(pose_str)
train_pair.append("{} {}/{}_{}_{}".format(
real_index, video_name, class_name, real_prefix,
inner_idx))
pair_set_file = os.path.join(pair_set_dir,
"train_{}.txt".format(class_name))
train_pair = sorted(train_pair)
def gen_poses():
pz = np.array([0,0,1])
pose_dict, quat_stat, trans_stat, new_points = stat_poses()
real_prefix_list = ['{:06d}'.format(i + 1) for i in range(NUM_IMAGES)]
sel_classes = classes
num_class = len(sel_classes)
# syn_pose_dict = {prefix: np.zeros((num_class, 7), dtype='float32')
# for prefix in real_prefix_list} # store poses
syn_pose_dict = {cls_name: np.zeros((NUM_IMAGES, 7)) for cls_name in sel_classes}
syn_pose_dir = os.path.join(cur_path, '..', 'data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/syn_poses_single')
mkdir_if_missing(syn_pose_dir)
for cls_i, cls_name in enumerate(sel_classes):
if cls_name != 'ape':
continue
src_quat_mean = quat_stat[cls_name]['quat_mean']
src_quat_std = quat_stat[cls_name]['quat_std']
src_trans_mean = trans_stat[cls_name]['trans_mean']
src_trans_std = trans_stat[cls_name]['trans_std']
deg_max = new_points[cls_name]['angle_max']
for i in tqdm(range(NUM_IMAGES)):
real_prefix = real_prefix_list[i]
# if cls_name == 'ape':
# continue
# randomly generate a pose
tgt_quat = np.random.normal(0, 1, 4)
tgt_quat = tgt_quat/np.linalg.norm(tgt_quat)
if tgt_quat[0] < 0:
tgt_quat *= -1
tgt_trans = np.random.normal(src_trans_mean, src_trans_std)
tgt_rot_m = se3.quat2mat(tgt_quat)
new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3,))
pz_mean = new_points[cls_name]['pz_mean']
deg = angle(new_pz, pz_mean)
# r_dist, t_dist = calc_rt_dist_q(tgt_quat, src_quat, tgt_trans, src_trans)
transform = np.matmul(K, tgt_trans.reshape(3, 1))
center_x = float(transform[0] / transform[2])
center_y = float(transform[1] / transform[2])
count = 0
while (deg > deg_max or not (48 < center_x < (640 - 48) and 48 < center_y < (480 - 48))):
# randomly generate a pose
tgt_quat = np.random.normal(0, 1, 4)
tgt_quat = tgt_quat / np.linalg.norm(tgt_quat)
if tgt_quat[0] < 0:
tgt_quat *= -1
tgt_trans = np.random.normal(src_trans_mean, src_trans_std)
tgt_rot_m = se3.quat2mat(tgt_quat)
new_pz = np.dot(tgt_rot_m, pz.reshape((-1, 1))).reshape((3,))
pz_mean = new_points[cls_name]['pz_mean']
deg = angle(new_pz, pz_mean)
transform = np.matmul(K, tgt_trans.reshape(3, 1))
center_x = float(transform[0] / transform[2])
center_y = float(transform[1] / transform[2])
count += 1
if count % 100 == 0:
print(real_prefix, cls_name, count,
"deg < deg_max={}: {}, 48 < center_x < (640-48): {}, 48 < center_y < (480-48): {}".format(
deg_max, deg <= deg_max, 48 < center_x < (640 - 48), 48 < center_y < (480 - 48)))
tgt_pose_q = np.zeros((7,), dtype='float32')
tgt_pose_q[:4] = tgt_quat
tgt_pose_q[4:] = tgt_trans
syn_pose_dict[cls_name][i, :] = tgt_pose_q
# write pose
poses_file = os.path.join(syn_pose_dir, 'LM6d_ds_v1_all_syn_pose.pkl')
with open(poses_file, 'wb') as f:
cPickle.dump(syn_pose_dict, f, 2)
def main():
gen_images = True ############
for class_idx, class_name in idx2class.items():
val_pair = []
print("start ", class_name)
if class_name in ['__back_ground__']:
continue
if not class_name in ['bowl', 'cup']:
continue
if gen_images:
# init render
model_dir = os.path.join(LM6d_root, 'models', class_name)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR,
ZFAR)
for set_type in ['all']:
with open(
os.path.join(real_set_root,
'{}_{}.txt'.format(class_name, 'all')),
'r') as f:
all_real_list = [x.strip() for x in f.readlines()]
with open(
os.path.join(real_set_root,
'{}_{}.txt'.format(class_name, 'test')),
'r') as f:
test_real_list = [x.strip() for x in f.readlines()]
with open(rendered_pose_path.format(set_type, class_name)) as f:
str_rendered_pose_list = [
x.strip().split(' ') for x in f.readlines()
]
rendered_pose_list = np.array(
[[float(x) for x in each_pose]
for each_pose in str_rendered_pose_list])
rendered_per_real = 1
assert (len(rendered_pose_list) == 1*len(all_real_list)), \
'{} vs {}'.format(len(rendered_pose_list), len(all_real_list))
for idx, real_index in enumerate(tqdm(all_real_list)):
video_name, real_prefix = real_index.split('/')
rendered_dir = os.path.join(rendered_root_dir, class_name)
mkdir_if_missing(rendered_dir)
for inner_idx in range(rendered_per_real):
if gen_images:
image_file = os.path.join(
rendered_dir,
'{}_{}-color.png'.format(real_prefix, inner_idx))
depth_file = os.path.join(
rendered_dir,
'{}_{}-depth.png'.format(real_prefix, inner_idx))
rendered_idx = idx * rendered_per_real + inner_idx
pose_rendered_q = rendered_pose_list[rendered_idx]
rgb_gl, depth_gl = render_machine.render(
pose_rendered_q[:4], pose_rendered_q[4:])
rgb_gl = rgb_gl.astype('uint8')
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(image_file, rgb_gl)
cv2.imwrite(depth_file, depth_gl)
pose_rendered_file = os.path.join(
rendered_dir,
'{}_{}-pose.txt'.format(real_prefix, inner_idx))
text_file = open(pose_rendered_file, 'w')
text_file.write("{}\n".format(class_idx))
pose_rendered_m = np.zeros((3, 4))
pose_rendered_m[:, :3] = RT_transform.quat2mat(
pose_rendered_q[:4])
pose_rendered_m[:, 3] = pose_rendered_q[4:]
pose_ori_m = pose_rendered_m
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}"\
.format(pose_ori_m[0, 0], pose_ori_m[0, 1], pose_ori_m[0, 2], pose_ori_m[0, 3],
pose_ori_m[1, 0], pose_ori_m[1, 1], pose_ori_m[1, 2], pose_ori_m[1, 3],
pose_ori_m[2, 0], pose_ori_m[2, 1], pose_ori_m[2, 2], pose_ori_m[2, 3])
text_file.write(pose_str)
if real_index in test_real_list:
if inner_idx == 0:
val_pair.append("{} {}/{}_{}".format(
real_index, class_name, real_prefix,
inner_idx))
test_pair_set_file = os.path.join(
pair_set_dir, "my_val_{}.txt".format(class_name))
val_pair.pose_est
with open(test_pair_set_file, "w") as text_file:
for x in val_pair:
text_file.write("{}\n".format(x))
print(class_name, " done")
def main():
gen_images = True
for class_idx, class_name in enumerate(tqdm(classes)):
train_pair = []
print("start ", class_name)
if class_name in ["__back_ground__"]:
continue
if gen_images:
# init render machine
model_dir = os.path.join(
cur_path,
"../data/LINEMOD_6D/LM6d_converted/models/{}".format(
class_name),
)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR,
ZFAR)
for set_type in ["NDtrain"]:
# observed index list
observed_list_path = os.path.join(
observed_set_dir, "NDtrain_observed_{}.txt".format(class_name))
with open(observed_list_path, "r") as f:
observed_list = [x.strip() for x in f.readlines()]
# rendered poses
rendered_pose_path = os.path.join(
rendered_pose_dir,
"LM6d_occ_dsm_{}_NDtrain_rendered_pose_{}.txt".format(
version, class_name),
)
with open(rendered_pose_path, "r") as f:
str_rendered_pose_list = [
x.strip().split(" ") for x in f.readlines()
]
rendered_pose_list = np.array(
[[float(x) for x in each_pose]
for each_pose in str_rendered_pose_list])
rendered_per_observed = 1
assert len(rendered_pose_list
) == 1 * len(observed_list), "{} vs {}".format(
len(rendered_pose_list), len(observed_list))
for idx, observed_index in enumerate(tqdm(observed_list)):
video_name, observed_prefix = observed_index.split(
"/") # ./prefix
rendered_dir = os.path.join(rendered_root_dir, video_name)
mkdir_if_missing(rendered_dir)
rendered_dir = os.path.join(rendered_dir, class_name)
mkdir_if_missing(rendered_dir)
for inner_idx in range(rendered_per_observed):
if gen_images:
image_file = os.path.join(
rendered_dir,
"{}_{}-color.png".format(observed_prefix,
inner_idx),
)
depth_file = os.path.join(
rendered_dir,
"{}_{}-depth.png".format(observed_prefix,
inner_idx),
)
# if os.path.exists(image_file) and os.path.exists(depth_file):
# continue
rendered_idx = idx * rendered_per_observed + inner_idx
pose_rendered_q = rendered_pose_list[rendered_idx]
rgb_gl, depth_gl = render_machine.render(
pose_rendered_q[:4], pose_rendered_q[4:])
rgb_gl = rgb_gl.astype("uint8")
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(image_file, rgb_gl)
cv2.imwrite(depth_file, depth_gl)
pose_rendered_file = os.path.join(
rendered_dir,
"{}_{}-pose.txt".format(observed_prefix,
inner_idx),
)
text_file = open(pose_rendered_file, "w")
text_file.write("{}\n".format(class2idx(class_name)))
pose_rendered_m = np.zeros((3, 4))
pose_rendered_m[:, :3] = se3.quat2mat(
pose_rendered_q[:4])
pose_rendered_m[:, 3] = pose_rendered_q[4:]
pose_ori_m = pose_rendered_m
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
pose_ori_m[0, 0],
pose_ori_m[0, 1],
pose_ori_m[0, 2],
pose_ori_m[0, 3],
pose_ori_m[1, 0],
pose_ori_m[1, 1],
pose_ori_m[1, 2],
pose_ori_m[1, 3],
pose_ori_m[2, 0],
pose_ori_m[2, 1],
pose_ori_m[2, 2],
pose_ori_m[2, 3],
)
text_file.write(pose_str)
train_pair.append("{} {}/{}_{}".format(
observed_index, class_name, observed_prefix,
inner_idx))
pair_set_file = os.path.join(image_set_dir,
"train_{}.txt".format(class_name))
train_pair = sorted(train_pair)
with open(pair_set_file, "w") as text_file:
for x in train_pair:
text_file.write("{}\n".format(x))
print(class_name, " done")
def main():
gen_images = (
True
) ##################################################################### gen_images
for class_idx, class_name in enumerate(class_list):
train_pair = []
# val_pair = []
if class_name in ["__back_ground__"]:
continue
if not class_name in sel_classes:
continue
print("start ", class_name)
if gen_images:
# init render
model_dir = os.path.join(TLESS_root, "models", class_name)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR,
ZFAR)
for set_type in ["train"]:
with open(
os.path.join(real_set_root,
"{}_{}.txt".format(class_name, "train")),
"r",
) as f:
train_real_list = [x.strip() for x in f.readlines()]
with open(rendered_pose_path.format(set_type, class_name)) as f:
str_rendered_pose_list = [
x.strip().split(" ") for x in f.readlines()
]
rendered_pose_list = np.array(
[[float(x) for x in each_pose]
for each_pose in str_rendered_pose_list])
rendered_per_real = 10
assert len(rendered_pose_list) == rendered_per_real * len(
train_real_list), "{} vs {}".format(len(rendered_pose_list),
len(train_real_list))
for idx, real_index in enumerate(tqdm(train_real_list)):
video_name, real_prefix = real_index.split("/")
rendered_dir = os.path.join(rendered_root_dir, class_name)
mkdir_if_missing(rendered_dir)
for inner_idx in range(rendered_per_real):
if gen_images:
# if gen_images and real_index in test_real_list and inner_idx == 0: # only generate my_val_v{}
image_file = os.path.join(
rendered_dir,
"{}_{}-color.png".format(real_prefix, inner_idx),
)
depth_file = os.path.join(
rendered_dir,
"{}_{}-depth.png".format(real_prefix, inner_idx),
)
rendered_idx = idx * rendered_per_real + inner_idx
pose_rendered_q = rendered_pose_list[rendered_idx]
rgb_gl, depth_gl = render_machine.render(
pose_rendered_q[:4], pose_rendered_q[4:])
rgb_gl = rgb_gl.astype("uint8")
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(image_file, rgb_gl)
cv2.imwrite(depth_file, depth_gl)
pose_rendered_file = os.path.join(
rendered_dir,
"{}_{}-pose.txt".format(real_prefix, inner_idx),
)
text_file = open(pose_rendered_file, "w")
text_file.write("{}\n".format(class_idx))
pose_rendered_m = np.zeros((3, 4))
pose_rendered_m[:, :3] = se3.quat2mat(
pose_rendered_q[:4])
pose_rendered_m[:, 3] = pose_rendered_q[4:]
pose_ori_m = pose_rendered_m
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
pose_ori_m[0, 0],
pose_ori_m[0, 1],
pose_ori_m[0, 2],
pose_ori_m[0, 3],
pose_ori_m[1, 0],
pose_ori_m[1, 1],
pose_ori_m[1, 2],
pose_ori_m[1, 3],
pose_ori_m[2, 0],
pose_ori_m[2, 1],
pose_ori_m[2, 2],
pose_ori_m[2, 3],
)
text_file.write(pose_str)
train_pair.append("{} {}/{}_{}".format(
real_index, class_name, real_prefix, inner_idx))
train_pair_set_file = os.path.join(
pair_set_dir, "train_{}_{}.txt".format(version, class_name))
train_pair = sorted(train_pair)
with open(train_pair_set_file, "w") as text_file:
for x in train_pair:
text_file.write("{}\n".format(x))
print(class_name, " done")
def pose_q2m(pose_q):
pose = np.zeros((3, 4), dtype="float32")
pose[:3, :3] = RT_transform.quat2mat(pose_q[:4])
pose[:3, 3] = pose_q[4:]
return pose