From 4f9b2cdc7ec947c35b2bce01de16675f66079151 Mon Sep 17 00:00:00 2001 From: Zhensheng Yuan <564361+yzslab@users.noreply.github.com> Date: Fri, 5 Jul 2024 13:00:58 +0800 Subject: [PATCH] Add round half up mode to colmap --- README.md | 3 +++ internal/dataparsers/colmap_dataparser.py | 10 ++++++++-- utils/eval_mipnerf360.py | 8 ++------ 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 2c3758ba..5bf01ba5 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,9 @@ You can use `utils/image_downsample.py` to downsample your images, e.g. 4x downs --data.parser.down_sample_factor 4 \ ... ``` + +Rounding mode is specified by `--data.parser.down_sample_rounding_model`. Available values are `floor`, `round`, `round_half_up`, `ceil`. Default is `round`. + * Load large dataset without OOM ```bash ... fit \ diff --git a/internal/dataparsers/colmap_dataparser.py b/internal/dataparsers/colmap_dataparser.py index 31ee8023..9bc7f92c 100644 --- a/internal/dataparsers/colmap_dataparser.py +++ b/internal/dataparsers/colmap_dataparser.py @@ -54,7 +54,7 @@ class Colmap(DataParserConfig): down_sample_factor: int = 1 - down_sample_rounding_model: Literal["floor", "round", "ceil"] = "round" + down_sample_rounding_model: Literal["floor", "round", "round_half_up", "ceil"] = "round" def instantiate(self, path: str, output_path: str, global_rank: int) -> DataParser: return ColmapDataParser(path, output_path, global_rank, self) @@ -68,6 +68,9 @@ def __init__(self, path: str, output_path: str, global_rank: int, params: Colmap self.global_rank = global_rank self.params = params + def _round_half_up(self, i: torch.Tensor): + return torch.floor(i + 0.5) + def detect_sparse_model_dir(self) -> str: if os.path.isdir(os.path.join(self.path, "sparse", "0")): return os.path.join(self.path, "sparse", "0") @@ -344,7 +347,10 @@ def get_outputs(self) -> DataParserOutputs: # recalculate intrinsics if down sample enabled if self.params.down_sample_factor != 1: - rounding_func = getattr(torch, self.params.down_sample_rounding_model) + if self.params.down_sample_rounding_model == "round_half_up": + rounding_func = self._round_half_up + else: + rounding_func = getattr(torch, self.params.down_sample_rounding_model) down_sampled_width = rounding_func(width.to(torch.float) / self.params.down_sample_factor) down_sampled_height = rounding_func(height.to(torch.float) / self.params.down_sample_factor) width_scale_factor = down_sampled_width / width diff --git a/utils/eval_mipnerf360.py b/utils/eval_mipnerf360.py index 753d81fd..0ecfbe4b 100644 --- a/utils/eval_mipnerf360.py +++ b/utils/eval_mipnerf360.py @@ -19,10 +19,6 @@ def start(command: str, scene: str, extra_args: list = None): - down_sample_rounding_model = "ceil" - if scene == "garden": - down_sample_rounding_model = "round" - arg_list = [ "python", "main.py", @@ -31,7 +27,7 @@ def start(command: str, scene: str, extra_args: list = None): "--data.parser", "Colmap", "--data.parser.down_sample_factor", "4", "--data.parser.split_mode", "experiment", - "--data.parser.down_sample_rounding_model", down_sample_rounding_model, + "--data.parser.down_sample_rounding_model", "round_half_up", "--cache_all_images", "--logger", "wandb", "--output", os.path.join("outputs", args.project), @@ -43,7 +39,7 @@ def start(command: str, scene: str, extra_args: list = None): if extra_args is not None: arg_list += extra_args - subprocess.call(arg_list) + return subprocess.call(arg_list) with tqdm(scenes) as t: