From 1dc246c1ffcee857223376955610f9ce68b165d9 Mon Sep 17 00:00:00 2001 From: AdamRashid96 Date: Mon, 12 Jun 2023 19:07:54 +0000 Subject: [PATCH] monodepth implemented --- nerfstudio/configs/method_configs.py | 2 + nerfstudio/data/datasets/depth_dataset.py | 61 ++++++++++++++++++++--- nerfstudio/model_components/losses.py | 13 +++++ nerfstudio/models/nerfacto.py | 15 +++++- pyproject.toml | 3 +- 5 files changed, 86 insertions(+), 8 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 027b01352..db903c4a2 100644 --- a/nerfstudio/configs/method_configs.py +++ b/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, diff --git a/nerfstudio/data/datasets/depth_dataset.py b/nerfstudio/data/datasets/depth_dataset.py index 1d64cdf1c..7ee4754f6 100644 --- a/nerfstudio/data/datasets/depth_dataset.py +++ b/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") diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py index 604f306e6..51e575088 100644 --- a/nerfstudio/model_components/losses.py +++ b/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]))) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 80ead5951..5cc281e03 100644 --- a/nerfstudio/models/nerfacto.py +++ b/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( diff --git a/pyproject.toml b/pyproject.toml index 02972b0e9..7420c9217 100644 --- a/pyproject.toml +++ b/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]