(args, train_dl, val_dl, hwf, i_split, near, far, render_poses=None, render_img=None)
| 80 | return loss, psnr |
| 81 | |
| 82 | def train_nerf(args, train_dl, val_dl, hwf, i_split, near, far, render_poses=None, render_img=None): |
| 83 | |
| 84 | i_train, i_val, i_test = i_split |
| 85 | # Cast intrinsics to right types |
| 86 | H, W, focal = hwf |
| 87 | H, W = int(H), int(W) |
| 88 | hwf = [H, W, focal] |
| 89 | |
| 90 | # Create log dir and copy the config file |
| 91 | basedir = args.basedir |
| 92 | expname = args.expname |
| 93 | os.makedirs(os.path.join(basedir, expname), exist_ok=True) |
| 94 | f = os.path.join(basedir, expname, 'args.txt') |
| 95 | with open(f, 'w') as file: |
| 96 | for arg in sorted(vars(args)): |
| 97 | attr = getattr(args, arg) |
| 98 | file.write('{} = {}\n'.format(arg, attr)) |
| 99 | if args.config is not None: |
| 100 | f = os.path.join(basedir, expname, 'config.txt') |
| 101 | with open(f, 'w') as file: |
| 102 | file.write(open(args.config, 'r').read()) |
| 103 | |
| 104 | # Create nerf model |
| 105 | render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer = create_nerf(args) |
| 106 | global_step = start |
| 107 | |
| 108 | bds_dict = { |
| 109 | 'near' : near, |
| 110 | 'far' : far, |
| 111 | } |
| 112 | render_kwargs_train.update(bds_dict) |
| 113 | render_kwargs_test.update(bds_dict) |
| 114 | if args.reduce_embedding==2: |
| 115 | render_kwargs_train['i_epoch'] = -1 |
| 116 | render_kwargs_test['i_epoch'] = -1 |
| 117 | |
| 118 | if args.render_test: |
| 119 | print('TRAIN views are', i_train) |
| 120 | print('TEST views are', i_test) |
| 121 | print('VAL views are', i_val) |
| 122 | if args.reduce_embedding==2: |
| 123 | render_kwargs_test['i_epoch'] = global_step |
| 124 | render_test(args, train_dl, val_dl, hwf, start, render_kwargs_test) |
| 125 | return |
| 126 | |
| 127 | # Prepare raybatch tensor if batching random rays |
| 128 | N_rand = args.N_rand |
| 129 | # use_batching = not args.no_batching |
| 130 | |
| 131 | N_epoch = args.epochs + 1 # epoch |
| 132 | print('Begin') |
| 133 | print('TRAIN views are', i_train) |
| 134 | print('TEST views are', i_test) |
| 135 | print('VAL views are', i_val) |
| 136 | |
| 137 | |
| 138 | # loss function |
| 139 | loss_func = loss_dict['nerfw'](coef=1) |
no test coverage detected