Integrating AerialMegaDepth dataset

data_processing/README.md

@@ -13,13 +13,14 @@ We also provide a HuggingFace dataset which contains the [pre-computed metadata]
 5. ✅ [ETH3D](https://www.eth3d.net/)
 6. ✅ [Mapillary Planet Scale Depth & Reconstructions](https://www.mapillary.com/dataset/depth) (MPSD)
 7. ✅ [MegaDepth (including Tanks & Temples)](https://www.cs.cornell.edu/projects/megadepth/)
-8. ✅ [MVS-Synth](https://phuang17.github.io/DeepMVS/mvs-synth.html)
-9. ✅ [Parallel Domain 4D](https://gcd.cs.columbia.edu/#datasets)
-10. ✅ [SAIL-VOS 3D](https://sailvos.web.illinois.edu/_site/_site/index.html)
-11. ✅ [ScanNet++ v2](https://kaldir.vc.in.tum.de/scannetpp/)
-12. ✅ [Spring](https://spring-benchmark.org/)
-13. ✅ [TartanAirV2 Wide Baseline](https://uniflowmatch.github.io/)
-14. ✅ [UnrealStereo4K](https://github.com/fabiotosi92/SMD-Nets)
+8. ✅ [AerialMegaDepth](https://aerial-megadepth.github.io/)
+9. ✅ [MVS-Synth](https://phuang17.github.io/DeepMVS/mvs-synth.html)
+10. ✅ [Parallel Domain 4D](https://gcd.cs.columbia.edu/#datasets)
+11. ✅ [SAIL-VOS 3D](https://sailvos.web.illinois.edu/_site/_site/index.html)
+12. ✅ [ScanNet++ v2](https://kaldir.vc.in.tum.de/scannetpp/)
+13. ✅ [Spring](https://spring-benchmark.org/)
+14. ✅ [TartanAirV2 Wide Baseline](https://uniflowmatch.github.io/)
+15. ✅ [UnrealStereo4K](https://github.com/fabiotosi92/SMD-Nets)
 
 ## Download Instructions:
 

data_processing/aggregate_scene_names.py

@@ -23,6 +23,7 @@
     DL3DV10KSplits,
     ETH3DSplits,
     MegaDepthSplits,
+    AerialMegaDepthSplits,
     MPSDSplits,
     ScanNetPPSplits,
     SpringSplits,
@@ -363,6 +364,24 @@ def aggregate(self):
         super().aggregate(val_split_scenes=self.dataset_split_info.val_split_scenes)
 
 
+class AerialMegaDepthAggregator(DatasetAggregator):
+    """Aggregator for AerialMegaDepth dataset."""
+
+    def __init__(self, root_dir, output_dir, covisibility_version_key="v0"):
+        super().__init__(
+            dataset_name="aerialmegadepth",
+            root_dir=root_dir,
+            output_dir=output_dir,
+            covisibility_version_key=covisibility_version_key,
+            depth_folder="depth",
+        )
+        self.dataset_split_info = AerialMegaDepthSplits()
+
+    def aggregate(self):
+        """Aggregate the AerialMegaDepth dataset."""
+        super().aggregate(val_split_scenes=self.dataset_split_info.val_split_scenes)
+
+
 class MPSDAggregator(DatasetAggregator):
     """Aggregator for MPSD dataset."""
 
@@ -551,6 +570,7 @@ def main():
             "dynamicreplica",
             "eth3d",
             "megadepth",
+            "aerialmegadepth",
             "mpsd",
             "mvs_synth",
             "paralleldomain4d",
@@ -567,6 +587,7 @@ def main():
             "dynamicreplica",
             "eth3d",
             "megadepth",
+            "aerialmegadepth",
             "mpsd",
             "mvs_synth",
             "paralleldomain4d",
@@ -637,6 +658,12 @@ def main():
                 root_dir=root_dir, output_dir=args.output_dir
             )
             aggregator.aggregate()
+        elif dataset == "aerialmegadepth":
+            # AerialMegaDepth
+            aggregator = AerialMegaDepthAggregator(
+                root_dir=root_dir, output_dir=args.output_dir
+            )
+            aggregator.aggregate()
         elif dataset == "mpsd":
             # MPSD
             raw_data_root_dir = os.path.join(args.raw_data_root, "mpsd")

data_processing/viz_data.py

@@ -202,6 +202,16 @@ def get_dataset_config(dataset_type):
             "confidence_key": None,
             "confidence_thres": 0.0,
         },
+        "aerialmegadepth": {
+            "root_dir": "/fsx/xrtech/data/aerialmegadepth",
+            "scene": "0001",
+            "depth_key": "depth",
+            "local_frame": False,
+            "viz_string": "WAI_Viz",
+            "load_skymask": False,
+            "confidence_key": None,
+            "confidence_thres": 0.0,
+        },
         "spring": {
             "root_dir": "/fsx/xrtech/data/spring",
             "scene": "0004",
@@ -331,6 +341,7 @@ def get_parser():
             "blendedmvs",
             "eth3d",
             "megadepth",
+            "aerialmegadepth",
             "spring",
             "mpsd",
             "ase",

data_processing/wai_processing/configs/conversion/aerialmegadepth.yaml

@@ -0,0 +1,9 @@
+original_root: # path of raw downloaded dataset
+root: # path of wai-formatted dataset
+
+dataset_name: aerialmegadepth
+version: 0.1
+overwrite: True
+
+scene_filters:
+  - process_state_not: [conversion, finished]

data_processing/wai_processing/configs/launch/aerialmegadepth.yaml

@@ -0,0 +1,24 @@
+stage: # set stage via CLI
+root: # path of wai-formatted dataset
+
+gpus: 0
+cpus: 10
+mem: 20
+scenes_per_job: 20
+conda_env: # pass the name of our conda environment
+nodelist:
+
+stages:
+  conversion:
+    script: conversion/aerialmegadepth.py
+    config: conversion/aerialmegadepth.yaml
+    scenes_per_job: 20 # fast
+  covisibility:
+    script: covisibility.py
+    config: covisibility/covisibility_gt_depth_224x224.yaml
+    gpus: 1
+  moge:
+    script: run_moge.py
+    config: moge/default.yaml
+    additional_cli_params: ['batch_size=1']
+    gpus: 1

data_processing/wai_processing/download_scripts/README.md

@@ -28,6 +28,10 @@ Use the provided bash script `download_mpsd.sh` and unzip the downloaded zip fil
 Use the provided python script `download_megadepth.py`. MegaDepth also includes all the Tanks & Temples scenes.
 **Source:** [MegaDepth Project](https://www.cs.cornell.edu/projects/megadepth/)
 
+## AerialMegaDepth
+Use the provided python script `download_aerialmegadepth.py`.
+**Source:** [AerialMegaDepth](https://aerial-megadepth.github.io/)
+
 ## MVS-Synth
 Use the provided python script `download_mvs_synth.py`.
 **Source:** [MVS-Synth Dataset](https://phuang17.github.io/DeepMVS/mvs-synth.html)

data_processing/wai_processing/download_scripts/download_aerialmegadepth.py

@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+"""
+Download AerialMegaDepth dataset from HuggingFace.
+
+References: https://github.com/kvuong2711/aerial-megadepth/blob/main/data_generation/download_data_hf.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+from pathlib import Path
+from huggingface_hub import snapshot_download
+from wai_processing.utils.download import (
+    extract_zip_archives,
+)
+
+
+# Configuration for AerialMegaDepth dataset
+REPO_ID = "kvuong2711/aerialmegadepth"
+ALLOW_PATTERNS = ("**.zip", "aerial_megadepth_all.npz")
+DEFAULT_MAX_WORKERS = 8
+ZIP_DIR_NAME = "aerialmegadepth_zip"
+EXTRACT_DIR_NAME = "aerialmegadepth"
+
+
+def download_archives(zip_dir: Path, max_workers: int):
+    """Download dataset archives into ``zip_dir``."""
+    zip_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Downloading {REPO_ID} archives to {zip_dir}...")
+    snapshot_download(
+        repo_id=REPO_ID,
+        repo_type="dataset",
+        local_dir=str(zip_dir),
+        max_workers=max_workers,
+        allow_patterns=list(ALLOW_PATTERNS),
+    )
+    print(f"Download complete!")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Download the Aerial MegaDepth dataset and optionally extract archives.",
+    )
+    parser.add_argument(
+        "--target_dir",
+        type=str,
+        required=True,
+        help="Base directory for downloaded archives and extracted data.",
+    )
+    parser.add_argument(
+        "--max_workers",
+        type=int,
+        default=DEFAULT_MAX_WORKERS,
+        help="Number of parallel workers used by the Hugging Face downloader.",
+    )
+    return parser
+
+
+def main():
+    parser = build_parser()
+    args = parser.parse_args()
+
+    target_dir = Path(args.target_dir)
+    zip_dir = target_dir / ZIP_DIR_NAME
+    extract_dir = target_dir / EXTRACT_DIR_NAME
+
+    # 1. Download zip files from huggingface
+    download_archives(zip_dir, max_workers=args.max_workers)
+
+    # 2. Extract zip files
+    extract_zip_archives(target_dir=zip_dir, output_dir=extract_dir, n_workers=args.max_workers)
+
+    # 3. Move the aerial_megadepth_all.npz to the extract_dir
+    shutil.move(zip_dir / "aerial_megadepth_all.npz", extract_dir / "aerial_megadepth_all.npz")
+
+    print("All tasks completed successfully.")
+
+
+if __name__ == "__main__":
+    main()