def __call__(self, est_data, gt_data, bins_tensor):
pitch_cls_loss, pitch_reg_loss = cls_reg_loss(est_data['pitch_cls_result'], gt_data['pitch_cls'], est_data['pitch_reg_result'], gt_data['pitch_reg'])
roll_cls_loss, roll_reg_loss = cls_reg_loss(est_data['roll_cls_result'], gt_data['roll_cls'], est_data['roll_reg_result'], gt_data['roll_reg'])
lo_ori_cls_loss, lo_ori_reg_loss = cls_reg_loss(est_data['lo_ori_cls_result'], gt_data['lo_ori_cls'], est_data['lo_ori_reg_result'], gt_data['lo_ori_reg'])
lo_centroid_loss = reg_criterion(est_data['lo_centroid_result'], gt_data['lo_centroid']) * cls_reg_ratio
lo_coeffs_loss = reg_criterion(est_data['lo_coeffs_result'], gt_data['lo_coeffs']) * cls_reg_ratio
lo_bdb3D_result = get_layout_bdb_sunrgbd(bins_tensor, est_data['lo_ori_reg_result'], gt_data['lo_ori_cls'], est_data['lo_centroid_result'],
est_data['lo_coeffs_result'])
# layout bounding box corner loss
lo_corner_loss = cls_reg_ratio * reg_criterion(lo_bdb3D_result, gt_data['lo_bdb3D'])
return {'pitch_cls_loss':pitch_cls_loss, 'pitch_reg_loss':pitch_reg_loss,
'roll_cls_loss':roll_cls_loss, 'roll_reg_loss':roll_reg_loss,
'lo_ori_cls_loss':lo_ori_cls_loss, 'lo_ori_reg_loss':lo_ori_reg_loss,
'lo_centroid_loss':lo_centroid_loss, 'lo_coeffs_loss':lo_coeffs_loss,
'lo_corner_loss':lo_corner_loss}, {'lo_bdb3D_result':lo_bdb3D_result}
def run(cfg):
'''Begin to run network.'''
checkpoint = CheckpointIO(cfg)
'''Mount external config data'''
cfg = mount_external_config(cfg)
'''Load save path'''
cfg.log_string('Data save path: %s' % (cfg.save_path))
'''Load device'''
cfg.log_string('Loading device settings.')
device = load_device(cfg)
'''Load net'''
cfg.log_string('Loading model.')
net = load_model(cfg, device=device)
checkpoint.register_modules(net=net)
cfg.log_string(net)
'''Load existing checkpoint'''
checkpoint.parse_checkpoint()
cfg.log_string('-' * 100)
'''Load data'''
cfg.log_string('Loading data.')
data = load_demo_data(cfg.config['demo_path'], device)
'''Run demo'''
net.train(cfg.config['mode'] == 'train')
with torch.no_grad():
start = time()
est_data = net(data)
end = time()
print('Time elapsed: %s.' % (end-start))
'''write and visualize outputs'''
from net_utils.libs import get_layout_bdb_sunrgbd, get_rotation_matix_result, get_bdb_evaluation
from scipy.io import savemat
from libs.tools import write_obj
lo_bdb3D_out = get_layout_bdb_sunrgbd(cfg.bins_tensor, est_data['lo_ori_reg_result'],
torch.argmax(est_data['lo_ori_cls_result'], 1),
est_data['lo_centroid_result'],
est_data['lo_coeffs_result'])
# camera orientation for evaluation
cam_R_out = get_rotation_matix_result(cfg.bins_tensor,
torch.argmax(est_data['pitch_cls_result'], 1), est_data['pitch_reg_result'],
torch.argmax(est_data['roll_cls_result'], 1), est_data['roll_reg_result'])
# projected center
P_result = torch.stack(((data['bdb2D_pos'][:, 0] + data['bdb2D_pos'][:, 2]) / 2 -
(data['bdb2D_pos'][:, 2] - data['bdb2D_pos'][:, 0]) * est_data['offset_2D_result'][:, 0],
(data['bdb2D_pos'][:, 1] + data['bdb2D_pos'][:, 3]) / 2 -
(data['bdb2D_pos'][:, 3] - data['bdb2D_pos'][:, 1]) * est_data['offset_2D_result'][:,1]), 1)
bdb3D_out_form_cpu, bdb3D_out = get_bdb_evaluation(cfg.bins_tensor,
torch.argmax(est_data['ori_cls_result'], 1),
est_data['ori_reg_result'],
torch.argmax(est_data['centroid_cls_result'], 1),
est_data['centroid_reg_result'],
data['size_cls'], est_data['size_reg_result'], P_result,
data['K'], cam_R_out, data['split'], return_bdb=True)
# save results
nyu40class_ids = [int(evaluate_bdb['classid']) for evaluate_bdb in bdb3D_out_form_cpu]
save_path = cfg.config['demo_path'].replace('inputs', 'outputs')
if not os.path.exists(save_path):
os.makedirs(save_path)
# save layout
savemat(os.path.join(save_path, 'layout.mat'),
mdict={'layout': lo_bdb3D_out[0, :, :].cpu().numpy()})
# save bounding boxes and camera poses
interval = data['split'][0].cpu().tolist()
current_cls = nyu40class_ids[interval[0]:interval[1]]
savemat(os.path.join(save_path, 'bdb_3d.mat'),
mdict={'bdb': bdb3D_out_form_cpu[interval[0]:interval[1]], 'class_id': current_cls})
savemat(os.path.join(save_path, 'r_ex.mat'),
mdict={'cam_R': cam_R_out[0, :, :].cpu().numpy()})
# save meshes
current_faces = est_data['out_faces'][interval[0]:interval[1]].cpu().numpy()
current_coordinates = est_data['meshes'].transpose(1, 2)[interval[0]:interval[1]].cpu().numpy()
for obj_id, obj_cls in enumerate(current_cls):
file_path = os.path.join(save_path, '%s_%s.obj' % (obj_id, obj_cls))
mesh_obj = {'v': current_coordinates[obj_id],
'f': current_faces[obj_id]}
write_obj(file_path, mesh_obj)
#########################################################################
#
# Visualization
#
#########################################################################
import scipy.io as sio
from utils.visualize import format_bbox, format_layout, format_mesh, Box
from glob import glob
pre_layout_data = sio.loadmat(os.path.join(save_path, 'layout.mat'))['layout']
pre_box_data = sio.loadmat(os.path.join(save_path, 'bdb_3d.mat'))
pre_boxes = format_bbox(pre_box_data, 'prediction')
pre_layout = format_layout(pre_layout_data)
pre_cam_R = sio.loadmat(os.path.join(save_path, 'r_ex.mat'))['cam_R']
vtk_objects, pre_boxes = format_mesh(glob(os.path.join(save_path, '*.obj')), pre_boxes)
image = np.array(Image.open(os.path.join(cfg.config['demo_path'], 'img.jpg')).convert('RGB'))
cam_K = np.loadtxt(os.path.join(cfg.config['demo_path'], 'cam_K.txt'))
scene_box = Box(image, None, cam_K, None, pre_cam_R, None, pre_layout, None, pre_boxes, 'prediction', output_mesh = vtk_objects)
scene_box.draw_projected_bdb3d('prediction', if_save=True, save_path = '%s/3dbbox.png' % (save_path))
scene_box.draw3D(if_save=True, save_path = '%s/recon.png' % (save_path))
def get_metric_values(self, est_data, gt_data):
''' Performs a evaluation step.
'''
# Layout IoU
lo_bdb3D_out = get_layout_bdb_sunrgbd(
self.cfg.bins_tensor, est_data['lo_ori_reg_result'],
torch.argmax(est_data['lo_ori_cls_result'], 1),
est_data['lo_centroid_result'], est_data['lo_coeffs_result'])
layout_iou = []
for index, sequence_id in enumerate(gt_data['sequence_id']):
lo_iou = get_iou_cuboid(
lo_bdb3D_out[index, :, :].cpu().numpy(),
gt_data['lo_bdb3D'][index, :, :].cpu().numpy())
layout_iou.append(lo_iou)
# camera orientation for evaluation
cam_R_out = get_rotation_matix_result(
self.cfg.bins_tensor, torch.argmax(est_data['pitch_cls_result'],
1),
est_data['pitch_reg_result'],
torch.argmax(est_data['roll_cls_result'],
1), est_data['roll_reg_result'])
# projected center
P_result = torch.stack(
((gt_data['bdb2D_pos'][:, 0] + gt_data['bdb2D_pos'][:, 2]) / 2 -
(gt_data['bdb2D_pos'][:, 2] - gt_data['bdb2D_pos'][:, 0]) *
est_data['offset_2D_result'][:, 0],
(gt_data['bdb2D_pos'][:, 1] + gt_data['bdb2D_pos'][:, 3]) / 2 -
(gt_data['bdb2D_pos'][:, 3] - gt_data['bdb2D_pos'][:, 1]) *
est_data['offset_2D_result'][:, 1]), 1)
bdb3D_out_form_cpu, bdb3D_out = get_bdb_evaluation(
self.cfg.bins_tensor,
torch.argmax(est_data['ori_cls_result'], 1),
est_data['ori_reg_result'],
torch.argmax(est_data['centroid_cls_result'], 1),
est_data['centroid_reg_result'],
gt_data['size_cls'],
est_data['size_reg_result'],
P_result,
gt_data['K'],
cam_R_out,
gt_data['split'],
return_bdb=True)
bdb2D_out = get_bdb_2d_result(bdb3D_out, cam_R_out, gt_data['K'],
gt_data['split'])
nyu40class_ids = []
IoU3D = []
IoU2D = []
for index, evaluate_bdb in enumerate(bdb3D_out_form_cpu):
NYU40CLASS_ID = int(evaluate_bdb['classid'])
iou_3D = get_iou_cuboid(
get_corners_of_bb3d_no_index(evaluate_bdb['basis'],
evaluate_bdb['coeffs'],
evaluate_bdb['centroid']),
gt_data['bdb3D'][index, :, :].cpu().numpy())
box1 = bdb2D_out[index, :].cpu().numpy()
box2 = gt_data['bdb2D_from_3D_gt'][index, :].cpu().numpy()
box1 = {'u1': box1[0], 'v1': box1[1], 'u2': box1[2], 'v2': box1[3]}
box2 = {'u1': box2[0], 'v1': box2[1], 'u2': box2[2], 'v2': box2[3]}
iou_2D = get_iou(box1, box2)
nyu40class_ids.append(NYU40CLASS_ID)
IoU3D.append(iou_3D)
IoU2D.append(iou_2D)
'''Save results'''
if self.cfg.config['log']['save_results']:
save_path = self.cfg.config['log']['vis_path']
for index, sequence_id in enumerate(gt_data['sequence_id']):
save_path_per_img = os.path.join(save_path,
str(sequence_id.item()))
if not os.path.exists(save_path_per_img):
os.mkdir(save_path_per_img)
# save layout results
savemat(
os.path.join(save_path_per_img, 'layout.mat'),
mdict={'layout': lo_bdb3D_out[index, :, :].cpu().numpy()})
# save bounding boxes and camera poses
interval = gt_data['split'][index].cpu().tolist()
current_cls = nyu40class_ids[interval[0]:interval[1]]
savemat(os.path.join(save_path_per_img, 'bdb_3d.mat'),
mdict={
'bdb': bdb3D_out_form_cpu[interval[0]:interval[1]],
'class_id': current_cls
})
savemat(os.path.join(save_path_per_img, 'r_ex.mat'),
mdict={'cam_R': cam_R_out[index, :, :].cpu().numpy()})
current_faces = est_data['out_faces'][
interval[0]:interval[1]].cpu().numpy()
current_coordinates = est_data['meshes'].transpose(
1, 2)[interval[0]:interval[1]].cpu().numpy()
for obj_id, obj_cls in enumerate(current_cls):
file_path = os.path.join(save_path_per_img,
'%s_%s.obj' % (obj_id, obj_cls))
mesh_obj = {
'v': current_coordinates[obj_id],
'f': current_faces[obj_id]
}
write_obj(file_path, mesh_obj)
metrics = {}
metrics['layout_iou'] = np.mean(layout_iou)
metrics['iou_3d'] = IoU3D
metrics['iou_2d'] = IoU2D
return metrics