fix examples

.gitattributes

@@ -36,3 +36,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 assets/test1.png filter=lfs diff=lfs merge=lfs -text
 assets/test2.png filter=lfs diff=lfs merge=lfs -text
 assets/test6.jpg filter=lfs diff=lfs merge=lfs -text
+assets/test4.jpg filter=lfs diff=lfs merge=lfs -text
+assets/test7.jpg filter=lfs diff=lfs merge=lfs -text
+assets/test8.jpeg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -238,9 +238,11 @@ with gr.Blocks(theme=theme, title="HORT: Monocular Hand-held Objects Reconstruct
                 ['/home/user/app/assets/test1.png'], 
                 ['/home/user/app/assets/test2.png'], 
                 ['/home/user/app/assets/test3.jpg'], 
-                ['/home/user/app/assets/test4.jpeg'],
+                ['/home/user/app/assets/test4.jpg'],
                 ['/home/user/app/assets/test5.jpeg'],
-                ['/home/user/app/assets/test6.jpg']
+                ['/home/user/app/assets/test6.jpg'],
+                ['/home/user/app/assets/test7.jpg'],
+                ['/home/user/app/assets/test8.jpeg']
                 ], 
                 inputs=input_image)
         

hort/models/tgs/models/snowflake/SPD_pp.py

@@ -33,7 +33,7 @@ class SPD_pp(nn.Module):
 
         self.mlp_delta = MLP_CONV(in_channel=128, layer_dims=[64, 3])
 
-    def forward(self, pcd_prev, feat_cond=None, K_prev=None):
+    def forward(self, pcd_prev, hand_pcd, hand_feat_cond, feat_cond=None, K_prev=None):
         """
         Args:
             pcd_prev: Tensor, (B, 3, N_prev)
@@ -45,15 +45,16 @@ class SPD_pp(nn.Module):
             K_curr: Tensor, displacement feature of current step, (B, 128, N_prev * up_factor)
         """
         b, _, n_prev = pcd_prev.shape
-        feat_1 = self.mlp_1(pcd_prev)
+        feat_1 = self.mlp_1(torch.cat([pcd_prev, hand_pcd], -1))
         feat_1 = torch.cat([feat_1,
                             torch.max(feat_1, 2, keepdim=True)[
                                 0].repeat((1, 1, feat_1.size(2))),
-                            feat_cond], 1) if self.global_feat else feat_1
+                            torch.cat([feat_cond, hand_feat_cond], -1)], 1) if self.global_feat else feat_1
         Q = self.mlp_2(feat_1)
 
         H = self.skip_transformer(
-                pcd_prev, K_prev if K_prev is not None else Q, Q)
+            torch.cat([pcd_prev, hand_pcd], -1), torch.cat([K_prev, hand_feat_cond], -1) if K_prev is not None else Q, Q)
+        H = H[:, :, :-778]
 
         feat_child = self.mlp_ps(H)
         feat_child = self.ps(feat_child)  # (B, 128, N_prev * up_factor)

hort/models/tgs/models/snowflake/model_spdpp.py

@@ -162,7 +162,7 @@ class Decoder(nn.Module):
         up_token =  F.interpolate(pcl_token, scale_factor=up_factor, mode='nearest')
         return up_token
 
-    def calculate_image_token(self, pcd, input_image_tokens, batch):
+    def calculate_image_token(self, pcd, hand_pcd, input_image_tokens, batch):
         """
         Args:
         """
@@ -174,7 +174,10 @@ class Decoder(nn.Module):
         local_features_proj = points_projection_v2(pcd * batch['scale'] + batch['trans'].unsqueeze(1), batch['intrinsic_cond'], local_features)
         local_features_proj = local_features_proj.permute(0, 2, 1).contiguous()
 
-        return local_features_proj
+        local_features_proj_hand = points_projection_v2(hand_pcd, batch['intrinsic_cond'], local_features)
+        local_features_proj_hand = local_features_proj_hand.permute(0, 2, 1).contiguous()
+ 
+        return local_features_proj, local_features_proj_hand
 
     def forward(self, x):
         """
@@ -184,6 +187,7 @@ class Decoder(nn.Module):
             # partial_coarse: Tensor, (b, n_coarse, 3)
         """
         points = x['points']
+        hand_points = x['hand_points']
         if self.token_type == 'pcl_token':
             feat_cond = x['pcl_token']
         elif self.token_type == 'image_token':
@@ -194,14 +198,15 @@ class Decoder(nn.Module):
 
         pcd = torch.permute(points, (0, 2, 1)).contiguous()
         pcl_up_scale = 1
+        hand_pcd = torch.permute((hand_points - x['trans'].unsqueeze(1)) / x['scale'], (0, 2, 1)).contiguous()
         for upper in self.uppers:
             if self.token_type == 'pcl_token':
                 up_cond = self.calculate_pcl_token(
                     feat_cond, pcl_up_scale)
                 pcl_up_scale *= upper.up_factor
             elif self.token_type == 'image_token':
-                up_cond = self.calculate_image_token(points, feat_cond, x)
-            pcd, feat_prev = upper(pcd, up_cond, feat_prev)
+                up_cond, up_cond_hand = self.calculate_image_token(points, hand_points, feat_cond, x)
+            pcd, feat_prev = upper(pcd, hand_pcd, up_cond_hand, up_cond, feat_prev)
             points = torch.permute(pcd, (0, 2, 1)).contiguous()
             arr_pcd.append(points)
         return arr_pcd