general save

models/pose_encoder.py

@@ -1,8 +1,8 @@
 import os
 import sys
 
-os.environ["CUDA_VISIBLE_DEVICES"] = '7'
-os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+# os.environ["CUDA_VISIBLE_DEVICES"] = '3'
+# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
 
 import torch
 from torch import nn, Tensor
@@ -77,12 +77,12 @@ class ConvNeXtReverseMaskPose(nn.Module):
         self.layers = nn.Sequential(*layers_list[:-1]).to(device)
         self.stage_pose = nn.Sequential(*layers_list[-1:]).to(device)
 
-        for cnf in block_setting[-1:]:
-            # Bottlenecks
-            stage: List[nn.Module] = []
-            for _ in range(cnf.num_layers):
-                stage.append(block(cnf.input_channels, layer_scale, 0))
-            self.stage_mask = nn.Sequential(*stage).to(device)
+        # for cnf in block_setting[-1:]:
+        #     # Bottlenecks
+        #     stage: List[nn.Module] = []
+        #     for _ in range(cnf.num_layers):
+        #         stage.append(block(cnf.input_channels, layer_scale, 0))
+        #     self.stage_mask = nn.Sequential(*stage).to(device)
 
         self.last = nn.Conv2d(
                     block_setting[-1].input_channels,
@@ -90,12 +90,12 @@ class ConvNeXtReverseMaskPose(nn.Module):
                     kernel_size=3,
                     padding='same',
                     ).to(device)
-        self.last_mask = nn.Conv2d(
-                    block_setting[-1].input_channels,
-                    1,
-                    kernel_size=3,
-                    padding='same',
-                    ).to(device)
+        # self.last_mask = nn.Conv2d(
+        #             block_setting[-1].input_channels,
+        #             1,
+        #             kernel_size=3,
+        #             padding='same',
+        #             ).to(device)
 
 
         for m in self.modules():
@@ -114,12 +114,13 @@ class ConvNeXtReverseMaskPose(nn.Module):
         x = self.first(input)
         f = self.layers(x)
         x = self.stage_pose(f)
-        pose = self.last(x)
+        pose = torch.sigmoid(self.last(x))
 
-        x = self.stage_mask(f)
-        mask = self.last_mask(x)
+        # x = self.stage_mask(f)
+        # mask = torch.sigmoid(self.last_mask(x))
         
-        return pose, mask
+        # return pose, mask
+        return pose, torch.zeros((x.shape[0], 1, x.shape[2], x.shape[3]), device=x.device)
 
     def forward(self, x) -> Tensor:
         return self._forward_impl(x)
@@ -178,11 +179,58 @@ class PoseEncoder(nn.Module):
         rotation = torch.squeeze(rotation, dim=1)
         rotation = torch.softmax(rotation, dim=1)
 
-        x = torch.einsum('bij,bkj->bik', x, rotation)
+        # 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 = self.decoder(x)
+        pose, mask_out = self.decoder(x)
+        # 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, rotation
+        # 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)
+
+class DumbPoseEncoder(nn.Module):
+    def __init__(self, in_channels=7, out_channels=3, device='cuda'):
+        super(DumbPoseEncoder, self).__init__()
+
+        # Encoder
+        self.encoder_conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
+        self.encoder_conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
+        self.encoder_conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+        self.encoder_conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
+        self.encoder_pool = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        # Decoder
+        self.decoder_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
+        self.decoder_upsample_2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv2 = nn.Conv2d(256, 128, kernel_size=3, padding=1)
+        self.decoder_upsample_3 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        self.decoder_conv3 = nn.Conv2d(128, 64, kernel_size=3, padding=1)
+        self.decoder_conv4 = nn.Conv2d(64, out_channels, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        # Encoder
+        mask, pose_in, pose_out = x
+        x = torch.cat([mask, pose_in, pose_out], dim=1)
+        x1 = nn.functional.relu(self.encoder_conv1(x))
+        x = self.encoder_pool(x1)
+        x = nn.functional.relu(self.encoder_conv2(x))
+        x = nn.functional.relu(self.encoder_conv3(x))
+        x = nn.functional.relu(self.encoder_conv4(x))
+
+        # Decoder
+        x = self.decoder_upsample(x)
+        x = nn.functional.relu(self.decoder_conv1(x))
+        x = self.decoder_upsample_2(x)
+        x = nn.functional.relu(self.decoder_conv2(x))
+        x = self.decoder_upsample_3(x)
+        x = nn.functional.relu(self.decoder_conv3(x))
+        x = nn.functional.sigmoid(self.decoder_conv4(x))
+
+        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)
 
 
 if __name__ == '__main__':