def gen_gt_observed():
for cls_idx, cls_name in idx2class.items():
print(cls_idx, cls_name)
# uncomment here to only generate data for ape
# if cls_name != 'ape':
# continue
with open(
os.path.join(observed_set_dir, "{}_all.txt".format(cls_name)),
"r") as f:
all_indices = [line.strip("\r\n") for line in f.readlines()]
# render machine
model_dir = os.path.join(LM6d_root, "models", cls_name)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
for observed_idx in tqdm(all_indices):
video_name, prefix = observed_idx.split("/")
# read pose -------------------------------------
observed_meta_path = os.path.join(
observed_data_root, "{}-meta.mat".format(observed_idx))
meta_data = sio.loadmat(observed_meta_path)
inner_id = np.where(
np.squeeze(meta_data["cls_indexes"]) == cls_idx)
if len(meta_data["poses"].shape) == 2:
pose = meta_data["poses"]
else:
pose = np.squeeze(meta_data["poses"][:, :, inner_id])
new_pose_path = os.path.join(gt_observed_root, cls_name,
"{}-pose.txt".format(prefix))
mkdir_if_missing(os.path.join(gt_observed_root, cls_name))
# write pose
write_pose_file(new_pose_path, cls_idx, pose)
# ----------------------render color, depth ------------
rgb_gl, depth_gl = render_machine.render(
RT_transform.mat2quat(pose[:3, :3]), pose[:, -1])
if any([x in observed_idx
for x in ["000128", "000256", "000512"]]):
rgb_gl = rgb_gl.astype("uint8")
render_color_path = os.path.join(gt_observed_root, cls_name,
"{}-color.png".format(prefix))
cv2.imwrite(render_color_path, rgb_gl)
# depth
depth_save = depth_gl * DEPTH_FACTOR
depth_save = depth_save.astype("uint16")
render_depth_path = os.path.join(gt_observed_root, cls_name,
"{}-depth.png".format(prefix))
cv2.imwrite(render_depth_path, depth_save)
# --------------------- render label ----------------------------------
render_label = depth_gl != 0
render_label = render_label.astype("uint8")
# write label
label_path = os.path.join(gt_observed_root, cls_name,
"{}-label.png".format(prefix))
cv2.imwrite(label_path, render_label)
def gen_render_real():
for cls_idx, cls_name in idx2class.items():
print(cls_idx, cls_name)
if cls_name == 'driller':
continue
with open(
os.path.join(real_set_dir,
'occLM_val_real_{}.txt'.format(cls_name)),
'r') as f:
all_indices = [line.strip('\r\n') for line in f.readlines()]
# render machine
model_dir = os.path.join(LM6d_root, 'models', cls_name)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
for real_idx in tqdm(all_indices):
video_name, prefix = real_idx.split('/') # video name is "test"
# read pose -------------------------------------
real_meta_path = os.path.join(real_data_root,
"02/{}-meta.mat".format(prefix))
meta_data = sio.loadmat(real_meta_path)
inner_id = np.where(
np.squeeze(meta_data['cls_indexes']) == cls_idx)
if len(meta_data['poses'].shape) == 2:
pose = meta_data['poses']
else:
pose = np.squeeze(meta_data['poses'][:, :, inner_id])
new_pose_path = os.path.join(render_real_root, cls_name,
"{}-pose.txt".format(prefix))
mkdir_if_missing(os.path.join(render_real_root, cls_name))
# write pose
write_pose_file(new_pose_path, cls_idx, pose)
# ----------------------render color, depth ------------
rgb_gl, depth_gl = render_machine.render(
RT_transform.mat2quat(pose[:3, :3]), pose[:, -1])
rgb_gl = rgb_gl.astype('uint8')
render_color_path = os.path.join(render_real_root, cls_name,
"{}-color.png".format(prefix))
cv2.imwrite(render_color_path, rgb_gl)
# depth
depth_save = depth_gl * DEPTH_FACTOR
depth_save = depth_save.astype('uint16')
render_depth_path = os.path.join(render_real_root, cls_name,
"{}-depth.png".format(prefix))
cv2.imwrite(render_depth_path, depth_save)
#--------------------- render label ----------------------------------
render_label = depth_gl != 0
render_label = render_label.astype('uint8')
# write label
label_path = os.path.join(render_real_root, cls_name,
"{}-label.png".format(prefix))
cv2.imwrite(label_path, render_label)
def render_pose(rendered_dir, count, class_name, fixed_transforms_dict,
camera_intrinsics, camera_pose, rotation_angles, location,
rotation):
width = 960
height = 540
K = np.array([[camera_intrinsics['fx'], 0, camera_intrinsics['cx']],
[0, camera_intrinsics['fy'], camera_intrinsics['cy']],
[0, 0, 1]])
# Check these TODO
ZNEAR = 0.1
ZFAR = 20
depth_factor = 1000
model_dir = os.path.join(LM6d_root, "models", class_name)
# model_dir = os.path.join(LM6d_root, "aligned_cm", class_name, "google_16k")
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
# camera_pose_matrix = np.zeros((4,4))
# camera_pose_matrix[:, 3] = [i/100 for i in camera_pose['location_worldframe']] + [1]
# camera_pose_matrix[:3, :3] = RT_transform.quat2mat(get_wxyz_quaternion(camera_pose['quaternion_xyzw_worldframe']))
# camera_pose_matrix[3,:3] = [0,0,0]
# print(camera_pose_matrix)
fixed_transform = np.transpose(np.array(fixed_transforms_dict[class_name]))
fixed_transform[:3, 3] = [i / 100 for i in fixed_transform[:3, 3]]
object_world_transform = np.zeros((4, 4))
object_world_transform[:3, :3] = RT_transform.euler2mat(
rotation_angles[0], rotation_angles[1], rotation_angles[2])
# object_world_transform[:3,:3] = RT_transform.euler2mat(0,0,0)
# object_world_transform[:3, :3] = RT_transform.quat2mat(get_wxyz_quaternion(rotation))
object_world_transform[:, 3] = [i / 100 for i in location] + [1]
# print(fixed_transform)
# total_transform = np.matmul(np.linalg.inv(camera_pose_matrix), object_world_transform)
# fixed_transform = np.matmul(m, fixed_transform)
total_transform = np.matmul(object_world_transform, fixed_transform)
# total_transform = object_world_transform
pose_rendered_q = RT_transform.mat2quat(total_transform[:3, :3]).tolist(
) + total_transform[:3, 3].flatten().tolist()
# pose_rendered_q = RT_transform.mat2quat(object_world_transform[:3,:3]).tolist() + object_world_transform[:3,3].flatten().tolist()
# print(pose_rendered_q)
# rendered_dir = '.'
# image_file = os.path.join(
# rendered_dir,
# "{}-{}-color.png".format(count, class_name),
# )
# depth_file = os.path.join(
# rendered_dir,
# "{}-{}-depth.png".format(count, class_name),
# )
rgb_gl, depth_gl = render_machine.render(pose_rendered_q[:4],
np.array(pose_rendered_q[4:]))
rgb_gl = rgb_gl.astype("uint8")
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
return rgb_gl, depth_gl
rendered_root_dir = os.path.join(
cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/data',
version)
pair_set_dir = os.path.join(
cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/image_set')
mkdir_if_missing(rendered_root_dir)
mkdir_if_missing(pair_set_dir)
all_pair = []
for small_class_idx, class_name in enumerate(tqdm(class_name_list)):
big_class_idx = class2big_idx[class_name]
# init render
model_dir = os.path.join(
cur_path,
'../data/LINEMOD_6D/LM6d_converted/models/{}'.format(class_name))
if gen_images:
render_machine = Render_Py(model_dir, K, width, height, ZNEAR,
ZFAR)
with open(keyframe_path.format(class_name)) as f:
real_index_list = [x.strip() for x in f.readlines()]
video_name_list = [x.split('/')[0] for x in real_index_list]
real_prefix_list = [x.split('/')[1] for x in real_index_list]
all_pair = []
for idx, real_index in enumerate(real_index_list):
rendered_dir = os.path.join(rendered_root_dir,
video_name_list[idx], class_name)
mkdir_if_missing(rendered_dir)
yu_idx = idx
yu_pred_file = os.path.join(yu_pred_dir, class_name,
"{:04d}.mat".format(yu_idx))
yu_pred = sio.loadmat(yu_pred_file)
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 adapt_real_train():
class_list = ["{:02d}".format(i) for i in range(1, 31)]
sel_classes = ["05", "06"]
width = 640 # 400
height = 480 # 400
depth_factor = 10000
K_0 = np.array(
[[1075.65091572, 0, 320.0], [0, 1073.90347929, 240.0], [0, 0, 1]]
) # Primesense
new_data_root = os.path.join(TLESS_root, "TLESS_render_v3/data/real")
mkdir_if_missing(new_data_root)
real_set_dir = os.path.join(TLESS_root, "TLESS_render_v3/image_set/real")
mkdir_if_missing(real_set_dir)
for cls_idx, cls_name in enumerate(class_list):
if not cls_name in sel_classes:
continue
print(cls_idx, cls_name)
model_dir = os.path.join(TLESS_root, "models", cls_name)
render_machine = Render_Py(model_dir, K_0, width, height, ZNEAR, ZFAR)
gt_path = os.path.join(
TLESS_root, "t-less_v2/train_primesense/{}/gt.yml".format(cls_name)
)
gt_dict = load_gt(gt_path)
info_path = os.path.join(
TLESS_root, "t-less_v2/train_primesense/{}/info.yml".format(cls_name)
)
info_dict = load_info(info_path)
real_indices = []
for img_id in tqdm(gt_dict.keys()):
R = np.array(gt_dict[img_id][0]["cam_R_m2c"]).reshape((3, 3))
t = np.array(gt_dict[img_id][0]["cam_t_m2c"]) / 1000.0
K = np.array(info_dict[img_id]["cam_K"]).reshape((3, 3))
# K[0, 2] += 120 # cx
# K[1, 2] += 40 # cy
pose = np.zeros((3, 4))
pose[:3, :3] = R
pose[:3, 3] = t
# print(pose)
# print(K)
K_diff = K_0 - K
cx_diff = K_diff[0, 2]
cy_diff = K_diff[1, 2]
px_diff = int(np.round(cx_diff))
py_diff = int(np.round(cy_diff))
# pose ----------------
pose_path = os.path.join(
new_data_root, cls_name, "{:06d}-pose.txt".format(img_id)
)
mkdir_if_missing(os.path.join(new_data_root, cls_name))
write_pose_file(pose_path, cls_idx, pose)
rgb_gl, depth_gl = render_machine.render(
pose[:3, :3], pose[:, -1], r_type="mat", K=K_0
)
rgb_gl = rgb_gl.astype("uint8")
# depth ------------------
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
depth_path = os.path.join(
new_data_root, cls_name, "{:06d}-depth.png".format(img_id)
)
cv2.imwrite(depth_path, depth_gl)
# label ---------------------
label_gl = np.zeros(depth_gl.shape)
label_gl[depth_gl != 0] = 1
label_path = os.path.join(
new_data_root, cls_name, "{:06d}-label.png".format(img_id)
)
cv2.imwrite(label_path, label_gl)
# real color ----------------------------
color_real = read_img(
os.path.join(
ori_train_data_root, cls_name, "rgb/{:04d}.png".format(img_id)
),
3,
)
# print(color_real.max(), color_real.min())
pad_real = np.zeros((480, 640, 3))
xs = 0
ys = 0
pad_real[xs : 400 + xs, ys : 400 + ys, :] = color_real
pad_real = pad_real.astype("uint8")
# translate image
M = np.float32([[1, 0, px_diff], [0, 1, py_diff]])
pad_real = cv2.warpAffine(pad_real, M, (640, 480))
color_path = os.path.join(
new_data_root, cls_name, "{:06d}-color.png".format(img_id)
)
cv2.imwrite(color_path, pad_real)
# real index
real_indices.append("{}/{:06d}".format(cls_name, img_id))
real_indices = sorted(real_indices)
real_set_file = os.path.join(real_set_dir, "{}_train.txt".format(cls_name))
with open(real_set_file, "w") as f:
for real_idx in real_indices:
f.write(real_idx + "\n")
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 gen_gt_observed():
with open(syn_poses_path, "rb") as f:
syn_pose_dict = cPickle.load(f)
for class_idx, class_name in enumerate(classes):
if class_name == "__back_ground__":
continue
# uncomment here to only generate data for ape
# if class_name not in ['ape']:
# continue
# init render machines
# brightness_ratios = [0.2, 0.25, 0.3, 0.35, 0.4] ###################
model_dir = os.path.join(LINEMOD_syn_root,
"models/{}".format(class_name))
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
# syn_poses_path = os.path.join(syn_poses_dir, 'LM6d_v1_all_rendered_pose_{}.txt'.format(class_name))
# syn_poses = np.loadtxt(syn_poses_path)
# print(syn_poses.shape) # nx7
syn_poses = syn_pose_dict[class_name]
num_poses = syn_poses.shape[0]
observed_index_list = [
"{}/{:06d}".format(class_name, i + 1) for i in range(num_poses)
]
# observed_set_path = os.path.join(
# image_set_dir,
# 'observed/LM_data_syn_train_observed_{}.txt'.format(class_name))
# mkdir_if_missing(os.path.join(image_set_dir, 'observed'))
# f_observed_set = open(observed_set_path, 'w')
for idx, observed_index in enumerate(tqdm(observed_index_list)):
# f_observed_set.write('{}\n'.format(observed_index))
# continue # just generate observed set file
prefix = observed_index.split("/")[1]
gt_observed_dir = os.path.join(gt_observed_root_dir, class_name)
mkdir_if_missing(gt_observed_dir)
gt_observed_color_file = os.path.join(gt_observed_dir,
prefix + "-color.png")
gt_observed_depth_file = os.path.join(gt_observed_dir,
prefix + "-depth.png")
gt_observed_pose_file = os.path.join(gt_observed_dir,
prefix + "-pose.txt")
# observed_label_file = os.path.join(observed_root_dir, video_name, prefix + "-label.png")
gt_observed_label_file = os.path.join(gt_observed_dir,
prefix + "-label.png")
pose_quat = syn_poses[idx, :]
pose = se3.se3_q2m(pose_quat)
# generate random light_position
if idx % 6 == 0:
light_position = [1, 0, 1]
elif idx % 6 == 1:
light_position = [1, 1, 1]
elif idx % 6 == 2:
light_position = [0, 1, 1]
elif idx % 6 == 3:
light_position = [-1, 1, 1]
elif idx % 6 == 4:
light_position = [-1, 0, 1]
elif idx % 6 == 5:
light_position = [0, 0, 1]
else:
raise Exception("???")
# print( "light_position a: {}".format(light_position))
light_position = np.array(light_position) * 0.5
# inverse yz
light_position[0] += pose[0, 3]
light_position[1] -= pose[1, 3]
light_position[2] -= pose[2, 3]
# print("light_position b: {}".format(light_position))
# get render result
rgb_gl, depth_gl = render_machine.render(pose[:3, :3],
pose[:, 3],
r_type="mat")
rgb_gl = rgb_gl.astype("uint8")
# gt_observed label
label_gl = np.zeros(depth_gl.shape)
# print('depth gl:', depth_gl.shape)
label_gl[depth_gl != 0] = 1
cv2.imwrite(gt_observed_color_file, rgb_gl)
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(gt_observed_depth_file, depth_gl)
cv2.imwrite(gt_observed_label_file, label_gl)
text_file = open(gt_observed_pose_file, "w")
text_file.write("{}\n".format(class_idx))
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
pose[0, 0],
pose[0, 1],
pose[0, 2],
pose[0, 3],
pose[1, 0],
pose[1, 1],
pose[1, 2],
pose[1, 3],
pose[2, 0],
pose[2, 1],
pose[2, 2],
pose[2, 3],
)
text_file.write(pose_str)
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 main(camera_params, env_params):
width = camera_params['camera_width']
height = camera_params['camera_height']
K = np.array([[camera_params['camera_fx'], 0, camera_params['camera_cx']], [0, camera_params['camera_fy'], camera_params['camera_cy']], [0, 0, 1]])
ZNEAR = camera_params['camera_znear']
ZFAR = camera_params['camera_zfar']
depth_factor = 1000
x_min = float(env_params['x_min'])
x_max = float(env_params['x_max']);
y_min = float(env_params['y_min']);
y_max = float(env_params['y_max']);
table_height = float(env_params['table_height']);
gen_images = True
pose_from_file = False
print("Camera Matrix:")
print(K)
# camera_pose = np.array([ \
# [0.868216, 6.3268e-06, 0.496186, 0.436202], \
# [-5.49302e-06, 1, -3.13929e-06, 0.0174911], \
# [-0.496186, 2.74908e-11, 0.868216, 0.709983], \
# [0, 0, 0, 1]])
# Camera to world transform
# camera_pose = np.array([ \
# [0.0068906 , -0.497786, 0.867272 , 0.435696], \
# [-0.999953, 0.0024452, 0.00934823, 0.0323318], \
# [-0.00677407, -0.867296, -0.497746, 0.710332], \
# [0, 0, 0, 1]])
camera_pose = np.array([ \
[0.00572327, -0.629604, 0.776895, 0.437408], \
[-0.999953, 0.00244603, 0.0093488, 0.0323317], \
[-0.00778635, -0.776912, -0.629561, 0.709281], \
[0, 0, 0, 1]])
#
# camera_pose = np.array([ \
# [0.778076, 6.3268e-06, 0.628171, 0.43785], \
# [-4.92271e-06, 1, -3.97433e-06, 0.0174995], \
# [ -0.628171, 2.70497e-11, 0.778076, 0.708856], \
# [ 0, 0, 0, 1]])
# cam_to_body = np.array([[ 0, 0, 1, 0],
# [-1, 0, 0, 0],
# [0, -1, 0, 0],
# [0, 0, 0, 1]]);
for class_idx, class_name in idx2class.items():
print("start ", class_idx, class_name)
if class_name in ["__back_ground__"]:
continue
if gen_images:
# init render
# model_dir = os.path.join(LM6d_root, "aligned_cm", class_name, "google_16k")
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"]:
rendered_pose_list = []
# For reading in Perch
rendered_pose_list_out = []
if pose_from_file:
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]
)
else:
for x in np.arange(x_min, x_max, float(env_params['search_resolution_translation'])):
for y in np.arange(y_min, y_max, float(env_params['search_resolution_translation'])):
for theta in np.arange(0, 2 * np.pi, float(env_params['search_resolution_yaw'])):
original_point = np.array([[x], [y], [table_height], [1]])
if class_name == "004_sugar_box":
# Add half the height of box to shift it up
point = np.array([[x], [y], [table_height+0.086], [1]])
if class_name == "035_power_drill":
point = np.array([[x], [y], [table_height], [1]])
# transformed_point = np.matmul(np.linalg.inv(camera_pose), point)
# transformed_rotation = np.matmul(np.linalg.inv(camera_pose[0:3, 0:3]), RT_transform.euler2mat(0,0,theta))
# transformed_rotation = np.linalg.inv(camera_pose)[0:3, 0:3]
# transformed_rotation = RT_transform.euler2mat(0,0,0)
# print(transformed_point)
object_world_transform = np.zeros((4,4))
if class_name == "004_sugar_box":
object_world_transform[:3,:3] = RT_transform.euler2mat(0,0,theta)
if class_name == "035_power_drill":
object_world_transform[:3,:3] = RT_transform.euler2mat(np.pi/2,0,theta)
object_world_transform[:4,3] = point.flatten()
# print(world_object_transform)
# First apply world to object transform on the object and then take it to camera frame
total_transform = np.matmul(np.linalg.inv(camera_pose), object_world_transform)
print(total_transform)
pose = RT_transform.mat2quat(total_transform[:3,:3]).tolist() + total_transform[:3,3].flatten().tolist()
# pose = RT_transform.mat2quat(transformed_rotation).tolist() + transformed_point.flatten()[0:3].tolist()
print(pose)
rendered_pose_list.append(pose)
# rendered_pose_list_out.append(point.flatten().tolist() + [0,0,theta])
rendered_pose_list_out.append(original_point.flatten().tolist() + [0,0,theta])
rendered_pose_list = np.array(rendered_pose_list)
rendered_pose_list_out = np.array(rendered_pose_list_out)
for idx, observed_pose in enumerate(tqdm(rendered_pose_list)):
# print(idx)
# print(observed_pose)
rendered_dir = os.path.join(rendered_root_dir, class_name)
mkdir_if_missing(rendered_dir)
if gen_images:
image_file = os.path.join(
rendered_dir,
"{}-color.png".format(idx),
)
depth_file = os.path.join(
rendered_dir,
"{}-depth.png".format(idx),
)
pose_rendered_q = observed_pose
# print(pose_rendered_q[4:])
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(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)
pose_rendered_file = os.path.join(
rendered_dir,
"poses.txt",
)
np.savetxt(pose_rendered_file, np.around(rendered_pose_list_out, 4))
# text_file = open(pose_rendered_file, "w")
# text_file.write(rendered_pose_list)
print(class_name, " done")
def main():
for cls_idx, cls_name in enumerate(tqdm(sel_classes)):
print(cls_idx, cls_name)
keyframe_path = os.path.join(
observed_set_dir, "train_observed_{}.txt".format(cls_name)
)
with open(keyframe_path) as f:
observed_index_list = [x.strip() for x in f.readlines()]
video_name_list = [x.split("/")[0] for x in observed_index_list]
observed_prefix_list = [x.split("/")[1] for x in observed_index_list]
# init renderer
model_dir = os.path.join(model_root, cls_name)
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
for idx, observed_index in enumerate(tqdm(observed_index_list)):
prefix = observed_prefix_list[idx]
video_name = video_name_list[idx]
gt_observed_dir = os.path.join(gt_observed_root_dir, cls_name)
mkdir_if_missing(gt_observed_dir)
gt_observed_dir = os.path.join(gt_observed_dir, video_name) # ./
mkdir_if_missing(gt_observed_dir)
# to be written
gt_observed_color_file = os.path.join(
gt_observed_dir, prefix + "-color.png"
)
gt_observed_depth_file = os.path.join(
gt_observed_dir, prefix + "-depth.png"
)
gt_observed_pose_file = os.path.join(gt_observed_dir, prefix + "-pose.txt")
gt_observed_label_file = os.path.join(
gt_observed_dir, prefix + "-label.png"
)
observed_pose_file = os.path.join(
observed_root_dir, video_name, prefix + "-poses.npy"
)
observed_poses = np.load(observed_pose_file)
observed_pose_dict = observed_poses.all()
# pprint(observed_pose_dict)
if cls_name not in observed_pose_dict.keys():
continue
pose = observed_pose_dict[cls_name]
rgb_gl, depth_gl = render_machine.render(
RT_transform.mat2quat(pose[:3, :3]), pose[:, -1]
)
rgb_gl = rgb_gl.astype("uint8")
label_gl = np.zeros(depth_gl.shape)
label_gl[depth_gl != 0] = 1
depth_gl = depth_gl * depth_factor
depth_gl = depth_gl.astype("uint16")
# write results
cv2.imwrite(gt_observed_color_file, rgb_gl)
cv2.imwrite(gt_observed_depth_file, depth_gl)
cv2.imwrite(gt_observed_label_file, label_gl)
text_file = open(gt_observed_pose_file, "w")
text_file.write("{}\n".format(cls_idx))
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}".format(
pose[0, 0],
pose[0, 1],
pose[0, 2],
pose[0, 3],
pose[1, 0],
pose[1, 1],
pose[1, 2],
pose[1, 3],
pose[2, 0],
pose[2, 1],
pose[2, 2],
pose[2, 3],
)
text_file.write(pose_str)
print(cls_name, " done")
q = np.zeros(4)
q[0] = np.cos(0.5 * angle)
q[1:] = np.sin(0.5 * angle) * rot_axis
if q[0] < 0:
q *= -1
# print('norm of q: ', LA.norm(q))
q = q / LA.norm(q)
# print('norm of q: ', LA.norm(q))
return q
if __name__ == "__main__":
cls_name = 'bowl'
model_folder = './data/LINEMOD_6D/models/{}'.format(cls_name)
print('init render machine...')
render_machine = Render_Py(model_folder, K, width, height, zNear, zFar)
def rotate(angle, rot_axis, pose_gt, p_center=np.array([0,0,0])):
rot_sym_q = angle_axis_to_quat(angle, rot_axis)
rot_sym_m = quat2mat(rot_sym_q)
# print(rot_sym_m)
rot_res = R_transform(pose_gt[:3, :3], rot_sym_m, rot_coord='model')
rot_res_q = mat2quat(rot_res)
rgb_gl, depth_gl = render_machine.render(rot_res_q, pose_gt[:, 3] + p_center)
rgb_gl = rgb_gl.astype('uint8')
return rgb_gl, depth_gl
pose_gt = np.loadtxt('./data/LM6d_render_v1/data/render_real/bowl/000001-pose.txt',
skiprows=1)
def gen_render_real():
syn_poses_dir = os.path.join(cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/syn_poses_single/')
# output path
render_real_root_dir = os.path.join(LINEMOD_root, 'LM6d_data_syn_light', 'data', 'render_real')
image_set_dir = os.path.join(LINEMOD_root, 'LM6d_data_syn_light/image_set')
mkdir_if_missing(render_real_root_dir)
mkdir_if_missing(image_set_dir)
syn_poses_path = os.path.join(syn_poses_dir, 'LM6d_ds_v1_all_syn_pose.pkl')
with open(syn_poses_path, 'rb') as f:
syn_pose_dict = cPickle.load(f)
for class_idx, class_name in enumerate(tqdm(classes)):
if class_name == '__back_ground__':
continue
if class_name not in ['ape']:
continue
# init render machines
# brightness_ratios = [0.2, 0.25, 0.3, 0.35, 0.4] ###################
model_dir = os.path.join(LINEMOD_root, 'models/{}'.format(class_name))
render_machine = Render_Py(model_dir, K, width, height, ZNEAR, ZFAR)
# syn_poses_path = os.path.join(syn_poses_dir, 'LM6d_v1_all_rendered_pose_{}.txt'.format(class_name))
# syn_poses = np.loadtxt(syn_poses_path)
# print(syn_poses.shape) # nx7
syn_poses = syn_pose_dict[class_name]
num_poses = syn_poses.shape[0]
real_index_list = ['{}/{:06d}'.format(class_name, i+1) for i in range(num_poses)]
# real_set_path = os.path.join(image_set_dir, 'real/LM_data_syn_train_real_{}.txt'.format(class_name))
# mkdir_if_missing(os.path.join(image_set_dir, 'real'))
# f_real_set = open(real_set_path, 'w')
all_pair = []
for idx, real_index in enumerate(real_index_list):
# f_real_set.write('{}\n'.format(real_index))
# continue # just generate real set file
prefix = real_index.split('/')[1]
video_name = real_index.split('/')[0]
render_real_dir = os.path.join(render_real_root_dir, class_name)
mkdir_if_missing(render_real_dir)
render_real_color_file = os.path.join(render_real_dir, prefix+"-color.png")
render_real_depth_file = os.path.join(render_real_dir, prefix+"-depth.png")
render_real_pose_file = os.path.join(render_real_dir, prefix+"-pose.txt")
# real_label_file = os.path.join(real_root_dir, video_name, prefix + "-label.png")
render_real_label_file = os.path.join(render_real_dir, prefix + "-label.png")
if idx % 500 == 0:
print(' ', class_name, idx, '/', len(real_index_list), ' ', real_index)
pose_quat = syn_poses[idx, :]
pose = se3.se3_q2m(pose_quat)
# generate random light_position
if idx%6 == 0:
light_position = [1, 0, 1]
elif idx%6 == 1:
light_position = [1, 1, 1]
elif idx%6 == 2:
light_position = [0, 1, 1]
elif idx%6 == 3:
light_position = [-1, 1, 1]
elif idx%6 == 4:
light_position = [-1, 0, 1]
elif idx%6 == 5:
light_position = [0, 0, 1]
else:
raise Exception("???")
# print( "light_position a: {}".format(light_position))
light_position=np.array(light_position)*0.5
# inverse yz
light_position[0] += pose[0, 3]
light_position[1] -= pose[1, 3]
light_position[2] -= pose[2, 3]
# print("light_position b: {}".format(light_position))
# randomly adjust color and intensity for light_intensity
colors = np.array([[0, 0, 1],
[0, 1, 0],
[0, 1, 1],
[1, 0, 0],
[1, 0, 1],
[1, 1, 0],
[1, 1, 1]])
intensity = np.random.uniform(0.9, 1.1, size=(3,))
colors_randk = random.randint(0, colors.shape[0] - 1)
light_intensity = colors[colors_randk] * intensity
# print('light intensity: ', light_intensity)
# print('brightness ratio:', brightness_ratios[rm_randk])
# get render result
rgb_gl, depth_gl = render_machine.render(pose[:3, :3], pose[:, 3], r_type='mat')
rgb_gl = rgb_gl.astype('uint8')
# render_real label
label_gl = np.zeros(depth_gl.shape)
# print('depth gl:', depth_gl.shape)
label_gl[depth_gl!=0] = 1
# import matplotlib.pyplot as plt
# fig = plt.figure()
# plt.axis('off')
# fig.add_subplot(1, 3, 1)
# plt.imshow(rgb_gl[:, :, [2,1,0]])
#
# fig.add_subplot(1, 3, 2)
# plt.imshow(depth_gl)
#
# fig.add_subplot(1, 3, 3)
# plt.imshow(label_gl)
#
# fig.suptitle('light position: {}\n light_intensity: {}\n brightness: {}'.format(light_position, light_intensity, brightness_ratios[rm_randk]))
# plt.show()
cv2.imwrite(render_real_color_file, rgb_gl)
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(render_real_depth_file, depth_gl)
cv2.imwrite(render_real_label_file, label_gl)
text_file = open(render_real_pose_file, 'w')
text_file.write("{}\n".format(class_idx))
pose_str = "{} {} {} {}\n{} {} {} {}\n{} {} {} {}" \
.format(pose[0, 0], pose[0, 1], pose[0, 2], pose[0, 3],
pose[1, 0], pose[1, 1], pose[1, 2], pose[1, 3],
pose[2, 0], pose[2, 1], pose[2, 2], pose[2, 3])
text_file.write(pose_str)
print(class_name, " done")