Update README.md

Author: zixunh

.gitignore

@@ -145,6 +145,7 @@ checkpoints/
 dataset/ycb/YCB_Video_Dataset
 
 # evaluation
+eval/*/debug/*
 eval/ycb/YCB_Video_toolbox/
 eval/ycb/eval_results*
 eval/dttd/eval_results*

README.md

@@ -91,9 +91,12 @@ bash eval.sh
 ```
 You can customize your own eval command, for example:
 ```bash
-python3 eval.py --dataset_root ../dataset/dttd_iphone/DTTD_IPhone_Dataset/root --model ../checkpoints/m8p4.pth --output eval_results --visualize 
+python eval.py --dataset_root ./dataset/dttd_iphone/DTTD_IPhone_Dataset/root\
+                --model ./checkpoints/m2p1.pth\
+                --base_latent 256 --embed_dim 512 --fusion_block_num 1 --layer_num_m 2 --layer_num_p 1\
+                --visualize --output eval_results_m8p4_model_filtered_best\
 ```
-
+To load model with filter-enhanced MLP, please add flag `--filter`.
 To visualize the attention map or/and the reduced geometric embeddings' distribution, you can add flag `--debug`.
 
 ### Eval
@@ -102,16 +105,25 @@ This is the [ToolBox](https://github.com/yuxng/YCB_Video_toolbox) that we used f
 ### Train
 To run training of our method, you can use:
 ```bash
-python train.py --dataset dttd_iphone --output_dir ./result/train_result --device 0 --batch_size 1 --lr 1e-6 --min_lr 1e-7 --warm_epoch 1 --pretrain ./checkpoints/m8p4_filter_modelrecon.pth
+python train.py --device 0 \
+    --dataset iphone --dataset_root ./dataset/dttd_iphone/DTTD_IPhone_Dataset/root --dataset_config ./dataset/dttd_iphone/dataset_config \
+    --output_dir ./result/result \
+    --base_latent 256 --embed_dim 512 --fusion_block_num 1 --layer_num_m 8 --layer_num_p 4 \
+    --recon_w 0.3 --recon_choice depth \
+    --loss adds --optim_batch 4 \
+    --start_epoch 0 \
+    --lr 1e-5 --min_lr 1e-6 --lr_rate 0.3 --decay_margin 0.033 --decay_rate 0.82 --nepoch 60 --warm_epoch 1 \
+    --filter_enhance \
 ```
+To train a smaller model, you can set flags `--layer_num_m 2 --layer_num_p 1`.
 To enable our method with depth robustifying modules, you can add flags `--filter_enhance` or/and `--recon_choice model`.
 
-To adjust the weight of Chamfer Distance Loss to 0.5, you can set flags `--reon_weight 0.5`.
+To adjust the weight of Chamfer Distance Loss to 0.5, you can set flags `--reon_w 0.5`.
 
 Our model is applicable on YCBV_Dataset and DTTD_v1 as well, please try following commands to run training of our method with other datasets (please ensure you download the dataset that you want to train on):
 ```bash
-python train.py --dataset ycb --output_dir ./result/train_result --device 0 --batch_size 1 --lr 1e-6 --min_lr 1e-7 --warm_epoch 1
-python train.py --dataset dttd --output_dir ./result/train_result --device 0 --batch_size 1 --lr 1e-6 --min_lr 1e-7 --warm_epoch 1
+python train.py --dataset ycb --output_dir ./result/train_result --device 0 --batch_size 1 --lr 8e-5 --min_lr 8e-6 --warm_epoch 1
+python train.py --dataset dttd --output_dir ./result/train_result --device 0 --batch_size 1 --lr 1e-5 --min_lr 1e-6 --warm_epoch 1
 ```
 
 ### Citation

model/model_utils.py

@@ -83,17 +83,52 @@ def __init__(self, max_seq_length, hidden_size, hidden_dropout_prob):
         self.out_dropout = nn.Dropout(hidden_dropout_prob)
         self.LayerNorm = LayerNorm(hidden_size, eps=1e-12)
 
-    def forward(self, input_tensor):
+    def forward(self, input_tensor, return_filtered=False):
         input_tensor = input_tensor.transpose(2, 1).contiguous()
         batch, seq_len, hidden = input_tensor.shape
         x = torch.fft.rfft(input_tensor, dim=1, norm='ortho')
         weight = torch.view_as_complex(self.complex_weight)
         x = x * weight
         sequence_emb_fft = torch.fft.irfft(x, n=seq_len, dim=1, norm='ortho')
+        if return_filtered: return sequence_emb_fft
         hidden_states = self.out_dropout(sequence_emb_fft)
         hidden_states = self.LayerNorm(hidden_states + input_tensor)
         # hidden_states = sequence_emb_fft + input_tensor
         return hidden_states.transpose(2, 1).contiguous()
+    
+    def visualize_frequency_domain(self, input_tensor):
+        import matplotlib.pyplot as plt
+        import matplotlib
+        font = {'family' : 'Times New Roman',
+                'weight' : 'bold',
+                'size'   : 22}
+
+        matplotlib.rc('font', **font)
+        output_tensor = self.forward(input_tensor, return_filtered=True)
+        print(output_tensor.shape)
+
+        input_tensor = input_tensor.cpu().detach().numpy()[0].T
+        output_tensor = output_tensor.cpu().detach().numpy()[0]
+        sequence_length = input_tensor.shape[0]
+
+        def save(input_tensor, fn, title):   
+            # Perform SVD
+            U, s, V = np.linalg.svd(input_tensor, full_matrices=False)
+
+            projected_tensor = input_tensor@U[0,:]
+
+            plt.figure(figsize=(10, 6))
+            plt.hist(projected_tensor, bins=1000, edgecolor='red')
+
+            plt.xlabel('Geometric Features Reduced to 1-dimensional')
+            plt.ylabel('Probability Density')
+            plt.title(title)
+            plt.grid(True)
+            plt.savefig(fn)
+
+        save(input_tensor, 'before_filtered.png', 'Before')
+        save(output_tensor, 'after_filtered.png', 'After')
+
 
 #  Transformer Customization
 future_mask = torch.triu(torch.zeros([1024, 1024]).fill_(float("-inf")), 1)

model/posefusion.py

@@ -197,7 +197,7 @@ def _make_layer(self, base_latent, embed_dim, n_layer1, n_layer2, require_adl):
 class PoseNet(nn.Module):
     def __init__(self, num_points, num_obj, \
                  base_latent=256, embedding_dim=512, fusion_block_num=1, layer_num_m=2, layer_num_p=4, \
-                 filter_enhance=True, require_adl=True):
+                 recon_choice='depth', filter_enhance=True, require_adl=False):
         super(PoseNet, self).__init__()
         self.num_points = num_points
         self.num_obj = num_obj
@@ -211,8 +211,10 @@ def __init__(self, num_points, num_obj, \
 
         # unimodal embedding
         self.cnn = ModifiedResnet(base_latent)
+        self.recon_choice = recon_choice
         self.ptnet = PointCloudAE(256, num_points, base_latent)
-        self.filter_enhance = None if not filter_enhance else FilterLayer(num_points, base_latent, 0.0)
+        self.modelnet = PointCloudAE(256, num_points, base_latent) if recon_choice=='both' else None
+        self.filter_enhance = FilterLayer(num_points, base_latent, 0.0) if filter_enhance else None
 
         # modality and position interaction
         self.fusion = PoseFusion(base_latent, embedding_dim, \
@@ -236,25 +238,35 @@ def forward(self, img, x, choose, obj, recon_ref=None):
         out_img = self.cnn(img) 
         bs, di, _, _ = out_img.size()
         emb = out_img.view(bs, di, -1)
+        robust_loss = 0
 
         # selection of rgb color embedding
         choose = choose.repeat(1, di, 1) 
         rgb_emb = torch.gather(emb, 2, choose).contiguous()
 
         # depth map / point cloud (embedding)
-        pt_feat, pt_emb, pt_recon, extra_loss = self.ptnet(x, None, recon_ref)
+        if self.recon_choice == 'both':
+            object_geo = recon_ref[1]
+            recon_ref = recon_ref[0]
+        pt_feat, pt_emb, pt_recon, cdl_0 = self.ptnet(x, None, recon_ref)
+        robust_loss+=cdl_0
         pt_emb = self.ptnet.latent(pt_feat, pt_emb)
         if self.filter_enhance is not None:
             pt_emb = self.filter_enhance(pt_emb)
-        
+        if self.recon_choice == 'both':
+            _, obj_emb, _, cdl_1 = self.modelnet(x, None, object_geo)
+            robust_loss+=cdl_1
+            pt_emb += obj_emb
         feat = self.fusion(rgb_emb, pt_emb) 
+
         if self.require_adl:
-            extra_loss += feat[1] 
+            adl = feat[1] 
             feat = feat[0]
+            robust_loss+=adl
 
         out_rx, out_tx, out_cx = self.posepred(feat, obj)
 
-        return out_rx, out_tx, out_cx, rgb_emb.detach(), pt_recon.detach(), extra_loss
+        return out_rx, out_tx, out_cx, rgb_emb.detach(), pt_recon.detach(), robust_loss
 
     def get_attention_map(self, img, x, choose):
         # rgb color embedding
@@ -275,4 +287,9 @@ def get_attention_map(self, img, x, choose):
         _, _, attn1, attn2 = self.fusion.layers[0](rgb_emb, pt_emb, require_attn=True)
 
         return attn1, attn2
- 
+    
+    def get_freq_domain(self, x):
+        pt_feat, pt_emb, _, _ = self.ptnet(x, None, None)
+        pt_emb = self.ptnet.latent(pt_feat, pt_emb)
+        assert self.filter_enhance is not None, "filter enhanced MLP is not applied."
+        freq_domain = self.filter_enhance.visualize_frequency_domain(pt_emb)