add view to input

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