upd

eval.py

@@ -28,7 +28,6 @@ def get_opts():
     parser.add_argument('--scene_name', type=str, default='test',
                         help='scene name, used as output folder name')
     parser.add_argument('--split', type=str, default='test',
-                        choices=['test', 'test_train'],
                         help='test or test_train')
     parser.add_argument('--img_wh', nargs="+", type=int, default=[800, 800],
                         help='resolution (img_w, img_h) of the image')

extract_color_mesh.py

@@ -114,7 +114,7 @@ def f(models, embeddings, rays, N_samples, N_importance, chunk, white_back):
     xmin, xmax = args.x_range
     ymin, ymax = args.y_range
     zmin, zmax = args.z_range
-    assert xmax-xmin == ymax-ymin == zmax-zmin, 'the ranges must have the same length!'
+    # assert xmax-xmin == ymax-ymin == zmax-zmin, 'the ranges must have the same length!'
     x = np.linspace(xmin, xmax, N)
     y = np.linspace(ymin, ymax, N)
     z = np.linspace(zmin, zmax, N)
@@ -209,7 +209,7 @@ def f(models, embeddings, rays, N_samples, N_importance, chunk, white_back):
 
         # Step 2. project the vertices onto each training image to infer the color
         print('Fusing colors ...')
-        for idx in tqdm(range(len(dataset.image_paths))):
+        for idx in tqdm(range(len(dataset.image_paths[:1]))):
             ## read image of this pose
             image = cv2.imread(dataset.image_paths[idx])[:,:,::-1]
             image = cv2.resize(image, tuple(args.img_wh))
@@ -265,6 +265,7 @@ def f(models, embeddings, rays, N_samples, N_importance, chunk, white_back):
                         args.chunk,
                         dataset.white_back)
             opacity = results['opacity_coarse'].cpu().numpy()[:, np.newaxis] # (N_vertices, 1)
+            opacity = np.nan_to_num(opacity, 1)
 
             non_occluded = np.ones_like(non_occluded_sum) * 0.1/depth # weight by inverse depth
                                                                     # near=more confident in color

train.py

@@ -134,9 +134,7 @@ def validation_step(self, batch, batch_nb):
             self.logger.experiment.add_images('val/GT_pred_depth',
                                                stack, self.global_step)
 
-        psnr_ = psnr(results[f'rgb_{typ}'], rgbs)
-        log['val_psnr'] = psnr_
-
+        log['val_psnr'] = psnr(results[f'rgb_{typ}'], rgbs)
         return log
 
     def validation_epoch_end(self, outputs):