def __init__(self, big_cfg, height, width):
self.big_cfg = big_cfg
self.model_dir = big_cfg.dataset.model_dir
self.rot_coord = big_cfg.network.ROT_COORD
self.pixel_means = big_cfg.network.PIXEL_MEANS[[2, 1, 0]]
self.pixel_means = self.pixel_means.reshape([3, 1,
1]).astype(np.float32)
self.K = big_cfg.dataset.INTRINSIC_MATRIX
self.T_means = big_cfg.dataset.trans_means
self.T_stds = big_cfg.dataset.trans_stds
self.height = height
self.width = width
self.zNear = big_cfg.dataset.ZNEAR
self.zFar = big_cfg.dataset.ZFAR
self.render_machine = None
if big_cfg.dataset.dataset.startswith("ModelNet"):
self.modelnet_root = big_cfg.modelnet_root
self.texture_path = os.path.join(self.modelnet_root,
"gray_texture.png")
from lib.render_glumpy.render_py_light_modelnet_multi import (
Render_Py_Light_ModelNet_Multi, )
self.model_path_list = [
os.path.join(self.model_dir, "{}.obj".format(model_name))
for model_name in big_cfg.dataset.class_name
]
self.render_machine = Render_Py_Light_ModelNet_Multi(
self.model_path_list,
self.texture_path,
self.K,
self.width,
self.height,
self.zNear,
self.zFar,
brightness_ratios=[0.7],
)
else:
self.render_machine = Render_Py(
self.model_dir,
big_cfg.dataset.class_name,
self.K,
self.width,
self.height,
self.zNear,
self.zFar,
)
self.reinit = True
self.batch_size = big_cfg.TRAIN.BATCH_PAIRS # will update according to data
self.Kinv = np.linalg.inv(np.matrix(self.K))
print("build render_machine: ", self.render_machine)
class batchUpdaterPyMulti:
def __init__(self, big_cfg, height, width):
self.big_cfg = big_cfg
self.model_dir = big_cfg.dataset.model_dir
self.rot_coord = big_cfg.network.ROT_COORD
self.pixel_means = big_cfg.network.PIXEL_MEANS[[2, 1, 0]]
self.pixel_means = self.pixel_means.reshape([3, 1,
1]).astype(np.float32)
self.K = big_cfg.dataset.INTRINSIC_MATRIX
self.T_means = big_cfg.dataset.trans_means
self.T_stds = big_cfg.dataset.trans_stds
self.height = height
self.width = width
self.zNear = big_cfg.dataset.ZNEAR
self.zFar = big_cfg.dataset.ZFAR
self.render_machine = None
if big_cfg.dataset.dataset.startswith("ModelNet"):
self.modelnet_root = big_cfg.modelnet_root
self.texture_path = os.path.join(self.modelnet_root,
"gray_texture.png")
from lib.render_glumpy.render_py_light_modelnet_multi import (
Render_Py_Light_ModelNet_Multi, )
self.model_path_list = [
os.path.join(self.model_dir, "{}.obj".format(model_name))
for model_name in big_cfg.dataset.class_name
]
self.render_machine = Render_Py_Light_ModelNet_Multi(
self.model_path_list,
self.texture_path,
self.K,
self.width,
self.height,
self.zNear,
self.zFar,
brightness_ratios=[0.7],
)
else:
self.render_machine = Render_Py(
self.model_dir,
big_cfg.dataset.class_name,
self.K,
self.width,
self.height,
self.zNear,
self.zFar,
)
self.reinit = True
self.batch_size = big_cfg.TRAIN.BATCH_PAIRS # will update according to data
self.Kinv = np.linalg.inv(np.matrix(self.K))
print("build render_machine: ", self.render_machine)
def get_names(self, big_cfg):
"""
:param small_cfg:
:return:
"""
pred = ["image_observed", "image_rendered"]
# pred = []
if big_cfg.network.PRED_FLOW:
pred.append("flow_est_crop")
pred.append("flow_loss")
pred.append("rot_est")
pred.append("rot_gt")
pred.append("trans_est")
pred.append("trans_gt")
if big_cfg.train_iter.SE3_DIST_LOSS:
pred.append("rot_loss")
pred.append("trans_loss")
if big_cfg.train_iter.SE3_PM_LOSS:
pred.append("point_matching_loss")
if self.big_cfg["network"]["PRED_MASK"]:
pred.append("zoom_mask_prob")
pred.append("zoom_mask_gt_observed")
pred.append("mask_pred") # unzoomed
return pred
def forward(self, data_batch, preds, big_cfg):
"""
:param data:
image_observed
image_rendered
depth_gt_observed
- depth_observed
- depth_rendered
- mask_real
src_pose
tgt_pose
:param label:
rot_i2r
trans_i2r
- flow_i2r
- flow_i2r_weights
- point_cloud_model
- point_cloud_weights
- point_cloud_real
:param preds:
image_observed
image_rendered
- flow_i2r_est
- flow_i2r_loss
rot_i2r
trans_i2r
- rot_i2r_loss
- trans_i2r_loss
- point_matching_loss
:return updated_batch:
"""
data_array = data_batch.data
label_array = data_batch.label
num_ctx = len(data_array)
pred_names = self.get_names(big_cfg)
init_time = 0
render_time = 0
image_time = 0
flow_time = 0
update_time = 0
mask_time = 0
io_time = 0
data_names = [x[0] for x in data_batch.provide_data[0]]
label_names = [x[0] for x in data_batch.provide_label[0]]
src_pose_all = [
data_array[ctx_i][data_names.index("src_pose")].asnumpy()
for ctx_i in range(num_ctx)
]
tgt_pose_all = [
data_array[ctx_i][data_names.index("tgt_pose")].asnumpy()
for ctx_i in range(num_ctx)
]
class_index_all = [
data_array[ctx_i][data_names.index("class_index")].asnumpy()
for ctx_i in range(num_ctx)
]
t = time.time()
# print("pred lens: {}".format(len(preds)))
# for i in preds:
# print(i[0].shape)
rot_est_all = [
preds[pred_names.index("rot_est")][ctx_i].asnumpy()
for ctx_i in range(num_ctx)
]
trans_est_all = [
preds[pred_names.index("trans_est")][ctx_i].asnumpy()
for ctx_i in range(num_ctx)
]
init_time += time.time() - t
if self.big_cfg.network.PRED_FLOW:
depth_gt_observed_all = [
data_array[ctx_i][data_names.index(
"depth_gt_observed")].asnumpy() for ctx_i in range(num_ctx)
]
for ctx_i in range(num_ctx):
batch_size = data_array[ctx_i][0].shape[0]
assert batch_size == self.batch_size, "{} vs. {}".format(
batch_size, self.batch_size)
cur_ctx = data_array[ctx_i][0].context
t = time.time()
src_pose = src_pose_all[
ctx_i] # data_array[ctx_i][data_names.index('src_pose')].asnumpy()
tgt_pose = tgt_pose_all[
ctx_i] # data_array[ctx_i][data_names.index('tgt_pose')].asnumpy()
if self.big_cfg.network.PRED_FLOW:
depth_gt_observed = depth_gt_observed_all[
ctx_i] # data_array[ctx_i][data_names.index('depth_gt_observed')] # ndarray
class_index = class_index_all[
ctx_i] # data_array[ctx_i][data_names.index('class_index')].asnumpy()
rot_est = rot_est_all[
ctx_i] # preds[pred_names.index('rot_est')][ctx_i].asnumpy()
trans_est = trans_est_all[
ctx_i] # preds[pred_names.index('trans_est')][ctx_i].asnumpy()
init_time += time.time() - t
refined_image_array = np.zeros(
(batch_size, 3, self.height, self.width))
refined_depth_array = np.zeros(
(batch_size, 1, self.height, self.width))
rot_res_array = np.zeros((batch_size, 4))
trans_res_array = np.zeros((batch_size, 3))
refined_pose_array = np.zeros((batch_size, 3, 4))
KT_array = np.zeros((batch_size, 3, 4))
for batch_idx in range(batch_size):
pre_pose = np.squeeze(src_pose[batch_idx])
r_delta = np.squeeze(rot_est[batch_idx])
t_delta = np.squeeze(trans_est[batch_idx])
refined_pose = RT_transform.RT_transform(
pre_pose,
r_delta,
t_delta,
self.T_means,
self.T_stds,
rot_coord=self.rot_coord,
)
t = time.time()
if not self.big_cfg.dataset.dataset.startswith("ModelNet"):
refined_image, refined_depth = self.render_machine.render(
class_index[batch_idx].astype("int"),
refined_pose[:3, :3],
refined_pose[:3, 3],
r_type="mat",
)
else:
idx = 2 # random.randint(0, 100)
# 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] += refined_pose[0, 3]
light_position[1] -= refined_pose[1, 3]
light_position[2] -= refined_pose[2, 3]
# print("light_position b: {}".format(light_position))
colors = np.array([1, 1, 1]) # white light
intensity = np.random.uniform(0.9, 1.1, size=(3, ))
colors_randk = 0 # 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 = 0 # random.randint(0, len(brightness_ratios) - 1)
refined_image, refined_depth = self.render_machine.render(
class_index[batch_idx].astype("int"),
refined_pose[:3, :3],
refined_pose[:3, 3],
light_position,
light_intensity,
brightness_k=rm_randk,
r_type="mat",
)
render_time += time.time() - t
# process refined_image
t = time.time()
refined_image = (refined_image[:, :, [2, 1, 0]].transpose(
[2, 0, 1]).astype(np.float32))
refined_image -= self.pixel_means
image_time += time.time() - t
# get se3_res
rot_res, trans_res = RT_transform.calc_RT_delta(
refined_pose,
np.squeeze(tgt_pose[batch_idx]),
self.T_means,
self.T_stds,
rot_coord=self.rot_coord,
rot_type="QUAT",
)
# print('{}, {}: {}, {}'.format(ctx_i, batch_idx, r_delta, rot_res))
refined_pose_array[batch_idx] = refined_pose
refined_image_array[batch_idx] = refined_image
refined_depth_array[batch_idx] = refined_depth.reshape(
(1, self.height, self.width))
rot_res_array[batch_idx] = rot_res
trans_res_array[batch_idx] = trans_res
se3_m = np.zeros([3, 4])
se3_rotm, se3_t = RT_transform.calc_se3(
refined_pose, np.squeeze(tgt_pose[batch_idx]))
se3_m[:, :3] = se3_rotm
se3_m[:, 3] = se3_t
KT_array[batch_idx] = np.dot(self.K, se3_m)
if self.big_cfg.network.PRED_MASK:
t = time.time()
refined_mask_rendered_array = np.zeros(
refined_depth_array.shape)
refined_mask_rendered_array[
refined_depth_array >
0.2] = 1 # if the mask_rendered input is depth
mask_time += time.time() - t
update_package = {
"image_rendered": refined_image_array,
"depth_rendered": refined_depth_array,
"src_pose": refined_pose_array,
"rot": rot_res_array,
"trans": trans_res_array,
}
if self.big_cfg.network.PRED_FLOW:
t = time.time()
gpu_flow_machine = gpu_flow_wrapper(cur_ctx.device_id)
# import matplotlib.pyplot as plt
# plt.figure()
# plt.subplot(1,2,1)
# plt.imshow(refined_depth_array[0,0])
# plt.subplot(1,2,2)
# plt.imshow(depth_gt_observed[0,0])
# plt.show()
refined_flow, refined_flow_valid = gpu_flow_machine(
refined_depth_array.astype(np.float32),
depth_gt_observed.astype(np.float32),
KT_array.astype(np.float32),
np.array(self.Kinv).astype(np.float32),
)
# problem with py3
# print('updater, flow: ', refined_flow.shape, np.unique(refined_flow))
# print('updater, flow weights: ', refined_flow_valid.shape, np.unique(refined_flow_valid))
# print('KT: ', KT_array[0])
# print('Kinv: ', self.Kinv)
flow_time += time.time() - t
refined_flow_weights = np.tile(refined_flow_valid,
[1, 2, 1, 1])
update_package["flow"] = refined_flow
update_package["flow_weights"] = refined_flow_weights
if self.big_cfg.network.INPUT_MASK:
update_package["mask_rendered"] = refined_mask_rendered_array
t = time.time()
data_array[ctx_i] = self.update_data_batch(data_array[ctx_i],
data_names,
update_package)
label_array[ctx_i] = self.update_data_batch(
label_array[ctx_i], label_names, update_package)
update_time += time.time() - t
t = time.time()
new_data_batch = mx.io.DataBatch(
data=data_array,
label=label_array,
pad=data_batch.pad,
index=data_batch.index,
provide_data=data_batch.provide_data,
provide_label=data_batch.provide_label,
)
io_time += time.time() - t
# print("---------------------------------")
# print("init_time: {:.3f} sec".format(init_time))
# print("render_time: {:.3f} sec".format(render_time))
# print("image_time: {:.3f} sec".format(image_time))
# print("flow_time: {:.3f} sec".format(flow_time))
# print("mask_time: {:.3f} sec".format(mask_time))
# print("update_time: {:.3f} sec".format(update_time))
# print("io_time: {:.3f} sec".format(io_time))
# print("all_time: {:.3f} sec".format(time.time() - t_all))
# print("---------------------------------")
return new_data_batch
def update_data_batch(self, data, data_names, update_package):
import mxnet.ndarray as nd
for blob_idx, blob_name in enumerate(data_names):
if blob_name not in update_package:
continue
# print('blob_idx: {}, blob_name: {} -- {}'.format(blob_idx, blob_name, np.max(update_package[blob_name])))
data[blob_idx] = nd.array(update_package[blob_name])
return data
def main():
sel_classes = classes
render_machine = Render_Py(model_dir, classes, K, width, height, ZNEAR,
ZFAR)
for cls_idx, cls_name in enumerate(classes):
if not cls_name in sel_classes:
continue
print(cls_idx, cls_name)
observed_indices = []
images = [
fn for fn in os.listdir(
os.path.join(LM6d_origin_root, '{:02d}'.format(
class2idx(cls_name)), 'rgb')) if '.png' in fn
]
images = sorted(images)
gt_path = os.path.join(LM6d_origin_root,
'{:02d}'.format(class2idx(cls_name)), 'gt.yml')
gt_dict = load_gt(gt_path)
info_path = os.path.join(LM6d_origin_root,
'{:02d}'.format(class2idx(cls_name)),
'info.yml')
info_dict = load_info(info_path)
for observed_img in tqdm(images):
old_color_path = os.path.join(LM6d_origin_root,
'{:02d}'.format(class2idx(cls_name)),
"rgb/{}".format(observed_img))
assert os.path.exists(old_color_path), old_color_path
old_depth_path = os.path.join(LM6d_origin_root,
'{:02d}'.format(class2idx(cls_name)),
"depth/{}".format(observed_img))
assert os.path.exists(old_depth_path), old_depth_path
img_id = int(observed_img.replace('.png', ''))
new_img_id = img_id + 1
# K
# K = np.array(info_dict[img_id]['cam_K']).reshape((3, 3))
color_img = cv2.imread(old_color_path, cv2.IMREAD_COLOR)
## depth
depth = read_img(old_depth_path, 1)
# print(np.max(depth), np.min(depth))
# print(color_img.shape)
new_color_path = os.path.join(
LM6d_new_root, '{:02d}'.format(class2idx(cls_name)),
"{:06d}-color.png".format(new_img_id))
new_depth_path = os.path.join(
LM6d_new_root, '{:02d}'.format(class2idx(cls_name)),
"{:06d}-depth.png".format(new_img_id))
mkdir_if_missing(os.path.dirname(new_color_path))
copyfile(old_color_path, new_color_path)
copyfile(old_depth_path, new_depth_path)
# meta and label
meta_dict = {}
num_instance = len(gt_dict[img_id])
meta_dict['cls_indexes'] = np.zeros((1, num_instance),
dtype=np.int32)
meta_dict['boxes'] = np.zeros((num_instance, 4), dtype='float32')
meta_dict['poses'] = np.zeros((3, 4, num_instance),
dtype='float32')
distances = []
label_dict = {}
for ins_id, instance in enumerate(gt_dict[img_id]):
obj_id = instance['obj_id']
meta_dict['cls_indexes'][0, ins_id] = obj_id
obj_bb = np.array(instance['obj_bb'])
meta_dict['boxes'][ins_id, :] = obj_bb
# pose
pose = np.zeros((3, 4))
R = np.array(instance['cam_R_m2c']).reshape((3, 3))
t = np.array(instance['cam_t_m2c']) / 1000. # mm -> m
pose[:3, :3] = R
pose[:3, 3] = t
distances.append(t[2])
meta_dict['poses'][:, :, ins_id] = pose
image_gl, depth_gl = render_machine.render(obj_id - 1,
pose[:3, :3],
pose[:3, 3],
r_type='mat')
image_gl = image_gl.astype('uint8')
label = np.zeros(depth_gl.shape)
label[depth_gl != 0] = 1
label_dict[obj_id] = label
meta_path = os.path.join(LM6d_new_root,
'{:02d}'.format(class2idx(cls_name)),
"{:06d}-meta.mat".format(new_img_id))
sio.savemat(meta_path, meta_dict)
dis_inds = sorted(
range(len(distances)),
key=lambda k: -distances[k]) # put deeper objects first
# label
res_label = np.zeros((480, 640))
for dis_id in dis_inds:
cls_id = meta_dict['cls_indexes'][0, dis_id]
tmp_label = label_dict[cls_id]
# label
res_label[tmp_label == 1] = cls_id
label_path = os.path.join(LM6d_new_root,
'{:02d}'.format(class2idx(cls_name)),
"{:06d}-label.png".format(new_img_id))
cv2.imwrite(label_path, res_label)
# observed idx
observed_indices.append("{:02d}/{:06d}".format(
class2idx(cls_name), new_img_id))
def main():
sel_classes = classes
model_dir = os.path.join(cur_dir, '../../data/LINEMOD_6D/LM6d_converted/models')
render_machine = Render_Py(model_dir, classes, K, width, height, ZNEAR, ZFAR)
for cls_idx, cls_name in enumerate(classes):
if not cls_name in sel_classes:
continue
print(cls_idx, cls_name)
real_indices = []
images = [fn for fn in os.listdir(os.path.join(LM6d_origin_root,
'{:02d}'.format(class2idx(cls_name)), 'rgb')) if '.png' in fn]
images = sorted(images)
gt_path = os.path.join(LM6d_origin_root, '{:02d}'.format(class2idx(cls_name)), 'gt.yml')
gt_dict = load_gt(gt_path)
info_path = os.path.join(LM6d_origin_root, '{:02d}'.format(class2idx(cls_name)), 'info.yml')
info_dict = load_info(info_path)
for real_img in tqdm(images):
old_color_path = os.path.join(LM6d_origin_root, '{:02d}'.format(class2idx(cls_name)), "rgb/{}".format(real_img))
assert os.path.exists(old_color_path), old_color_path
old_depth_path = os.path.join(LM6d_origin_root, '{:02d}'.format(class2idx(cls_name)), "depth/{}".format(real_img))
assert os.path.exists(old_depth_path), old_depth_path
img_id = int(real_img.replace('.png', ''))
new_img_id = img_id + 1
# K
# K = np.array(info_dict[img_id]['cam_K']).reshape((3, 3))
color_img = cv2.imread(old_color_path, cv2.IMREAD_COLOR)
## depth
depth = read_img(old_depth_path, 1)
# print(np.max(depth), np.min(depth))
# print(color_img.shape)
new_color_path = os.path.join(LM6d_new_root, '{:02d}'.format(class2idx(cls_name)),
"{:06d}-color.png".format(new_img_id))
new_depth_path = os.path.join(LM6d_new_root, '{:02d}'.format(class2idx(cls_name)),
"{:06d}-depth.png".format(new_img_id))
mkdir_if_missing(os.path.dirname(new_color_path))
copyfile(old_color_path, new_color_path)
copyfile(old_depth_path, new_depth_path)
# meta and label
meta_dict = {}
num_instance = len(gt_dict[img_id])
meta_dict['cls_indexes'] = np.zeros((1, num_instance), dtype=np.int32)
meta_dict['boxes'] = np.zeros((num_instance, 4), dtype='float32')
meta_dict['poses'] = np.zeros((3,4,num_instance), dtype='float32')
distances = []
label_dict = {}
for ins_id, instance in enumerate(gt_dict[img_id]):
obj_id = instance['obj_id']
meta_dict['cls_indexes'][0, ins_id] = obj_id
obj_bb = np.array(instance['obj_bb'])
meta_dict['boxes'][ins_id, :] = obj_bb
# pose
pose = np.zeros((3, 4))
R = np.array(instance['cam_R_m2c']).reshape((3, 3))
t = np.array(instance['cam_t_m2c']) / 1000. # mm -> m
pose[:3, :3] = R
pose[:3, 3] = t
distances.append(t[2])
meta_dict['poses'][:,:,ins_id] = pose
image_gl, depth_gl = render_machine.render(obj_id-1, pose[:3, :3], pose[:3, 3],
r_type='mat')
image_gl = image_gl.astype('uint8')
label = np.zeros(depth_gl.shape)
label[depth_gl!=0] = 1
label_dict[obj_id] = label
meta_path = os.path.join(LM6d_new_root, '{:02d}'.format(class2idx(cls_name)),
"{:06d}-meta.mat".format(new_img_id))
sio.savemat(meta_path, meta_dict)
dis_inds = sorted(range(len(distances)), key=lambda k: -distances[k]) # put deeper objects first
# label
res_label = np.zeros((480, 640))
for dis_id in dis_inds:
cls_id = meta_dict['cls_indexes'][0, dis_id]
tmp_label = label_dict[cls_id]
# label
res_label[tmp_label == 1] = cls_id
label_path = os.path.join(LM6d_new_root, '{:02d}'.format(class2idx(cls_name)),
"{:06d}-label.png".format(new_img_id))
cv2.imwrite(label_path, res_label)
def vis_check():
fig = plt.figure(figsize=(8, 6), dpi=120)
plt.subplot(2, 3, 1)
plt.imshow(color_img[:,:,[2,1,0]])
plt.title('color_img')
plt.subplot(2, 3, 2)
plt.imshow(depth_gl)
plt.title('depth')
plt.subplot(2, 3, 3)
plt.imshow(depth_gl)
plt.title('depth_gl')
plt.subplot(2, 3, 4)
plt.imshow(res_label)
plt.title('res_label')
plt.subplot(2,3,5)
label_v1_path = os.path.join('/data/wanggu/Storage/LINEMOD_SIXD_wods/LM6d_render_v1/data/real',
'{:02d}'.format(class2idx(cls_name)),
"{:06d}-label.png".format(new_img_id))
assert os.path.exists(label_v1_path), label_v1_path
label_v1 = read_img(label_v1_path, 1)
plt.imshow(label_v1)
plt.title('label_v1')
plt.show()
# vis_check()
# real idx
real_indices.append("{:02d}/{:06d}".format(class2idx(cls_name), new_img_id))
# one idx file for each video of each class
real_idx_file = os.path.join(real_set_dir, "{}_all.txt".format(cls_name))
with open(real_idx_file, 'w') as f:
for real_idx in real_indices:
f.write(real_idx + '\n')
def pred_eval(config,
predictor,
test_data,
imdb_test,
vis=False,
ignore_cache=None,
logger=None,
pairdb=None):
"""
wrapper for calculating offline validation for faster data analysis
in this example, all threshold are set by hand
:param predictor: Predictor
:param test_data: data iterator, must be non-shuffle
:param imdb_test: image database
:param vis: controls visualization
:param ignore_cache: ignore the saved cache file
:param logger: the logger instance
:return:
"""
logger.info(imdb_test.result_path)
logger.info("test iter size: {}".format(config.TEST.test_iter))
pose_err_file = os.path.join(
imdb_test.result_path,
imdb_test.name + "_pose_iter{}.pkl".format(config.TEST.test_iter))
if os.path.exists(pose_err_file) and not ignore_cache and not vis:
with open(pose_err_file, "rb") as fid:
if six.PY3:
[all_rot_err, all_trans_err, all_poses_est,
all_poses_gt] = cPickle.load(fid, encoding="latin1")
else:
[all_rot_err, all_trans_err, all_poses_est,
all_poses_gt] = cPickle.load(fid)
imdb_test.evaluate_pose(config, all_poses_est, all_poses_gt)
pose_add_plots_dir = os.path.join(imdb_test.result_path, "add_plots")
mkdir_p(pose_add_plots_dir)
imdb_test.evaluate_pose_add(config,
all_poses_est,
all_poses_gt,
output_dir=pose_add_plots_dir)
pose_arp2d_plots_dir = os.path.join(imdb_test.result_path,
"arp_2d_plots")
mkdir_p(pose_arp2d_plots_dir)
imdb_test.evaluate_pose_arp_2d(config,
all_poses_est,
all_poses_gt,
output_dir=pose_arp2d_plots_dir)
return
assert vis or not test_data.shuffle
assert config.TEST.BATCH_PAIRS == 1
if not isinstance(test_data, PrefetchingIter):
test_data = PrefetchingIter(test_data)
num_pairs = len(pairdb)
height = 480
width = 640
data_time, net_time, post_time = 0.0, 0.0, 0.0
sum_EPE_all = 0.0
num_inst_all = 0.0
sum_EPE_viz = 0.0
num_inst_viz = 0.0
sum_EPE_vizbg = 0.0
num_inst_vizbg = 0.0
sum_PoseErr = [
np.zeros((len(imdb_test.classes) + 1, 2))
for batch_idx in range(config.TEST.test_iter)
]
all_rot_err = [[[] for j in range(config.TEST.test_iter)]
for batch_idx in range(len(imdb_test.classes))
] # num_cls x test_iter
all_trans_err = [[[] for j in range(config.TEST.test_iter)]
for batch_idx in range(len(imdb_test.classes))]
all_poses_est = [[[] for j in range(config.TEST.test_iter)]
for batch_idx in range(len(imdb_test.classes))]
all_poses_gt = [[[] for j in range(config.TEST.test_iter)]
for batch_idx in range(len(imdb_test.classes))]
num_inst = np.zeros(len(imdb_test.classes) + 1)
K = config.dataset.INTRINSIC_MATRIX
if (config.TEST.test_iter > 1 or config.TEST.VISUALIZE) and True:
print(
"************* start setup render_glumpy environment... ******************"
)
if config.dataset.dataset.startswith("ModelNet"):
from lib.render_glumpy.render_py_light_modelnet_multi import Render_Py_Light_ModelNet_Multi
modelnet_root = config.modelnet_root
texture_path = os.path.join(modelnet_root, "gray_texture.png")
model_path_list = [
os.path.join(config.dataset.model_dir,
"{}.obj".format(model_name))
for model_name in config.dataset.class_name
]
render_machine = Render_Py_Light_ModelNet_Multi(
model_path_list,
texture_path,
K,
width,
height,
config.dataset.ZNEAR,
config.dataset.ZFAR,
brightness_ratios=[0.7],
)
else:
render_machine = Render_Py(
config.dataset.model_dir,
config.dataset.class_name,
K,
width,
height,
config.dataset.ZNEAR,
config.dataset.ZFAR,
)
def render(render_machine, pose, cls_idx, K=None):
if config.dataset.dataset.startswith("ModelNet"):
idx = 2
# 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]
colors = np.array([1, 1, 1]) # white light
intensity = np.random.uniform(0.9, 1.1, size=(3, ))
colors_randk = 0
light_intensity = colors[colors_randk] * intensity
# randomly choose a render machine
rm_randk = 0 # random.randint(0, len(brightness_ratios) - 1)
rgb_gl, depth_gl = render_machine.render(
cls_idx,
pose[:3, :3],
pose[:3, 3],
light_position,
light_intensity,
brightness_k=rm_randk,
r_type="mat",
)
rgb_gl = rgb_gl.astype("uint8")
else:
rgb_gl, depth_gl = render_machine.render(cls_idx,
pose[:3, :3],
pose[:, 3],
r_type="mat",
K=K)
rgb_gl = rgb_gl.astype("uint8")
return rgb_gl, depth_gl
print(
"***************setup render_glumpy environment succeed ******************"
)
if config.TEST.PRECOMPUTED_ICP:
print("precomputed_ICP")
config.TEST.test_iter = 1
all_rot_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
all_trans_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
all_poses_est = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
all_poses_gt = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
xy_trans_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
z_trans_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
for idx in range(len(pairdb)):
pose_path = pairdb[idx]["depth_rendered"][:-10] + "-pose_icp.txt"
pose_rendered_update = np.loadtxt(pose_path, skiprows=1)
pose_observed = pairdb[idx]["pose_observed"]
r_dist_est, t_dist_est = calc_rt_dist_m(pose_rendered_update,
pose_observed)
xy_dist = np.linalg.norm(pose_rendered_update[:2, -1] -
pose_observed[:2, -1])
z_dist = np.linalg.norm(pose_rendered_update[-1, -1] -
pose_observed[-1, -1])
print(
"{}: r_dist_est: {}, t_dist_est: {}, xy_dist: {}, z_dist: {}".
format(idx, r_dist_est, t_dist_est, xy_dist, z_dist))
class_id = imdb_test.classes.index(pairdb[idx]["gt_class"])
# store poses estimation and gt
all_poses_est[class_id][0].append(pose_rendered_update)
all_poses_gt[class_id][0].append(pairdb[idx]["pose_observed"])
all_rot_err[class_id][0].append(r_dist_est)
all_trans_err[class_id][0].append(t_dist_est)
xy_trans_err[class_id][0].append(xy_dist)
z_trans_err[class_id][0].append(z_dist)
all_rot_err = np.array(all_rot_err)
all_trans_err = np.array(all_trans_err)
print("rot = {} +/- {}".format(np.mean(all_rot_err[class_id][0]),
np.std(all_rot_err[class_id][0])))
print("trans = {} +/- {}".format(np.mean(all_trans_err[class_id][0]),
np.std(all_trans_err[class_id][0])))
num_list = all_trans_err[class_id][0]
print("xyz: {:.2f} +/- {:.2f}".format(
np.mean(num_list) * 100,
np.std(num_list) * 100))
num_list = xy_trans_err[class_id][0]
print("xy: {:.2f} +/- {:.2f}".format(
np.mean(num_list) * 100,
np.std(num_list) * 100))
num_list = z_trans_err[class_id][0]
print("z: {:.2f} +/- {:.2f}".format(
np.mean(num_list) * 100,
np.std(num_list) * 100))
imdb_test.evaluate_pose(config, all_poses_est, all_poses_gt)
pose_add_plots_dir = os.path.join(imdb_test.result_path,
"add_plots_precomputed_ICP")
mkdir_p(pose_add_plots_dir)
imdb_test.evaluate_pose_add(config,
all_poses_est,
all_poses_gt,
output_dir=pose_add_plots_dir)
pose_arp2d_plots_dir = os.path.join(imdb_test.result_path,
"arp_2d_plots_precomputed_ICP")
mkdir_p(pose_arp2d_plots_dir)
imdb_test.evaluate_pose_arp_2d(config,
all_poses_est,
all_poses_gt,
output_dir=pose_arp2d_plots_dir)
return
if config.TEST.BEFORE_ICP:
print("before_ICP")
config.TEST.test_iter = 1
all_rot_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
all_trans_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
all_poses_est = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
all_poses_gt = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
xy_trans_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
z_trans_err = [[[] for j in range(1)]
for batch_idx in range(len(imdb_test.classes))]
for idx in range(len(pairdb)):
pose_path = pairdb[idx]["depth_rendered"][:-10] + "-pose.txt"
pose_rendered_update = np.loadtxt(pose_path, skiprows=1)
pose_observed = pairdb[idx]["pose_observed"]
r_dist_est, t_dist_est = calc_rt_dist_m(pose_rendered_update,
pose_observed)
xy_dist = np.linalg.norm(pose_rendered_update[:2, -1] -
pose_observed[:2, -1])
z_dist = np.linalg.norm(pose_rendered_update[-1, -1] -
pose_observed[-1, -1])
class_id = imdb_test.classes.index(pairdb[idx]["gt_class"])
# store poses estimation and gt
all_poses_est[class_id][0].append(pose_rendered_update)
all_poses_gt[class_id][0].append(pairdb[idx]["pose_observed"])
all_rot_err[class_id][0].append(r_dist_est)
all_trans_err[class_id][0].append(t_dist_est)
xy_trans_err[class_id][0].append(xy_dist)
z_trans_err[class_id][0].append(z_dist)
all_trans_err = np.array(all_trans_err)
imdb_test.evaluate_pose(config, all_poses_est, all_poses_gt)
pose_add_plots_dir = os.path.join(imdb_test.result_path,
"add_plots_before_ICP")
mkdir_p(pose_add_plots_dir)
imdb_test.evaluate_pose_add(config,
all_poses_est,
all_poses_gt,
output_dir=pose_add_plots_dir)
pose_arp2d_plots_dir = os.path.join(imdb_test.result_path,
"arp_2d_plots_before_ICP")
mkdir_p(pose_arp2d_plots_dir)
imdb_test.evaluate_pose_arp_2d(config,
all_poses_est,
all_poses_gt,
output_dir=pose_arp2d_plots_dir)
return
# ------------------------------------------------------------------------------
t_start = time.time()
t = time.time()
for idx, data_batch in enumerate(test_data):
if np.sum(pairdb[idx]
["pose_rendered"]) == -12: # NO POINT VALID IN INIT POSE
print(idx)
class_id = imdb_test.classes.index(pairdb[idx]["gt_class"])
for pose_iter_idx in range(config.TEST.test_iter):
all_poses_est[class_id][pose_iter_idx].append(
pairdb[idx]["pose_rendered"])
all_poses_gt[class_id][pose_iter_idx].append(
pairdb[idx]["pose_observed"])
r_dist = 1000
t_dist = 1000
all_rot_err[class_id][pose_iter_idx].append(r_dist)
all_trans_err[class_id][pose_iter_idx].append(t_dist)
sum_PoseErr[pose_iter_idx][class_id, :] += np.array(
[r_dist, t_dist])
sum_PoseErr[pose_iter_idx][-1, :] += np.array([r_dist, t_dist])
# post process
if idx % 50 == 0:
logger.info(
"testing {}/{} data {:.4f}s net {:.4f}s calc_gt {:.4f}s".
format(
(idx + 1),
num_pairs,
data_time / ((idx + 1) * test_data.batch_size),
net_time / ((idx + 1) * test_data.batch_size),
post_time / ((idx + 1) * test_data.batch_size),
))
print("in test: NO POINT_VALID IN rendered")
continue
data_time += time.time() - t
t = time.time()
pose_rendered = pairdb[idx]["pose_rendered"]
if np.sum(pose_rendered) == -12:
print(idx)
class_id = imdb_test.classes.index(pairdb[idx]["gt_class"])
num_inst[class_id] += 1
num_inst[-1] += 1
for pose_iter_idx in range(config.TEST.test_iter):
all_poses_est[class_id][pose_iter_idx].append(pose_rendered)
all_poses_gt[class_id][pose_iter_idx].append(
pairdb[idx]["pose_observed"])
# post process
if idx % 50 == 0:
logger.info(
"testing {}/{} data {:.4f}s net {:.4f}s calc_gt {:.4f}s".
format(
(idx + 1),
num_pairs,
data_time / ((idx + 1) * test_data.batch_size),
net_time / ((idx + 1) * test_data.batch_size),
post_time / ((idx + 1) * test_data.batch_size),
))
t = time.time()
continue
output_all = predictor.predict(data_batch)
net_time += time.time() - t
t = time.time()
rst_iter = []
for output in output_all:
cur_rst = {}
cur_rst["se3"] = np.squeeze(
output["se3_output"].asnumpy()).astype("float32")
if not config.TEST.FAST_TEST and config.network.PRED_FLOW:
cur_rst["flow"] = np.squeeze(
output["flow_est_crop_output"].asnumpy().transpose(
(2, 3, 1, 0))).astype("float16")
else:
cur_rst["flow"] = None
if config.network.PRED_MASK and config.TEST.UPDATE_MASK not in [
"init", "box_rendered"
]:
mask_pred = np.squeeze(
output["mask_observed_pred_output"].asnumpy()).astype(
"float32")
cur_rst["mask_pred"] = mask_pred
rst_iter.append(cur_rst)
post_time += time.time() - t
# sample_ratio = 1 # 0.01
for batch_idx in range(0, test_data.batch_size):
# if config.TEST.VISUALIZE and not (r_dist>15 and t_dist>0.05):
# continue # 3388, 5326
# calculate the flow error --------------------------------------------
t = time.time()
if config.network.PRED_FLOW and not config.TEST.FAST_TEST:
# evaluate optical flow
flow_gt = par_generate_gt(config, pairdb[idx])
if config.network.PRED_FLOW:
all_diff = calc_EPE_one_pair(rst_iter[batch_idx], flow_gt,
"flow")
sum_EPE_all += all_diff["epe_all"]
num_inst_all += all_diff["num_all"]
sum_EPE_viz += all_diff["epe_viz"]
num_inst_viz += all_diff["num_viz"]
sum_EPE_vizbg += all_diff["epe_vizbg"]
num_inst_vizbg += all_diff["num_vizbg"]
# calculate the se3 error ---------------------------------------------
# evaluate se3 estimation
pose_rendered = pairdb[idx]["pose_rendered"]
class_id = imdb_test.classes.index(pairdb[idx]["gt_class"])
num_inst[class_id] += 1
num_inst[-1] += 1
post_time += time.time() - t
# iterative refine se3 estimation --------------------------------------------------
for pose_iter_idx in range(config.TEST.test_iter):
t = time.time()
pose_rendered_update = RT_transform(
pose_rendered,
rst_iter[0]["se3"][:-3],
rst_iter[0]["se3"][-3:],
config.dataset.trans_means,
config.dataset.trans_stds,
config.network.ROT_COORD,
)
# calculate error
r_dist, t_dist = calc_rt_dist_m(pose_rendered_update,
pairdb[idx]["pose_observed"])
# store poses estimation and gt
all_poses_est[class_id][pose_iter_idx].append(
pose_rendered_update)
all_poses_gt[class_id][pose_iter_idx].append(
pairdb[idx]["pose_observed"])
all_rot_err[class_id][pose_iter_idx].append(r_dist)
all_trans_err[class_id][pose_iter_idx].append(t_dist)
sum_PoseErr[pose_iter_idx][class_id, :] += np.array(
[r_dist, t_dist])
sum_PoseErr[pose_iter_idx][-1, :] += np.array([r_dist, t_dist])
if config.TEST.VISUALIZE:
print("idx {}, iter {}: rError: {}, tError: {}".format(
idx + batch_idx, pose_iter_idx + 1, r_dist, t_dist))
post_time += time.time() - t
# # if more than one iteration
if pose_iter_idx < (config.TEST.test_iter -
1) or config.TEST.VISUALIZE:
t = time.time()
# get refined image
K_path = pairdb[idx]["image_observed"][:-10] + "-K.txt"
if os.path.exists(K_path):
K = np.loadtxt(K_path)
image_refined, depth_refined = render(
render_machine,
pose_rendered_update,
config.dataset.class_name.index(
pairdb[idx]["gt_class"]),
K=K,
)
image_refined = image_refined[:, :, :3]
# update minibatch
update_package = [{
"image_rendered": image_refined,
"src_pose": pose_rendered_update
}]
if config.network.INPUT_DEPTH:
update_package[0]["depth_rendered"] = depth_refined
if config.network.INPUT_MASK:
mask_rendered_refined = np.zeros(depth_refined.shape)
mask_rendered_refined[depth_refined > 0.2] = 1
update_package[0][
"mask_rendered"] = mask_rendered_refined
if config.network.PRED_MASK:
# init, box_rendered, mask_rendered, box_observed, mask_observed
if config.TEST.UPDATE_MASK == "box_rendered":
input_names = [
blob_name[0]
for blob_name in data_batch.provide_data[0]
]
update_package[0][
"mask_observed"] = np.squeeze(
data_batch.data[0][input_names.index(
"mask_rendered")].asnumpy()
[batch_idx]) # noqa
elif config.TEST.UPDATE_MASK == "init":
pass
else:
raise Exception(
"Unknown UPDATE_MASK type: {}".format(
config.network.UPDATE_MASK))
pose_rendered = pose_rendered_update
data_batch = update_data_batch(config, data_batch,
update_package)
data_time += time.time() - t
# forward and get rst
if pose_iter_idx < config.TEST.test_iter - 1:
t = time.time()
output_all = predictor.predict(data_batch)
net_time += time.time() - t
t = time.time()
rst_iter = []
for output in output_all:
cur_rst = {}
if config.network.REGRESSOR_NUM == 1:
cur_rst["se3"] = np.squeeze(
output["se3_output"].asnumpy()).astype(
"float32")
if not config.TEST.FAST_TEST and config.network.PRED_FLOW:
cur_rst["flow"] = np.squeeze(
output["flow_est_crop_output"].asnumpy().
transpose((2, 3, 1, 0))).astype("float16")
else:
cur_rst["flow"] = None
if config.network.PRED_MASK and config.TEST.UPDATE_MASK not in [
"init", "box_rendered"
]:
mask_pred = np.squeeze(
output["mask_observed_pred_output"].
asnumpy()).astype("float32")
cur_rst["mask_pred"] = mask_pred
rst_iter.append(cur_rst)
post_time += time.time() - t
# post process
if idx % 50 == 0:
logger.info(
"testing {}/{} data {:.4f}s net {:.4f}s calc_gt {:.4f}s".
format(
(idx + 1),
num_pairs,
data_time / ((idx + 1) * test_data.batch_size),
net_time / ((idx + 1) * test_data.batch_size),
post_time / ((idx + 1) * test_data.batch_size),
))
t = time.time()
all_rot_err = np.array(all_rot_err)
all_trans_err = np.array(all_trans_err)
# save inference results
if not config.TEST.VISUALIZE:
with open(pose_err_file, "wb") as f:
logger.info("saving result cache to {}".format(pose_err_file))
cPickle.dump(
[all_rot_err, all_trans_err, all_poses_est, all_poses_gt],
f,
protocol=2)
logger.info("done")
if config.network.PRED_FLOW:
logger.info("evaluate flow:")
logger.info("EPE all: {}".format(sum_EPE_all / max(num_inst_all, 1.0)))
logger.info("EPE ignore unvisible: {}".format(
sum_EPE_vizbg / max(num_inst_vizbg, 1.0)))
logger.info("EPE visible: {}".format(sum_EPE_viz /
max(num_inst_viz, 1.0)))
logger.info("evaluate pose:")
imdb_test.evaluate_pose(config, all_poses_est, all_poses_gt)
# evaluate pose add
pose_add_plots_dir = os.path.join(imdb_test.result_path, "add_plots")
mkdir_p(pose_add_plots_dir)
imdb_test.evaluate_pose_add(config,
all_poses_est,
all_poses_gt,
output_dir=pose_add_plots_dir)
pose_arp2d_plots_dir = os.path.join(imdb_test.result_path, "arp_2d_plots")
mkdir_p(pose_arp2d_plots_dir)
imdb_test.evaluate_pose_arp_2d(config,
all_poses_est,
all_poses_gt,
output_dir=pose_arp2d_plots_dir)
logger.info("using {} seconds in total".format(time.time() - t_start))
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__":
big_classes = sel_classes
classes = ['024_bowl', '036_wood_block',
'051_large_clamp', '052_extra_large_clamp', '061_foam_brick']
model_folder = './data/LOV/models'
print('init render machine...')
render_machine = Render_Py(model_folder, big_classes, K, width, height, zNear, zFar)
for cls_idx, cls_name in enumerate(big_classes):
if cls_name in classes:
continue
print(cls_name)
with open('./data/render_v5/image_set/train_{}.txt'.format(cls_name), 'r') as f:
real_indices = [line.strip().split()[0] for line in f.readlines()]
img_indices = []
for i in [0, 100]:
img_indices.append(real_indices[i])
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)