def collate_fn(self, batch):
obj_data = []
for data_n in batch:
_, _, obs = data_n
obj_data_ = parse_obs_data(obs, parse_joints=True)
obj_data.append(obj_data_)
obj_data = tc.concatenate(obj_data)
return obj_data
def get_predictions(self,
obj_predictor,
use_gt_joints=True,
predict_keypoints=True,
n_iterations=10):
predictions = defaultdict(list)
for data in tqdm(self.dataloader):
images = data['images'].cuda().float().permute(0, 3, 1, 2) / 255
detections = data['detections']
cameras = data['cameras'].cuda().float()
joints = data['joints']
stacked_joints = dict()
for k, v in joints.items():
stacked_joints[k] = torch.stack(v).cuda().float()
joints = stacked_joints
assert len(obj_predictor.model.robots) == 1
robot = obj_predictor.model.robots[0]
urdf_layer = obj_predictor.model.get_urdf_layer(robot.name)
joints_tensor = urdf_layer.to_tensor(joints)
detections.register_tensor('joints', joints_tensor)
detections.infos['joint_names'] = urdf_layer.joint_names[
None].repeat(len(detections), axis=0).tolist()
K = cameras.K.float()
detections = detections.cuda().float()
_, preds = obj_predictor.get_predictions(
images=images,
K=K,
detections=detections,
use_gt_joints=use_gt_joints,
gt_joints=joints_tensor,
n_iterations=n_iterations,
)
if predict_keypoints:
for pred_k, pred_v in preds.items():
TCO_keypoints_3d, keypoints_2d = obj_predictor.pred_keypoints(
pred_v)
orig_wh = cameras.orig_wh[pred_v.infos.batch_im_id]
cropresize_bboxes = cameras.cropresize_bboxes[
pred_v.infos.batch_im_id]
wh = torch.zeros_like(orig_wh)
wh[:, 1] = images.shape[2]
wh[:, 0] = images.shape[3]
keypoints_2d = cropresize_backtransform_points2d(
orig_wh, cropresize_bboxes, wh, keypoints_2d)
pred_v.register_tensor('keypoints_2d', keypoints_2d)
pred_v.register_tensor('TCO_keypoints_3d',
TCO_keypoints_3d)
predictions[pred_k].append(pred_v)
for k, v in predictions.items():
predictions[k] = tc.concatenate(predictions[k])
return predictions
def batched_model_predictions(self,
model,
images,
K,
obj_data,
n_iterations=1):
timer = Timer()
timer.start()
ids = torch.arange(len(obj_data))
ds = TensorDataset(ids)
dl = DataLoader(ds, batch_size=self.bsz_objects)
preds = defaultdict(list)
for (batch_ids, ) in dl:
timer.resume()
obj_inputs = obj_data[batch_ids.numpy()]
labels = obj_inputs.infos['label'].values
im_ids = obj_inputs.infos['batch_im_id'].values
images_ = images[im_ids]
K_ = K[im_ids]
TCO_input = obj_inputs.poses
outputs = model(images=images_,
K=K_,
TCO=TCO_input,
n_iterations=n_iterations,
labels=labels)
timer.pause()
for n in range(1, n_iterations + 1):
iter_outputs = outputs[f'iteration={n}']
infos = obj_inputs.infos
batch_preds = tc.PandasTensorCollection(
infos,
poses=iter_outputs['TCO_output'],
poses_input=iter_outputs['TCO_input'],
K_crop=iter_outputs['K_crop'],
boxes_rend=iter_outputs['boxes_rend'],
boxes_crop=iter_outputs['boxes_crop'])
preds[f'iteration={n}'].append(batch_preds)
logger.debug(
f'Pose prediction on {len(obj_data)} detections (n_iterations={n_iterations}): {timer.stop()}'
)
preds = dict(preds)
for k, v in preds.items():
preds[k] = tc.concatenate(v)
return preds
def get_predictions(self,
obj_predictor,
use_gt_joints=True,
predict_keypoints=True,
n_iterations_init=10,
n_iterations_new_image=1):
predictions = defaultdict(list)
prev_preds = None
pred_key = None
for frame_id, data in enumerate(tqdm(self.dataloader)):
detections = data['detections']
images = data['images'].cuda().float().permute(0, 3, 1, 2) / 255
cameras = data['cameras'].cuda().float()
joints = data['joints']
stacked_joints = dict()
for k, v in joints.items():
stacked_joints[k] = torch.stack(v).cuda().float()
joints = stacked_joints
assert len(obj_predictor.model.robots) == 1
robot = obj_predictor.model.robots[0]
urdf_layer = obj_predictor.model.get_urdf_layer(robot.name)
joints_tensor = urdf_layer.to_tensor(joints)
detections.register_tensor('joints', joints_tensor)
detections.infos['joint_names'] = urdf_layer.joint_names[
None].repeat(len(detections), axis=0).tolist()
K = cameras.K.float()
detections = detections.cuda().float()
assert len(images) == 1
if frame_id == 0:
_, current_preds = obj_predictor.get_predictions(
images=images,
K=K,
detections=detections,
use_gt_joints=use_gt_joints,
gt_joints=joints_tensor,
n_iterations=n_iterations_init,
)
pred_key = f'iteration={n_iterations_init}'
else:
_, current_preds = obj_predictor.get_predictions(
images=images,
K=K,
use_gt_joints=use_gt_joints,
gt_joints=joints_tensor,
n_iterations=n_iterations_new_image,
data_TCO_init=prev_preds[f'{pred_key}'],
)
pred_key = f'iteration={n_iterations_new_image}'
prev_preds = current_preds
current_pred = current_preds[pred_key]
for k in ('view_id', 'scene_id', 'batch_im_id'):
current_pred.infos[k] = detections.infos[k]
preds = {'online': current_pred}
if predict_keypoints:
for pred_k, pred_v in preds.items():
TCO_keypoints_3d, keypoints_2d = obj_predictor.pred_keypoints(
pred_v)
orig_wh = cameras.orig_wh[pred_v.infos.batch_im_id]
cropresize_bboxes = cameras.cropresize_bboxes[
pred_v.infos.batch_im_id]
wh = torch.zeros_like(orig_wh)
wh[:, 1] = images.shape[2]
wh[:, 0] = images.shape[3]
keypoints_2d = cropresize_backtransform_points2d(
orig_wh, cropresize_bboxes, wh, keypoints_2d)
pred_v.register_tensor('keypoints_2d', keypoints_2d)
pred_v.register_tensor('TCO_keypoints_3d',
TCO_keypoints_3d)
predictions[pred_k].append(pred_v)
for k, v in predictions.items():
predictions[k] = tc.concatenate(predictions[k])
return predictions
def batched_model_predictions(self,
images,
K,
obj_data,
n_iterations=1,
update_obj_infos=True):
timer = Timer()
timer.start()
ids = torch.arange(len(obj_data))
ds = TensorDataset(ids)
dl = DataLoader(ds, batch_size=self.bsz_objects)
preds = defaultdict(list)
for (batch_ids, ) in dl:
obj_inputs = obj_data[batch_ids.numpy()]
im_ids = obj_inputs.infos.batch_im_id.values
images_ = images[im_ids]
K_ = K[im_ids]
TCO_input, obj_infos_input = data_to_pose_model_inputs(obj_inputs)
outputs = self.model(images=images_,
K=K_,
TCO=TCO_input,
obj_infos=obj_infos_input,
n_iterations=n_iterations,
update_obj_infos=update_obj_infos,
deterministic=True)
for n in range(1, n_iterations + 1):
iter_outputs = outputs[f'iteration={n}']
bsz = len(images_)
obj_outputs = iter_outputs['obj_infos_output']
obj_inputs_ = iter_outputs['obj_infos_input']
q_input = {
k: torch.cat(
[obj_inputs_[n]['joints'][k] for n in range(bsz)],
dim=0)
for k in obj_inputs_[0]['joints'].keys()
}
q_pred = {
k: torch.cat(
[obj_outputs[n]['joints'][k] for n in range(bsz)],
dim=0)
for k in obj_outputs[0]['joints'].keys()
}
q_pred = self.urdf_layer.to_tensor(q_pred)
infos = obj_inputs.infos
data = tc.PandasTensorCollection(
infos,
poses=iter_outputs['TCO_output'],
K=iter_outputs['K_input'],
joints=q_pred,
K_crop=iter_outputs['K_crop'])
preds[f'iteration={n}'].append(data)
for k, v in preds.items():
preds[k] = tc.concatenate(v)
logger.debug(
f'Pose prediction on {len(obj_data)} detections (n_iterations={n_iterations}) (joint_update={update_obj_infos}): {timer.stop()}'
)
return preds