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 data'''
cfg.log_string('Loading dataset.')
test_loader = load_dataloader(cfg, mode='test')
'''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()
'''Load tester'''
cfg.log_string('Loading tester.')
tester = load_tester(cfg=cfg, net=net, device=device)
'''Start to test'''
cfg.log_string('Start to test.')
cfg.log_string('Total number of parameters in {0:s}: {1:d}.'.format(
cfg.config['method'], sum(p.numel() for p in net.parameters())))
test(cfg=cfg, tester=tester, test_loader=test_loader)
cfg.write_config()
cfg.log_string('Testing finished.')
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()
'''Load data'''
cfg.log_string('Loading data.')
input_data = load_demo_data(cfg, device)
'''Run demo'''
net.train(cfg.config['mode'] == 'train')
start = time()
our_data = generate(cfg, net.module, input_data, post_processing=False)
end = time()
print('Time elapsed: %s.' % (end - start))
'''Save visualization'''
scene_name = os.path.splitext(os.path.basename(cfg.config['demo_path']))[0]
output_dir = os.path.join('demo/outputs', scene_name)
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
save_visualization(cfg, input_data, our_data, output_dir)
visualize(output_dir, offline=False)
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 data'''
cfg.log_string('Loading dataset.')
train_loader = load_dataloader(cfg, mode='train')
val_loader = load_dataloader(cfg, mode='val')
'''Load net'''
cfg.log_string('Loading model.')
net = load_model(cfg, device=device)
checkpoint.register_modules(net=net)
cfg.log_string(net)
'''Load optimizer'''
cfg.log_string('Loading optimizer.')
optimizer = load_optimizer(config=cfg.config, net=net)
checkpoint.register_modules(optimizer=optimizer)
'''Load scheduler'''
cfg.log_string('Loading optimizer scheduler.')
scheduler = load_scheduler(config=cfg.config, optimizer=optimizer)
checkpoint.register_modules(scheduler=scheduler)
'''Check existing checkpoint (resume or finetune)'''
checkpoint.parse_checkpoint()
'''BN momentum scheduler'''
bnm_scheduler = load_bnm_scheduler(cfg=cfg, net=net, start_epoch=scheduler.last_epoch)
'''Load trainer'''
cfg.log_string('Loading trainer.')
trainer = load_trainer(cfg=cfg, net=net, optimizer=optimizer, device=device)
'''Start to train'''
cfg.log_string('Start to train.')
cfg.log_string('Total number of parameters in {0:s}: {1:d}.'.format(cfg.config['method'], sum(p.numel() for p in net.parameters())))
train(cfg=cfg, trainer=trainer, scheduler=scheduler, bnm_scheduler=bnm_scheduler, checkpoint=checkpoint, train_loader=train_loader, val_loader=val_loader)
cfg.log_string('Training finished.')
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))