Feature/dpt

Models/dpt/models.py

@@ -49,7 +49,7 @@ def __init__(
         self.pretrained, self.scratch = _make_encoder(
             backbone,
             features,
-            False,  # Set to true of you want to train from scratch, uses ImageNet weights
+            True,  # Set to true of you want to train from scratch, uses ImageNet weights
             groups=1,
             expand=False,
             exportable=False,
@@ -120,34 +120,4 @@ def forward(self, x):
             depth = 1.0 / depth
             return depth
         else:
-            return inv_depth
-
-
-class DPTSegmentationModel(DPT):
-    def __init__(self, num_classes, path=None, **kwargs):
-
-        features = kwargs["features"] if "features" in kwargs else 256
-
-        kwargs["use_bn"] = True
-
-        head = nn.Sequential(
-            nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False),
-            nn.BatchNorm2d(features),
-            nn.ReLU(True),
-            nn.Dropout(0.1, False),
-            nn.Conv2d(features, num_classes, kernel_size=1),
-            Interpolate(scale_factor=2, mode="bilinear", align_corners=True),
-        )
-
-        super().__init__(head, **kwargs)
-
-        self.auxlayer = nn.Sequential(
-            nn.Conv2d(features, features, kernel_size=3, padding=1, bias=False),
-            nn.BatchNorm2d(features),
-            nn.ReLU(True),
-            nn.Dropout(0.1, False),
-            nn.Conv2d(features, num_classes, kernel_size=1),
-        )
-
-        if path is not None:
-            self.load(path)
+            return inv_depth

Trainer.py

@@ -10,10 +10,11 @@
 from utils import secondsToHM
 
 from Models.EncoderModel import EncoderModel
-from Models.DecoderModel import DepthDecoderModel, PoseDecoderModel
+from Models.DecoderModel import PoseDecoderModel
 from Models.BackprojectDepth import BackprojectDepth
 from Models.Project3D import Project3D
 from Models.DisparityAdjustmentV2 import DisparityAdjustment
+from Models.dpt.models import DPTDepthModel
 
 from Losses.SSIM import SSIM
 
@@ -37,17 +38,17 @@ def __init__(self,config):
         self.trainableParameters = []
 
         # Depth Estimation Model Initialization
+        self.models["depth"] = DPTDepthModel(backbone="vitl16_384",non_negative=True,enable_attention_hooks=False,)
+        self.models["depth"] = self.models["depth"].to(self.device)
+        self.trainableParameters += list(self.models["depth"].parameters())
+        self.totalTrainableParams += sum(p.numel() for p in self.models["depth"].parameters() if p.requires_grad)
+
+        # Pose Estimation Model Initialization
         self.models["encoder"] = EncoderModel(50)
         self.models["encoder"] = self.models["encoder"].to(self.device)
         self.trainableParameters += list(self.models["encoder"].parameters())
         self.totalTrainableParams += sum(p.numel() for p in self.models["encoder"].parameters() if p.requires_grad)
-
-        self.models["decoder"] = DepthDecoderModel(self.models["encoder"].numChannels)
-        self.models["decoder"] = self.models["decoder"].to(self.device)
-        self.trainableParameters += list(self.models["decoder"].parameters())
-        self.totalTrainableParams += sum(p.numel() for p in self.models["decoder"].parameters() if p.requires_grad)
-
-        # Pose Estimation Model Initialization
+
         self.models["pose"] = PoseDecoderModel(self.models["encoder"].numChannels, 2, 1)
         self.models["pose"] = self.models["pose"].to(self.device)
         self.trainableParameters += list(self.models["pose"].parameters())
@@ -337,7 +338,7 @@ def processBatch(self, inputs):
         features = {}
         for i, frameIdx in enumerate(self.frameIdxs):
             features[frameIdx] = [f[i] for f in allFrameFeatures]
-        outputs = self.models["decoder"](features[0])
+        outputs = self.models["depth"](features[0])
         outputs.update(self.predictPoses(inputs, features))
         self.generateImagePredictions(inputs, outputs)
         losses = self.computeLosses(inputs, outputs)