add new branch for normals

models/convnext.py

@@ -241,12 +241,25 @@ class ConvNeXtReverseBase2(nn.Module):
                     ).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_normals = nn.Sequential(*stage).to(device)
+
         self.last = nn.Conv2d(
                     block_setting[-1].input_channels,
                     1,
                     kernel_size=3,
                     padding='same',
                     ).to(device)
+        self.last_normals = nn.Conv2d(
+                    block_setting[-1].input_channels,
+                    3,
+                    kernel_size=3,
+                    padding='same',
+                    ).to(device)
         # self.last = nn.Sequential(
         #                 Permute([0, 2, 3, 1]),
         #                 nn.Linear(in_features=block_setting[-1].input_channels, out_features=1, bias=True),
@@ -272,10 +285,15 @@ class ConvNeXtReverseBase2(nn.Module):
             x = self.upsamples[i](x)
         
         x = self.stages[-2](x)
-        x = self.upsamples[-1](x)
-        x = self.stages[-1](x)
-        x = self.last(x)
-        return x
+        f = self.upsamples[-1](x)
+
+        x = self.stages[-1](f)
+        depth = self.last(x)
+
+        x = self.stage_normals(f)
+        normals = self.last_normals(x)
+        
+        return depth, normals
 
     def forward(self, x) -> Tensor:
         return self._forward_impl(x)