def download_dataset(self, dataset_dir, dataset_url):
"""Downloads and extracts dataset.
Args:
dataset_dir (str): dataset directory.
dataset_url (str): url to download dataset.
"""
if osp.exists(dataset_dir):
return
if dataset_url is None:
raise RuntimeError(
'{} dataset needs to be manually '
'prepared, please follow the '
'document to prepare this dataset'.format(
self.__class__.__name__
)
)
print('Creating directory "{}"'.format(dataset_dir))
util.mkdir_if_missing(dataset_dir)
fpath = osp.join(dataset_dir, osp.basename(dataset_url))
print(
'Downloading {} dataset to "{}"'.format(
self.__class__.__name__, dataset_dir
)
)
print('Extracting "{}"'.format(fpath))
try:
tar = tarfile.open(fpath)
tar.extractall(path=dataset_dir)
tar.close()
except:
zip_ref = zipfile.ZipFile(fpath, 'r')
zip_ref.extractall(dataset_dir)
zip_ref.close()
print('{} dataset is ready'.format(self.__class__.__name__))
ZNEAR = 0.25
ZFAR = 6.0
depth_factor = 1000
LINEMOD_root = os.path.join(cur_path,
'../data/LINEMOD_6D/LM6d_converted/LM6d_refine')
gt_observed_dir = os.path.join(LINEMOD_root, 'data/gt_observed')
observed_set_dir = os.path.join(LINEMOD_root, 'image_set/observed')
# output path
LINEMOD_syn_root = os.path.join(
cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_refine_syn')
pose_dir = os.path.join(LINEMOD_syn_root,
'poses') # single object in each image
mkdir_if_missing(pose_dir)
print("target path: {}".format(pose_dir))
NUM_IMAGES = 10000
def angle(u, v):
c = np.dot(u, v) / (np.linalg.norm(u) * np.linalg.norm(v)
) # -> cosine of the angle
rad = np.arccos(np.clip(c, -1, 1))
deg = rad / np.pi * 180
return deg
def stat_poses():
pz = np.array([0, 0, 1])
ZNEAR = 0.25
ZFAR = 6.0
depth_factor = 1000
gen_images = True
# output_path
version = 'rendered_Yu_v02'
real_root_dir = os.path.join(
cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/data/real')
real_meta_path = "%s/{}-meta.mat" % (real_root_dir)
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]
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")
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():
train_pair = []
val_pair = []
print("start ", 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 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(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{}
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)
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 gen_observed():
# output path
observed_root_dir = os.path.join(LINEMOD_syn_root, 'data', 'observed')
image_set_dir = os.path.join(LINEMOD_syn_root, 'image_set')
mkdir_if_missing(observed_root_dir)
mkdir_if_missing(image_set_dir)
syn_poses_path = os.path.join(observed_pose_dir,
'LM6d_ds_train_observed_pose_all.pkl')
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', class_name)
render_machine = Render_Py_Light(model_dir, K, width, height, ZNEAR,
ZFAR, brightness_ratios)
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/LM6d_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))
prefix = observed_index.split('/')[1]
observed_dir = os.path.join(observed_root_dir, class_name)
mkdir_if_missing(observed_dir)
observed_color_file = os.path.join(observed_dir,
prefix + "-color.png")
observed_depth_file = os.path.join(observed_dir,
prefix + "-depth.png")
observed_pose_file = os.path.join(observed_dir,
prefix + "-pose.txt")
observed_label_file = os.path.join(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("???")
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]
# 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
# randomly choose a render machine
rm_randk = random.randint(0, len(brightness_ratios) - 1)
# get render result
rgb_gl, depth_gl = render_machine.render(se3.mat2quat(
pose[:3, :3]),
pose[:, -1],
light_position,
light_intensity,
brightness_k=rm_randk)
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(observed_color_file, rgb_gl)
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(observed_depth_file, depth_gl)
cv2.imwrite(observed_label_file, label_gl)
text_file = open(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 gen_real():
syn_poses_dir = os.path.join(
cur_path,
'../data/LINEMOD_6D/LM6d_converted/LM6d_render_v1/syn_poses_single/')
# output path
real_root_dir = os.path.join(LINEMOD_root, 'LM6d_data_syn_light', 'data',
'real')
image_set_dir = os.path.join(LINEMOD_root, 'LM6d_data_syn_light/image_set')
mkdir_if_missing(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', class_name)
render_machine = Render_Py_Light(model_dir, K, width, height, ZNEAR,
ZFAR, brightness_ratios)
# 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/LM6d_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]
real_dir = os.path.join(real_root_dir, class_name)
mkdir_if_missing(real_dir)
real_color_file = os.path.join(real_dir, prefix + "-color.png")
real_depth_file = os.path.join(real_dir, prefix + "-depth.png")
real_pose_file = os.path.join(real_dir, prefix + "-pose.txt")
# real_label_file = os.path.join(real_root_dir, video_name, prefix + "-label.png")
real_label_file = os.path.join(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)
# randomly choose a render machine
rm_randk = random.randint(0, len(brightness_ratios) - 1)
# print('brightness ratio:', brightness_ratios[rm_randk])
# get render result
rgb_gl, depth_gl = render_machine.render(se3.mat2quat(
pose[:3, :3]),
pose[:, -1],
light_position,
light_intensity,
brightness_k=rm_randk)
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(real_color_file, rgb_gl)
depth_gl = (depth_gl * depth_factor).astype(np.uint16)
cv2.imwrite(real_depth_file, depth_gl)
cv2.imwrite(real_label_file, label_gl)
text_file = open(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")
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')
all_pair = []
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]
video_name = observed_index.split('/')[0]
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")
# config for renderer
width = 640
height = 480
K = np.array([[572.4114, 0, 325.2611], [0, 573.57043, 242.04899], [0, 0, 1]]) # for lm
ZNEAR = 0.25
ZFAR = 6.0
depth_factor = 1000
LINEMOD_root = os.path.join(cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_refine')
LINEMOD_syn_root = os.path.join(cur_path, '../data/LINEMOD_6D/LM6d_converted/LM6d_refine_syn')
syn_poses_path = os.path.join(LINEMOD_syn_root, 'poses/LM6d_ds_train_observed_pose_all.pkl')
# output path
gt_observed_root_dir = os.path.join(LINEMOD_syn_root, 'data', 'gt_observed')
mkdir_if_missing(gt_observed_root_dir)
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))
