Update convert.py

convert.py

@@ -183,42 +183,40 @@ class Converter(nn.Module):
         
         return color, alpha
 
-    def fit_nerf(self, iters=512, resolution=128):
-
+     def fit_nerf(self, iters=512, resolution=128, batch_size=32):
         self.opt.output_size = resolution
-
+    
         optimizer = torch.optim.Adam([
             {'params': self.encoder_density.parameters(), 'lr': 1e-2},
             {'params': self.encoder.parameters(), 'lr': 1e-2},
             {'params': self.mlp_density.parameters(), 'lr': 1e-3},
             {'params': self.mlp.parameters(), 'lr': 1e-3},
         ])
-
+    
         print(f"[INFO] fitting nerf...")
         pbar = tqdm.trange(iters)
         for i in pbar:
-
-            ver = np.random.randint(-45, 45)
-            hor = np.random.randint(-180, 180)
-            rad = np.random.uniform(1.5, 3.0)
+            poses = []
+            for _ in range(batch_size):
+                ver = np.random.randint(-45, 45)
+                hor = np.random.randint(-180, 180)
+                rad = np.random.uniform(1.5, 3.0)
+                poses.append(orbit_camera(ver, hor, rad))
             
-            pose = orbit_camera(ver, hor, rad)
+            poses = np.stack(poses)
             
-            image_gt, alpha_gt = self.render_gs(pose)
-            image_pred, alpha_pred = self.render_nerf(pose)
-
-            # if i % 200 == 0:
-            #     kiui.vis.plot_image(image_gt, alpha_gt, image_pred, alpha_pred)
+            image_gt, alpha_gt = self.render_gs(poses)
+            image_pred, alpha_pred = self.render_nerf(poses)
             
             loss_mse = F.mse_loss(image_pred, image_gt) + 0.1 * F.mse_loss(alpha_pred, alpha_gt)
-            loss = loss_mse #+ 0.1 * self.encoder_density.tv_loss() #+ 0.0001 * self.encoder_density.density_loss()
-
+            loss = loss_mse
+    
             loss.backward()
             self.encoder_density.grad_total_variation(1e-8)
         
             optimizer.step()
             optimizer.zero_grad()
-
+    
             pbar.set_description(f"MSE = {loss_mse.item():.6f}")
         
         print(f"[INFO] finished fitting nerf!")
@@ -266,6 +264,7 @@ class Converter(nn.Module):
 
         # init mesh from nerf
         grid_size = 256
+        chunk_size = 64
         sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32)
 
         S = 128
@@ -275,14 +274,18 @@ class Converter(nn.Module):
         Y = torch.linspace(-1, 1, grid_size).split(S)
         Z = torch.linspace(-1, 1, grid_size).split(S)
 
-        for xi, xs in enumerate(X):
-            for yi, ys in enumerate(Y):
-                for zi, zs in enumerate(Z):
+        for xi in range(0, grid_size, chunk_size):
+            for yi in range(0, grid_size, chunk_size):
+                for zi in range(0, grid_size, chunk_size):
+                    xs = torch.linspace(-1, 1, chunk_size)
+                    ys = torch.linspace(-1, 1, chunk_size)
+                    zs = torch.linspace(-1, 1, chunk_size)
+    
                     xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij')
-                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3]
+                    pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1)
                     val = self.get_density(pts.to(self.device))
-                    sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z]
-
+                    sigmas[xi:xi+chunk_size, yi:yi+chunk_size, zi:zi+chunk_size] = val.reshape(chunk_size, chunk_size, chunk_size).detach().cpu().numpy()
+        
         print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})')
 
         vertices, triangles = mcubes.marching_cubes(sigmas, density_thresh)