pose pipeline trial

models/pose_encoder.py

@@ -168,7 +168,7 @@ class PoseEncoder(nn.Module):
             CNBlockConfig(16, None, 3),
         ]
 
-        self.encoder = ConvNeXt(block_setting, input_features=4, device=device)
+        self.encoder = ConvNeXt(block_setting, input_features=7, device=device)
         self.decoder = ConvNeXtReverseMaskPose(block_setting_reverse, device=device)
         self.rotation_encoder = ConvNeXt(rotation_block_setting, input_features=7, device=device)
         self.rotation_decoder = ConvNeXt(rotation_block_setting_reverse, reversed=True, device=device)
@@ -177,8 +177,8 @@ class PoseEncoder(nn.Module):
 
 
     def forward(self, input) -> torch.Tensor:
-        mask, pose_in, pose_out = input
-        features, _ = self.encoder(torch.cat([mask, pose_in], dim=1))
+        mask, pose_in, pose_in_back, pose_out, mask_out = input
+        features, _ = self.encoder(torch.cat([mask, pose_in, pose_in_back], dim=1))
         x = features.view(features.size(0), features.size(1), -1)
         y, _ = self.rotation_encoder(torch.cat([mask, pose_in, pose_out], dim=1))
         rotation = self.rotation_decoder(y)
@@ -188,14 +188,15 @@ class PoseEncoder(nn.Module):
         # x = torch.einsum('bij,bkj->bik', x, rotation)
         x = torch.einsum('bij,bjk->bik', x, rotation)
         x = x.view(x.size(0), x.size(1), 16, 16)
-        pose, mask_out = self.decoder(x)
+        pose, mask_out_pred = self.decoder(x)
+        pose = pose * mask_out
         # mask_out = torch.round(mask_out)
         # mask_out = (mask_out > 0.5).to(torch.int64)
         # pose = pose * (mask_out > 0.5).to(torch.float32)
 
-        # return pose, mask_out, rotation
-        x = pose
-        return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
+        return pose, mask_out_pred, rotation
+        # x = pose
+        # return x, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device), torch.zeros((x.shape[0], 256, 256), device=x.device)
 
 class DumbPoseEncoder(nn.Module):
     def __init__(self, in_channels=7, out_channels=3, device='cuda'):