new reversed convnext

models/depth_estim.py

@@ -6,21 +6,38 @@ import torch.nn as nn
 from torch.nn import functional as F
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from models.convnext import ConvNeXt
+from models.convnext import ConvNeXt, CNBlockConfig
 
 
 class DepthEstimation(nn.Module):
     def __init__(self, device='cuda'):
         super().__init__()
 
-        self.convnext = ConvNeXt()
+        
+        block_setting = [
+            CNBlockConfig(96, 192, 3),
+            CNBlockConfig(192, 384, 3),
+            CNBlockConfig(384, 768, 9),
+            CNBlockConfig(768, None, 3),
+        ]
+        
+        block_setting_reverse = [
+            CNBlockConfig(768, 384, 3),
+            CNBlockConfig(384, 192, 9),
+            CNBlockConfig(192, 96, 3),
+            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(27, 9, 3),
+            CNBlockConfig(9, None, 3),
+        ]
+
+        self.convnext = ConvNeXt(block_setting)
         self.flatten = nn.Flatten()
-        self.decoder = ConvNeXt(reversed=True)
+        self.decoder = ConvNeXt(block_setting_reverse, reversed=True)
 
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.convnext(x)
+        x, features = self.convnext(x)
         # x = self.flatten(x)
-        x = self.decoder(x)
+        x = self.decoder([x, features])
         
         return x