We need a new query function, Coarse = NeRF, Fine = NeRF-W Inputs: inputs: torch.Tensor() [N_rays,N_samples,3] viewdirs: torch.Tensor() [N_rays, 3] ts: latent code from img_idxs [N_rays] fn: NeRFW object embed_fn: embedder for position embeddirs_
(inputs, viewdirs, ts, fn, embed_fn, embeddirs_fn,
typ, embedding_a, embedding_t, output_transient,
netchunk=1024*64, test_time=False)
| 13 | to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8) |
| 14 | |
| 15 | def run_network_NeRFW(inputs, viewdirs, ts, fn, embed_fn, embeddirs_fn, |
| 16 | typ, embedding_a, embedding_t, output_transient, |
| 17 | netchunk=1024*64, test_time=False): |
| 18 | ''' We need a new query function, Coarse = NeRF, Fine = NeRF-W |
| 19 | Inputs: |
| 20 | inputs: torch.Tensor() [N_rays,N_samples,3] |
| 21 | viewdirs: torch.Tensor() [N_rays, 3] |
| 22 | ts: latent code from img_idxs [N_rays] |
| 23 | fn: NeRFW object |
| 24 | embed_fn: embedder for position |
| 25 | embeddirs_fn: embedder for view directions |
| 26 | typ: 'coarse' or 'fine' |
| 27 | embedding_a: NeRFW appearance embedding layer |
| 28 | embedding_t: NeRFW transient embedding layer |
| 29 | output_transient: True/False |
| 30 | netchunk: chunk size to inference |
| 31 | test_time: True/False |
| 32 | ''' |
| 33 | out_chunks = [] |
| 34 | N_rays, N_samples = inputs.shape[0], inputs.shape[1] |
| 35 | |
| 36 | # embed inputs like NeRF |
| 37 | if typ == 'coarse' and test_time: |
| 38 | inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]]) |
| 39 | |
| 40 | # Feed NeRF coarse train |
| 41 | for i in range(0, inputs_flat.shape[0], netchunk): |
| 42 | embedded_inputs = [embed_fn(inputs_flat[i: i+netchunk])] |
| 43 | out_chunks += [fn(torch.cat(embedded_inputs, 1), sigma_only=True)] |
| 44 | out = torch.cat(out_chunks, 0) # [N_rays*N_samples, 4] |
| 45 | out = torch.reshape(out, list(inputs.shape[:-1]) + [out.shape[-1]]) # [N_rays, N_samples, 4] |
| 46 | return out |
| 47 | if typ == 'coarse': # case: coarse + train |
| 48 | inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]]) |
| 49 | |
| 50 | input_dirs = viewdirs[:,None].expand(inputs.shape) |
| 51 | input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]]) |
| 52 | |
| 53 | # Feed NeRF coarse train |
| 54 | for i in range(0, inputs_flat.shape[0], netchunk): |
| 55 | embedded_inputs = [embed_fn(inputs_flat[i: i+netchunk]), embeddirs_fn(input_dirs_flat[i:i+netchunk])] |
| 56 | out_chunks += [fn(torch.cat(embedded_inputs, 1), output_transient=output_transient)] |
| 57 | |
| 58 | out = torch.cat(out_chunks, 0) # [N_rays*N_samples, 4] |
| 59 | out = torch.reshape(out, list(inputs.shape[:-1]) + [out.shape[-1]]) # [N_rays, N_samples, 4] |
| 60 | return out |
| 61 | |
| 62 | elif typ == 'fine': |
| 63 | inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]]) |
| 64 | |
| 65 | input_dirs = viewdirs[:,None].expand(inputs.shape) |
| 66 | input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]]) |
| 67 | |
| 68 | # embed encode_a and encode_t |
| 69 | a_embedded = embedding_a(ts.long()) |
| 70 | |
| 71 | # if output_transient: |
| 72 | t_embedded = embedding_t(ts.long()) |