add view to input

data/dataset.py

@@ -192,6 +192,8 @@ class Dataset(Dataset):
 
         rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
 
+        view = torch.tensor(np.concatenate((np.array([info['zrot']], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
+
         # image = self.transforms(image)
         data = { 'mesh_posed': mesh_posed,
                  'depth_image': depth_image, 
@@ -201,6 +203,7 @@ class Dataset(Dataset):
                  'depth_mask_garment': depth_mask_garment,
                  'depth_mask_body': depth_mask_body,
                  'rgb_image': rgb_image,
+                 'view': view,
                 }
 
         self.cache[frame] = data

models/depth_estim.py

@@ -8,9 +8,26 @@ from torch.nn import functional as F
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from models.convnext import ConvNeXt, CNBlockConfig
 
+class ViewEncoder(nn.Module):
+    def __init__(self, features=128, device='cuda'):
+        super().__init__()
+        self.features: int = features
+
+        self.lin1 = nn.LazyLinear(features).to(device)
+        self.act1 = nn.GELU().to(device)
+        self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
+    
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        x: torch.Tensor = self.lin1(input)
+        x = self.act1(x)
+        x = self.lin2(x)
+        x = x.view(x.shape[0],self.features,15,20)
+
+        return x
 
 class DepthEstimation(nn.Module):
-    def __init__(self, device='cuda'):
+    def __init__(self, view_features=128, device='cuda'):
         super().__init__()
 
         
@@ -22,7 +39,7 @@ class DepthEstimation(nn.Module):
         ]
         
         block_setting_reverse = [
-            CNBlockConfig(768, 384, 3),
+            CNBlockConfig(768 + view_features, 384, 3),
             CNBlockConfig(384, 192, 9),
             CNBlockConfig(192, 96, 3),
             CNBlockConfig(96, 27, 3),
@@ -34,10 +51,14 @@ class DepthEstimation(nn.Module):
         self.flatten = nn.Flatten()
         self.decoder = ConvNeXt(block_setting_reverse, reversed=True, device=device)
 
+        self.view_encoder = ViewEncoder(view_features)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x, features = self.convnext(x)
+
+    def forward(self, input) -> torch.Tensor:
+        image, view = input
+        x, features = self.convnext(image)
+        view_features = self.view_encoder(view)
         # x = self.flatten(x)
-        x = self.decoder([x, features])
+        x = self.decoder([torch.concat((x, view_features), axis=1),features])
         
         return x

test.py

@@ -23,7 +23,7 @@ def train_step(batch):
     model.train() # train mode is on by default after construction
 
     # 1. Forward pass
-    y_pred = model(batch['rgb_image'])
+    y_pred = model([batch['rgb_image'], batch['view']])
 
     # 2. Calculate loss
     loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
@@ -45,7 +45,7 @@ def valid_step(batch):
     model.eval() # put the model in evaluation mode for testing (inference)
     # 1. Forward pass
     with torch.inference_mode():
-        pred = model(batch['rgb_image'])
+        pred = model([batch['rgb_image'], batch['view']])
     
         # 2. Calculate the loss
         loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])