pengsida/NeuralBody

1import torch.nn as nn2import torch3from lib.config import cfg4from .embedder import get_embedder5import torch.nn.functional as F6 7 8class Nerf(nn.Module):9    def __init__(self,10                 D=8,11                 W=256,12                 input_ch=3,13                 input_ch_views=3,14                 skips=[4],15                 use_viewdirs=False):16        """17        """18        super(Nerf, self).__init__()19 20        self.D = D21        self.W = W22        self.input_ch = input_ch23        self.input_ch_views = input_ch_views24        self.skips = skips25        self.use_viewdirs = use_viewdirs26 27        self.pts_linears = nn.ModuleList([nn.Linear(input_ch, W)] + [28            nn.Linear(W, W) if i not in29            self.skips else nn.Linear(W + input_ch, W) for i in range(D - 1)30        ])31 32        ### Implementation according to the official code release (https://github.com/bmild/nerf/blob/master/run_nerf_helpers.py#L104-L105)33        self.views_linears = nn.ModuleList(34            [nn.Linear(input_ch_views + W, W // 2)])35 36        ### Implementation according to the paper37        # self.views_linears = nn.ModuleList(38        #     [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)])39 40        if self.use_viewdirs:41            self.feature_linear = nn.Linear(W, W)42            self.alpha_linear = nn.Linear(W, 1)43            self.rgb_linear = nn.Linear(W // 2, 3)44 45    def forward(self, x):46        input_pts, input_views = torch.split(47            x, [self.input_ch, self.input_ch_views], dim=-1)48        h = input_pts49        for i, l in enumerate(self.pts_linears):50            h = self.pts_linears[i](h)51            h = F.relu(h)52            if i in self.skips:53                h = torch.cat([input_pts, h], -1)54 55        if self.use_viewdirs:56            alpha = self.alpha_linear(h)57            feature = self.feature_linear(h)58            h = torch.cat([feature, input_views], -1)59 60            for i, l in enumerate(self.views_linears):61                h = self.views_linears[i](h)62                h = F.relu(h)63 64            rgb = self.rgb_linear(h)65            outputs = torch.cat([rgb, alpha], -1)66        else:67            outputs = self.output_linear(h)68 69        return outputs70 71    def load_weights_from_keras(self, weights):72        assert self.use_viewdirs, "Not implemented if use_viewdirs=False"73 74        # Load pts_linears75        for i in range(self.D):76            idx_pts_linears = 2 * i77            self.pts_linears[i].weight.data = torch.from_numpy(78                np.transpose(weights[idx_pts_linears]))79            self.pts_linears[i].bias.data = torch.from_numpy(80                np.transpose(weights[idx_pts_linears + 1]))81 82        # Load feature_linear83        idx_feature_linear = 2 * self.D84        self.feature_linear.weight.data = torch.from_numpy(85            np.transpose(weights[idx_feature_linear]))86        self.feature_linear.bias.data = torch.from_numpy(87            np.transpose(weights[idx_feature_linear + 1]))88 89        # Load views_linears90        idx_views_linears = 2 * self.D + 291        self.views_linears[0].weight.data = torch.from_numpy(92            np.transpose(weights[idx_views_linears]))93        self.views_linears[0].bias.data = torch.from_numpy(94            np.transpose(weights[idx_views_linears + 1]))95 96        # Load rgb_linear97        idx_rbg_linear = 2 * self.D + 498        self.rgb_linear.weight.data = torch.from_numpy(99            np.transpose(weights[idx_rbg_linear]))100        self.rgb_linear.bias.data = torch.from_numpy(101            np.transpose(weights[idx_rbg_linear + 1]))102 103        # Load alpha_linear104        idx_alpha_linear = 2 * self.D + 6105        self.alpha_linear.weight.data = torch.from_numpy(106            np.transpose(weights[idx_alpha_linear]))107        self.alpha_linear.bias.data = torch.from_numpy(108            np.transpose(weights[idx_alpha_linear + 1]))109 110 111class Network(nn.Module):112    def __init__(self):113        super(Network, self).__init__()114 115        self.embed_fn, input_ch = get_embedder(cfg.xyz_res)116        self.embeddirs_fn, input_ch_views = get_embedder(cfg.view_res)117 118        skips = [4]119        self.model = Nerf(D=cfg.netdepth,120                          W=cfg.netwidth,121                          input_ch=input_ch,122                          skips=skips,123                          input_ch_views=input_ch_views,124                          use_viewdirs=cfg.use_viewdirs)125 126        self.model_fine = Nerf(D=cfg.netdepth_fine,127                               W=cfg.netwidth_fine,128                               input_ch=input_ch,129                               skips=skips,130                               input_ch_views=input_ch_views,131                               use_viewdirs=cfg.use_viewdirs)132 133    def batchify(self, fn, chunk):134        """Constructs a version of 'fn' that applies to smaller batches.135        """136        def ret(inputs):137            return torch.cat([fn(inputs[i:i+chunk]) for i in range(0, inputs.shape[0], chunk)], 0)138        return ret139 140    def forward(self, inputs, viewdirs, model=''):141        """Prepares inputs and applies network 'fn'.142        """143        if model == 'fine':144            fn = self.model_fine145        else:146            fn = self.model147 148        inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])149        embedded = self.embed_fn(inputs_flat)150 151        input_dirs = viewdirs[:,None].expand(inputs.shape)152        input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]])153        embedded_dirs = self.embeddirs_fn(input_dirs_flat)154        embedded = torch.cat([embedded, embedded_dirs], -1)155 156        outputs_flat = self.batchify(fn, cfg.netchunk)(embedded)157        outputs = torch.reshape(outputs_flat, list(inputs.shape[:-1]) + [outputs_flat.shape[-1]])158        return outputs159