Simplify the combine_results script

lamar/combine_results.py

@@ -5,18 +5,24 @@
 
 from . import logger
 
+LOCATIONS = ["CAB", "LIN", "HGE"]
+DEVICES = ["hololens", "phone"]
+
+
 def assert_valid_txt_path(path: Path):
     """Assert that the txt path is valid."""
     assert path.exists()
     assert path.is_file()
     assert path.suffix == ".txt"
 
 
-def combine_results(description_path: Path, results_paths: dict[str, Path], output_dir: Path):
+def combine_results(
+    description_path: Path, results_paths: dict[str, Path | None], output_dir: Path
+):
     # Generate timestamp for the zip file name.
     timestamp = datetime.datetime.now().strftime("%Y.%m.%d_%H.%M.%S")
     zip_filename = output_dir / f"submission_{timestamp}.zip"
-    
+
     # Create the zip file with existing paths.
     logger.info(f"Creating zip file at {zip_filename}")
     output_dir.mkdir(parents=True, exist_ok=True)
@@ -39,53 +45,27 @@ def combine_results(description_path: Path, results_paths: dict[str, Path], outp
     logger.info(f"Successfully created zip file at {zip_filename}")
 
 
-
-if __name__ == '__main__':
+if __name__ == "__main__":
     parser = argparse.ArgumentParser(
-        description="Combine estimated poses from multiple scenes / devices in a zip file for evaluation."
+        description=(
+            "Combine the estimated poses from multiple scenes / devices in a "
+            "zip file for evaluation on CodaBench."
+        )
     )
     parser.add_argument(
         "--description_path",
         type=Path,
         required=True,
-        help="Path to a text file containing description of the submission.",
-    )
-    parser.add_argument(
-        "--cab_phone_path",
-        type=Path,
-        default=None,
-        help="File containing the CAB phone estimated poses.",
-    )
-    parser.add_argument(
-        "--lin_phone_path",
-        type=Path,
-        default=None,
-        help="File containing the LIN phone estimated poses.",
-    )
-    parser.add_argument(
-        "--hge_phone_path",
-        type=Path,
-        default=None,
-        help="File containing the HGE phone estimated poses.",
-    )
-    parser.add_argument(
-        "--cab_hololens_path",
-        type=Path,
-        default=None,
-        help="File containing the CAB HoloLens estimated poses.",
-    )
-    parser.add_argument(
-        "--lin_hololens_path",
-        type=Path,
-        default=None,
-        help="File containing the LIN HoloLens estimated poses.",
-    )
-    parser.add_argument(
-        "--hge_hololens_path",
-        type=Path,
-        default=None,
-        help="File containing the HGE HoloLens estimated poses.",
+        help="Path to a text file containing the description of the submission.",
     )
+    for location in LOCATIONS:
+        for device in DEVICES:
+            parser.add_argument(
+                f"--{location.lower()}_{device}_path",
+                type=Path,
+                default=None,
+                help=f"File containing the estimated poses for {location} {device}.",
+            )
     parser.add_argument(
         "--output_dir",
         type=Path,
@@ -95,6 +75,7 @@ def combine_results(description_path: Path, results_paths: dict[str, Path], outp
     args = parser.parse_args().__dict__
     output_dir = args.pop("output_dir")
     description_path = args.pop("description_path")
+    assert len(args) == len(LOCATIONS) * len(DEVICES)
 
     combine_results(
         description_path=description_path,