monodepth implemented

nerfstudio/configs/method_configs.py

@@ -96,6 +96,8 @@ method_configs["nerfacto"] = TrainerConfig(
     mixed_precision=True,
     pipeline=VanillaPipelineConfig(
         datamanager=VanillaDataManagerConfig(
+            _target=VanillaDataManager[DepthDataset],
+            patch_size=2,
             dataparser=NerfstudioDataParserConfig(),
             train_num_rays_per_batch=4096,
             eval_num_rays_per_batch=4096,

nerfstudio/data/datasets/depth_dataset.py

@@ -16,11 +16,21 @@
 Depth dataset.
 """
 
+import os
 from typing import Dict
 
+import numpy as np
+
 from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs
 from nerfstudio.data.datasets.base_dataset import InputDataset
 from nerfstudio.data.utils.data_utils import get_depth_image_from_path
+from PIL import Image
+import torch
+from rich.progress import Console, track
+from pathlib import Path
+import json
+
+
 
 
 class DepthDataset(InputDataset):
@@ -31,14 +41,45 @@ class DepthDataset(InputDataset):
         scale_factor: The scaling factor for the dataparser outputs.
     """
 
-    exclude_batch_keys_from_device = InputDataset.exclude_batch_keys_from_device + ["depth_image"]
-
     def __init__(self, dataparser_outputs: DataparserOutputs, scale_factor: float = 1.0):
         super().__init__(dataparser_outputs, scale_factor)
-        assert (
-            "depth_filenames" in dataparser_outputs.metadata.keys()
-            and dataparser_outputs.metadata["depth_filenames"] is not None
-        )
+        #if there are no depth images than we want to generate them all with zoe depth
+        if len(dataparser_outputs.image_filenames) > 0 and ("depth_filenames" not in dataparser_outputs.metadata.keys() or dataparser_outputs.metadata["depth_filenames"] is None):
+            depth_paths = []
+            tranforms = self._find_transform(dataparser_outputs.image_filenames[0])
+            data = dataparser_outputs.image_filenames[0].parent
+            meta = json.load(open(tranforms, "r"))
+            frames = meta['frames']
+            filenames = [data / frames[j]['file_path'].split('/')[-1] for j in range(len(frames))]
+            os.makedirs(dataparser_outputs.image_filenames[0].parent / "depth", exist_ok=True)
+            repo = "isl-org/ZoeDepth"
+            self.zoe = torch.compile(torch.hub.load(repo, "ZoeD_NK", pretrained=True).cuda())
+
+            for i in track(range(len(filenames)), description="Generating depth images"):
+                image_filename = filenames[i]
+                pil_image = Image.open(image_filename)
+                image = np.array(pil_image, dtype="uint8")  # shape is (h, w) or (h, w, 3 or 4)
+                if len(image.shape) == 2:
+                    image = image[:, :, None].repeat(3, axis=2)
+                image = torch.from_numpy(image.astype("float32") / 255.0)
+
+                #BAD: FIX BY FINDING DEVICE
+                with torch.no_grad():
+                    image = torch.permute(image, (2, 0, 1)).unsqueeze(0).cuda()
+                    depth_numpy = self.zoe.infer(image).squeeze().unsqueeze(-1).cpu().numpy()
+                    
+                depth_paths.append(image_filename.parent / "depth" / f"depth{i}.npy")
+                meta['frames'][i]['depth_file_path'] = str(Path(meta['frames'][i]['file_path']).parent / "depth" / f"depth{i}.npy")
+                np.save(depth_paths[-1], depth_numpy)
+
+            json.dump(meta, open(tranforms, "w"))
+
+
+            dataparser_outputs.metadata["depth_filenames"] = depth_paths
+            self.metadata["depth_filenames"] = depth_paths
+            dataparser_outputs.metadata["depth_unit_scale_factor"] = 1.0
+            self.metadata["depth_unit_scale_factor"] = 1.0
+            
         self.depth_filenames = self.metadata["depth_filenames"]
         self.depth_unit_scale_factor = self.metadata["depth_unit_scale_factor"]
 
@@ -54,3 +95,11 @@ class DepthDataset(InputDataset):
         )
 
         return {"depth_image": depth_image}
+
+    def _find_transform(self, image_path : Path) -> Path:
+        while image_path.parent != image_path:
+            transform_path = image_path.parent / "transforms.json"
+            if transform_path.exists():
+                return transform_path
+            image_path = image_path.parent
+        raise FileNotFoundError("Could not find transforms.json in any parent directory of image_path")

nerfstudio/model_components/losses.py

@@ -560,3 +560,16 @@ def scale_gradients_by_distance_squared(
     for key, value in field_outputs.items():
         out[key], _ = cast(Tuple[Tensor, Tensor], _GradientScaler.apply(value, scaling))
     return out
+
+
+def depth_ranking_loss(rendered_depth, gt_depth):
+    """
+    Depth ranking loss as described in the SparseNeRF paper
+    Assumes that the layout of the batch comes from a PairPixelSampler, so that adjacent samples in the gt_depth 
+    and rendered_depth are from pixels with a radius of each other
+    """
+    m = 1e-4
+    dpt_diff = gt_depth[::2, :] - gt_depth[1::2, :]
+    out_diff = rendered_depth[::2, :] - rendered_depth[1::2, :] + m
+    differing_signs = torch.sign(dpt_diff) != torch.sign(out_diff)
+    return torch.nanmean((out_diff[differing_signs] * torch.sign(out_diff[differing_signs])))

nerfstudio/models/nerfacto.py

@@ -45,6 +45,8 @@ from nerfstudio.model_components.losses import (
     orientation_loss,
     pred_normal_loss,
     scale_gradients_by_distance_squared,
+    depth_ranking_loss,
+
 )
 from nerfstudio.model_components.ray_samplers import (
     ProposalNetworkSampler,
@@ -223,7 +225,8 @@ class NerfactoModel(Model):
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=self.config.background_color)
         self.renderer_accumulation = AccumulationRenderer()
-        self.renderer_depth = DepthRenderer()
+        self.renderer_depth = DepthRenderer('expected')
+        self.renderer_depth_med = DepthRenderer()
         self.renderer_normals = NormalsRenderer()
 
         # shaders
@@ -236,6 +239,8 @@ class NerfactoModel(Model):
         self.psnr = PeakSignalNoiseRatio(data_range=1.0)
         self.ssim = structural_similarity_index_measure
         self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True)
+        self.step=0
+
 
     def get_param_groups(self) -> Dict[str, List[Parameter]]:
         param_groups = {}
@@ -254,6 +259,8 @@ class NerfactoModel(Model):
             def set_anneal(step):
                 # https://arxiv.org/pdf/2111.12077.pdf eq. 18
                 train_frac = np.clip(step / N, 0, 1)
+                self.step=step
+
 
                 def bias(x, b):
                     return b * x / ((b - 1) * x + 1)
@@ -290,12 +297,15 @@ class NerfactoModel(Model):
 
         rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
         depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
+        with torch.no_grad():
+            depth_med = self.renderer_depth_med(weights=weights, ray_samples=ray_samples)
         accumulation = self.renderer_accumulation(weights=weights)
 
         outputs = {
             "rgb": rgb,
             "accumulation": accumulation,
             "depth": depth,
+            "depth_med": depth_med,
         }
 
         if self.config.predict_normals:
@@ -352,6 +362,9 @@ class NerfactoModel(Model):
                 loss_dict["pred_normal_loss"] = self.config.pred_normal_loss_mult * torch.mean(
                     outputs["rendered_pred_normal_loss"]
                 )
+            
+            if 'depth_image' in batch and self.step>=500:
+                loss_dict['depth_ranking'] = np.interp(self.step,[500,1000],[0,.2])*depth_ranking_loss(outputs['depth'],batch['depth_image'])
         return loss_dict
 
     def get_image_metrics_and_images(

pyproject.toml

@@ -55,7 +55,8 @@ dependencies = [
     "nuscenes-devkit>=1.1.1",
     "wandb>=0.13.3",
     "xatlas",
-    "trimesh>=3.20.2"
+    "trimesh>=3.20.2",
+    "timm==0.6.7"
 ]
 
 [project.urls]