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))