pengsida/NeuralBody

1from lib.config import cfg, args2 3 4def run_dataset():5    from lib.datasets import make_data_loader6    import tqdm7 8    cfg.train.num_workers = 09    data_loader = make_data_loader(cfg, is_train=False)10    for batch in tqdm.tqdm(data_loader):11        pass12 13 14def run_network():15    from lib.networks import make_network16    from lib.datasets import make_data_loader17    from lib.utils.net_utils import load_network18    import tqdm19    import torch20    import time21 22    network = make_network(cfg).cuda()23    load_network(network, cfg.trained_model_dir, epoch=cfg.test.epoch)24    network.eval()25 26    data_loader = make_data_loader(cfg, is_train=False)27    total_time = 028    for batch in tqdm.tqdm(data_loader):29        for k in batch:30            if k != 'meta':31                batch[k] = batch[k].cuda()32        with torch.no_grad():33            torch.cuda.synchronize()34            start = time.time()35            network(batch)36            torch.cuda.synchronize()37            total_time += time.time() - start38    print(total_time / len(data_loader))39 40 41def run_evaluate():42    from lib.datasets import make_data_loader43    from lib.evaluators import make_evaluator44    import tqdm45    import torch46    from lib.networks import make_network47    from lib.utils import net_utils48    from lib.networks.renderer import make_renderer49 50    cfg.perturb = 051 52    network = make_network(cfg).cuda()53    net_utils.load_network(network,54                           cfg.trained_model_dir,55                           resume=cfg.resume,56                           epoch=cfg.test.epoch)57    network.train()58 59    data_loader = make_data_loader(cfg, is_train=False)60    renderer = make_renderer(cfg, network)61    evaluator = make_evaluator(cfg)62    for batch in tqdm.tqdm(data_loader):63        for k in batch:64            if k != 'meta':65                batch[k] = batch[k].cuda()66        with torch.no_grad():67            output = renderer.render(batch)68        evaluator.evaluate(output, batch)69    evaluator.summarize()70 71 72def run_visualize():73    from lib.networks import make_network74    from lib.datasets import make_data_loader75    from lib.utils.net_utils import load_network76    from lib.utils import net_utils77    import tqdm78    import torch79    from lib.visualizers import make_visualizer80    from lib.networks.renderer import make_renderer81 82    cfg.perturb = 083 84    network = make_network(cfg).cuda()85    load_network(network,86                 cfg.trained_model_dir,87                 resume=cfg.resume,88                 epoch=cfg.test.epoch)89    network.train()90 91    data_loader = make_data_loader(cfg, is_train=False)92    renderer = make_renderer(cfg, network)93    visualizer = make_visualizer(cfg)94    for batch in tqdm.tqdm(data_loader):95        for k in batch:96            if k != 'meta':97                batch[k] = batch[k].cuda()98        with torch.no_grad():99            output = renderer.render(batch)100            visualizer.visualize(output, batch)101 102 103def run_light_stage():104    from lib.utils.light_stage import ply_to_occupancy105    ply_to_occupancy.ply_to_occupancy()106    # ply_to_occupancy.create_voxel_off()107 108 109def run_evaluate_nv():110    from lib.datasets import make_data_loader111    from lib.evaluators import make_evaluator112    import tqdm113    from lib.utils import net_utils114 115    data_loader = make_data_loader(cfg, is_train=False)116    evaluator = make_evaluator(cfg)117    for batch in tqdm.tqdm(data_loader):118        for k in batch:119            if k != 'meta':120                batch[k] = batch[k].cuda()121        evaluator.evaluate(batch)122    evaluator.summarize()123 124 125if __name__ == '__main__':126    globals()['run_' + args.type]()127