def get_data_pair_test_batch(pairdb, config):
"""
return a dict of train batch
:param segdb: ['image', 'flipped']
:param config: the config setting
:return: data, label, im_info
"""
im_observed, im_rendered, scale_ind_list = get_pair_image(
pairdb, config, "test")
if config.network.INPUT_DEPTH:
depth_observed, depth_rendered = get_pair_depth(
pairdb, config, scale_ind_list, "test")
if config.network.INPUT_MASK:
mask_observed, _, mask_rendered = get_pair_mask(
pairdb, config, scale_ind_list, "test")
im_info = [
np.array([pairdb[i]["height"], pairdb[i]["width"]], dtype=np.float32)
for i in range(len(pairdb))
]
num_pair = len(pairdb)
for i in range(num_pair):
class_index_tensor = [[] for i in range(num_pair)]
class_index_tensor[i] = np.array(
config.dataset.class_name.index(pairdb[i]["gt_class"])).reshape(1)
class_index_array = my_tensor_vstack(class_index_tensor)
data = []
for i in range(len(pairdb)):
cur_batch = {
"image_observed": im_observed[i],
"image_rendered": im_rendered[i],
"src_pose": np.array(pairdb[i]["pose_rendered"]).reshape(
(1, 3, 4)),
"class_index": class_index_array,
}
if config.network.INPUT_DEPTH:
cur_batch["depth_observed"] = depth_observed[i]
cur_batch["depth_rendered"] = depth_rendered[i]
if config.network.INPUT_MASK:
cur_batch["mask_observed"] = mask_observed[i]
cur_batch["mask_rendered"] = mask_rendered[i]
data.append(cur_batch)
label = {}
return data, label, im_info
def get_data_pair_train_batch(pairdb, config):
"""
return a dict of train batch
:param pairdb: ['image_observed', 'image_rendered', 'height', 'width',
'depth_observed', 'depth_rendered', 'pose_observed', 'pose_rendered',
'flipped']
:param config: ['INTRINSIC_MATRIX', 'DEPTH_FACTOR', ...]
:return: data: ['image_observed', 'image_rendered', 'depth_observed', 'depth_rendered']
label: ['flow', 'flow_weights'(RT)]
"""
num_pair = len(pairdb)
random_k = np.random.randint(18)
im_observed, im_rendered, scale_ind_list = get_pair_image(
pairdb, config, phase="train", random_k=random_k)
depth_gt_observed = get_gt_observed_depth(pairdb,
config,
scale_ind_list,
random_k=random_k)
if config.network.INPUT_DEPTH:
depth_observed, depth_rendered = get_pair_depth(pairdb,
config,
scale_ind_list,
phase="train",
random_k=random_k)
if config.network.PRED_MASK or config.network.INPUT_MASK:
mask_observed, mask_gt_observed, mask_rendered = get_pair_mask(
pairdb, config, scale_ind_list, phase="train", random_k=random_k)
im_observed_array = my_tensor_vstack(im_observed)
im_rendered_array = my_tensor_vstack(im_rendered)
depth_gt_observed_array = my_tensor_vstack(depth_gt_observed)
if config.network.INPUT_DEPTH:
depth_observed_array = my_tensor_vstack(depth_observed)
depth_rendered_array = my_tensor_vstack(depth_rendered)
if config.network.INPUT_MASK or config.network.PRED_MASK:
mask_observed_array = my_tensor_vstack(mask_observed)
mask_gt_observed_array = my_tensor_vstack(mask_gt_observed)
mask_rendered_array = my_tensor_vstack(mask_rendered)
if config.network.PRED_FLOW:
flow_i2r, flow_i2r_weights, X_rendered_tensor, X_observed_tensor = get_pair_flow(
pairdb, config, scale_ind_list, phase="train", random_k=random_k)
flow_i2r_array = my_tensor_vstack(flow_i2r)
flow_i2r_weights_array = my_tensor_vstack(flow_i2r_weights)
if config.train_iter.SE3_PM_LOSS:
X_obj, X_obj_weights = get_point_cloud_model(config, pairdb)
X_obj_array = my_tensor_vstack(X_obj)
X_obj_weights_array = my_tensor_vstack(X_obj_weights)
# SE3
rot_i2r_tensor = [[] for i in range(num_pair)]
trans_i2r_tensor = [[] for i in range(num_pair)]
src_pose_tensor = [[] for i in range(num_pair)]
tgt_pose_tensor = [[] for i in range(num_pair)]
tgt_points_tensor = [[] for i in range(num_pair)]
for i in range(num_pair):
rot_i2r, trans_i2r = calc_RT_delta(
pairdb[i]["pose_rendered"],
pairdb[i]["pose_observed"],
config.dataset.trans_means,
config.dataset.trans_stds,
config.network.ROT_COORD,
config.network.ROT_TYPE,
)
rot_i2r_tensor[i] = np.array(rot_i2r).reshape((1, -1))
trans_i2r_tensor[i] = np.array(trans_i2r).reshape((1, -1))
src_pose_tensor[i] = np.array(pairdb[i]["pose_rendered"]).reshape(
(1, 3, 4))
tgt_pose_tensor[i] = np.array(pairdb[i]["pose_observed"]).reshape(
(1, 3, 4))
if config.train_iter.SE3_PM_LOSS:
tgt_points = get_point_cloud_observed(
config,
X_obj_array[i],
pose_observed=np.array(pairdb[i]["pose_observed"]),
)
tgt_points_tensor[i] = np.expand_dims(tgt_points,
axis=0) # (1, 3, 3000)
rot_i2r_array = my_tensor_vstack(rot_i2r_tensor)
trans_i2r_array = my_tensor_vstack(trans_i2r_tensor)
src_pose_array = my_tensor_vstack(src_pose_tensor) # before refinment
tgt_pose_array = my_tensor_vstack(tgt_pose_tensor) # refinement target
tgt_points_array = my_tensor_vstack(tgt_points_tensor)
for i in range(num_pair):
class_index_tensor = [[] for i in range(num_pair)]
class_index_tensor[i] = np.array(
config.dataset.class_name.index(
pairdb[i]["gt_class"])).reshape(num_pair)
class_index_array = my_tensor_vstack(class_index_tensor)
data = {
"image_observed": im_observed_array,
"image_rendered": im_rendered_array,
"depth_gt_observed": depth_gt_observed_array,
"class_index": class_index_array,
"src_pose": src_pose_array,
"tgt_pose": tgt_pose_array,
}
if config.network.INPUT_DEPTH:
data["depth_observed"] = depth_observed_array
data["depth_rendered"] = depth_rendered_array
if config.network.INPUT_MASK:
data["mask_observed"] = mask_observed_array
data["mask_rendered"] = mask_rendered_array
label = {"rot": rot_i2r_array, "trans": trans_i2r_array}
if config.network.PRED_MASK:
label["mask_gt_observed"] = mask_gt_observed_array
if config.network.PRED_FLOW:
label["flow"] = flow_i2r_array
label["flow_weights"] = flow_i2r_weights_array
if config.train_iter.SE3_PM_LOSS:
label["point_cloud_model"] = X_obj_array
label["point_cloud_weights"] = X_obj_weights_array
label["point_cloud_observed"] = tgt_points_array
return {"data": data, "label": label}