Fix evaluation distance code.

Author: Phuc Le

data.py

@@ -63,21 +63,20 @@ def collate_fn(self, batch):
             batch: (list) [chunkSize, ]
                 each element of list batch has shape
                 [viewN, height, width, channels]
-        Return: {}
-            inputImage: [batchSize, height, width, channels]
+        Return: Tensor 
+            inputImage: [batchSize, channels, height, width]
             targetTrans: [batchSize, novelN, 4]
-            depthGT: [batchSize, novelN, height, width, 1]
-            maskGT: [batchSize, novelN, height, width, 1]
+            depthGT: [batchSize, novelN, 1, height, width]
+            maskGT: [batchSize, novelN, 1, height, width]
         """
         # Shape: [chunkSize, viewN, height, width, channels] 
         batch_n = {key: np.array([d[key] for d in batch]) for key in batch[0]}
         # Shape: [batchSize,]
         modelIdx = np.random.permutation(self.cfg.chunkSize)[:self.cfg.batchSize]
         # Shape: [batchSize, novelN]
         modelIdxTile = np.tile(modelIdx, [self.cfg.novelN, 1]).T
-        # 24 is the number of rendered images for a single CAD models
         # Shape: [batchSize,]
-        angleIdx = np.random.randint(24, size=[self.cfg.batchSize])
+        angleIdx = np.random.randint(self.cfg.inputViewN, size=[self.cfg.batchSize])
         # Shape: [batchSize, novelN]
         sampleIdx = np.random.randint(
             self.cfg.sampleN, size=[self.cfg.batchSize, self.cfg.novelN])

scripts/evaluate.sh

@@ -3,7 +3,7 @@
 # 	--chunkSize 32 --batchSize 32 \
 # 	--gpu 0
 
-python evaluate.py --model ORIG_STG2 --experiment sgd_trueWD_restart \
-	--loadPath ORIG_STG2_sgd_trueWD_restart \
+python evaluate.py --model ORIG_STG2 --experiment sgd_trueWD_restart_cont \
+	--loadPath ORIG_STG2_sgd_trueWD_restart_cont \
 	--chunkSize 32 --batchSize 32 \
 	--gpu 0

scripts/evaluate_dist.sh

@@ -2,12 +2,12 @@
 # 	--loadPath ORIG_STG2_adam_trueWD_cyclical \
 # 	--chunkSize 32 --batchSize 32 
 
-# python evaluate_dist.py --model ORIG_STG2 --experiment sgd_trueWD_restart \
-# 	--loadPath ORIG_STG2_sgd_trueWD_restart \
-# 	--chunkSize 32 --batchSize 32 \
-# 	--gpu 0
-
-python evaluate_dist.py --model ORIG_STG2 --experiment orig_tf \
-	--loadPath "~/3D-point-cloud-generation/results_0/orig-ft_it100000" \
+python evaluate_dist.py --model ORIG_STG2 --experiment sgd_trueWD_restart_cont \
+	--loadPath ORIG_STG2_sgd_trueWD_restart_cont \
 	--chunkSize 32 --batchSize 32 \
-	--gpu 1
+	--gpu 0
+
+# python evaluate_dist.py --model ORIG_STG2 --experiment orig_tf \
+# 	--loadPath "~/3D-point-cloud-generation/results_0/orig-ft_it100000" \
+# 	--chunkSize 32 --batchSize 32 \
+# 	--gpu 1

scripts/train_stg2.sh

@@ -35,9 +35,9 @@
 # 	--gpu 1
 
 # Continue training
-python train_stg2.py --model ORIG_STG2 --experiment sgd_trueWD_restart_cont \
-	--loadPath ORIG_STG2_sgd_trueWD_restart_cont \
-	--chunkSize 32 --batchSize 32 --saveEpoch 10 \
+python train_stg2.py --model ORIG_STG2 --experiment sgd_trueWD_restart_cont1 \
+	--loadPath ORIG_STG1_sgd_trueWD_restart \
+	--chunkSize 100 --batchSize 32 --saveEpoch 10 \
 	--optim sgd --trueWD 1e-4 --lr 1e-1 \
 	--lrSched restart --T_0 5 --T_mult 2 --lrBase 5e-3 \
 	--gpu 1

trainer.py

@@ -487,17 +487,18 @@ def eval(self, model):
             XYZid, ML = transform.fuse3D(
                 self.cfg, XYZ, maskLogit, fuseTrans) # [B,3,VHW],[B,1,VHW]
 
+            XYZid, ML = XYZid.permute([0, 2, 1]), ML.squeeze()
             for a in range(self.cfg.inputViewN):
-                xyz = XYZid[a].transpose(0,1) #[VHW, 3]
-                ml = ML[a].reshape([-1]) #[VHW]
+                xyz = XYZid[a] #[VHW, 3]
+                ml = ML[a] #[VHW]
                 points24[a, 0] = (xyz[ml > 0]).detach().cpu().numpy()
 
-            pointMeanN = np.array([len(p) for p in points24]).mean()
+            pointMeanN = np.array([len(p) for p in points24[:, 0]]).mean()
             scipy.io.savemat(
                 f"{self.result_path}/{self.CADs[i]}.mat", 
                 {"image": cad["image_in"], "pointcloud": points24})
 
-            print(f"Save pointcloud to {self.result_path}/{self.CADs[i]}.mat")
+            print(f"{pointMeanN:.2f} points save to {self.result_path}/{self.CADs[i]}.mat")
             self.history.append(
                 {"cad": self.CADs[i], "average points": pointMeanN})
 

transform.py

@@ -116,18 +116,23 @@ def render2D(cfg, XYZid, ML, renderTrans):  # [B,1,VHW]
     _, MLnewInvis = torch.unbind(invisFloat, dim=1) # [U-U']
 
     # map to upsampled inverse depth and mask (visible)
-    scatterIdx = torch.stack(
-        [batchIdxVis, novelIdxVis, YnewVis, XnewVis], dim=1)  # [U,4]
+    # scatterIdx = torch.stack(
+    #     [batchIdxVis, novelIdxVis, YnewVis, XnewVis], dim=1)  # [U,4]
     upNewiZMLCnt = torch.zeros([cfg.batchSize, cfg.novelN, 3,
                                  cfg.H*cfg.upscale, cfg.W*cfg.upscale]
                                 ).to(cfg.device) #[B,N,3,uH,uW]
     countOnes = torch.ones_like(iZnewVis)
     scatteriZMLCnt = torch.stack([iZnewVis, MLnewVis, countOnes], dim=1) #[U,3]
-    upNewiZMLCnt[scatterIdx[:,0],
-                 scatterIdx[:,1],
+    # upNewiZMLCnt[scatterIdx[:,0],
+    #              scatterIdx[:,1],
+    #              :,
+    #              scatterIdx[:,2],
+    #              scatterIdx[:,3]] = scatteriZMLCnt
+    upNewiZMLCnt[batchIdxVis,
+                 novelIdxVis,
                  :,
-                 scatterIdx[:,2],
-                 scatterIdx[:,3]] = scatteriZMLCnt
+                 YnewVis,
+                 XnewVis] = scatteriZMLCnt
     upNewiZMLCnt = upNewiZMLCnt.reshape([cfg.batchSize * cfg.novelN,
                                          3,
                                          cfg.H * cfg.upscale,
@@ -145,7 +150,7 @@ def render2D(cfg, XYZid, ML, renderTrans):  # [B,1,VHW]
     upNewML = torch.zeros([cfg.batchSize, cfg.novelN, 1,
                            cfg.H*cfg.upscale, cfg.W*cfg.upscale]
                           ).to(cfg.device) # [B,N,1,uH,uW]
-    scatterML = torch.stack([MLnewInvis], dim=1)  # [U,1]
+    scatterML = MLnewInvis.unsqueeze(-1)  # [U,1]
     upNewML[scatterIdx[:,0],
             scatterIdx[:,1],
             :,