apply alphamask reconstruction

TensoRF/configs/chair.txt

@@ -3,13 +3,15 @@ datadir = ../../nerf_synthetic/chair
 expname =  tensorf_lego_VM
 basedir = ./log
 
+alpha_offset = 1e-4
+
 n_iters = 30000
 batch_size = 4096
 
 N_voxel_init = 2097156 # 128**3
 N_voxel_final = 27000000 # 300**3
 upsamp_list = [2000, 3000, 4000, 5500, 7000]
-update_AlphaMask_list = [2000, 4000]
+update_AlphaMask_list = [2000, 4000, 6000, 11000, 16000, 21000, 26000] 
 
 N_vis = 5
 vis_every = 10000

TensoRF/configs/lego.txt

@@ -4,13 +4,15 @@ datadir = ./data/nerf_synthetic/lego
 expname =  tensorf_lego_VM
 basedir = ./log
 
+alpha_offset = 1e-4
+
 n_iters = 30000
 batch_size = 4096
 
 N_voxel_init = 2097156 # 128**3
 N_voxel_final = 27000000 # 300**3
 upsamp_list = [2000,3000,4000,5500,7000]
-update_AlphaMask_list = [2000,4000]
+update_AlphaMask_list = [2000, 4000, 6000, 11000, 16000, 21000, 26000] 
 
 N_vis = 5
 vis_every = 10000

TensoRF/models/tensorBase.py

@@ -154,7 +154,7 @@ def __init__(self, aabb, gridSize, device, density_n_comp=8,
                  alphaMask_thres=0.001, distance_scale=25,
                  rayMarch_weight_thres=0.0001, pos_pe=6, view_pe=6, fea_pe=6,
                  featureC=128, step_ratio=2.0, fea2denseAct='softplus',
-                 grid_bit=32):
+                 grid_bit=32, alpha_offset = 0):
         super(TensorBase, self).__init__()
 
         self.density_n_comp = density_n_comp
@@ -165,7 +165,8 @@ def __init__(self, aabb, gridSize, device, density_n_comp=8,
         self.device=device
 
         self.density_shift = density_shift
-        self.alphaMask_thres = alphaMask_thres
+        self.alphaMask_thres = alphaMask_thres - alpha_offset
+        self.alpha_offset = alpha_offset
         self.distance_scale = distance_scale
         self.rayMarch_weight_thres = rayMarch_weight_thres
         self.fea2denseAct = fea2denseAct
@@ -345,6 +346,7 @@ def updateAlphaMask(self, gridSize=(200,200,200)):
         alpha = alpha.clamp(0,1).transpose(0,2).contiguous()[None,None]
         total_voxels = gridSize[0] * gridSize[1] * gridSize[2]
 
+        self.alphaMask_thres += self.alpha_offset
         ks = 3
         alpha = F.max_pool3d(alpha, kernel_size=ks, padding=ks // 2, stride=1).view(gridSize[::-1])
         alpha[alpha>=self.alphaMask_thres] = 1

TensoRF/opt.py

@@ -78,6 +78,10 @@ def config_parser(cmd=None):
     parser.add_argument("--use_dwt", action='store_true')
     parser.add_argument("--dwt_level", type=int, default=2)
 
+    # Alpha mask
+    parser.add_argument("--alpha_offset", type=float, default=0.0,
+                        help='add to alphamask threshold')
+
     # encoding option
     parser.add_argument("--compress", type=int, default=0)
     parser.add_argument("--decompress", type=int, default=0)

TensoRF/train.py

@@ -146,7 +146,7 @@ def reconstruction(args):
             fea2denseAct=args.fea2denseAct,
             grid_bit=args.grid_bit,
             use_mask=args.use_mask,
-            use_dwt=args.use_dwt, dwt_level=args.dwt_level)
+            use_dwt=args.use_dwt, dwt_level=args.dwt_level, alpha_offset=args.alpha_offset)
 
     # print(tensorf)
     print(f'{sum([p.numel() for p in tensorf.parameters()])*32/8_388_608}MB')
@@ -261,6 +261,13 @@ def reconstruction(args):
                 # update volume resolution
                 reso_mask = reso_cur
 
+            if args.alpha_offset > 0:
+                if iteration != update_AlphaMask_list[0]:
+                    tensorf.alphaMask = None
+
+                if iteration == update_AlphaMask_list[3]:
+                    tensorf.alpha_offset = 0
+
             new_aabb = tensorf.updateAlphaMask(tuple(reso_mask))
 
             if iteration == update_AlphaMask_list[0]:
@@ -322,6 +329,13 @@ def reconstruction(args):
                 f'(G ({args.grid_bit}bit): {grid_bytes/1_048_576:.3f}MB) '
                 f'(N: {non_grid_bytes/1_048_576:3f}MB)')
 
+        # Alpha mask reconstruction
+        _, _, Z, Y, X = tensorf.alphaMask.alpha_volume.shape
+        tensorf.alphaMask = None
+        tensorf.alpha_offset = 0
+        tensorf.updateAlphaMask((X,Y,Z)) 
+
+
     if args.render_train:
         os.makedirs(f'{logfolder}/imgs_train_all', exist_ok=True)
         train_dataset = dataset(args.datadir, split='train',