diff --git a/lamar/combine_results.py b/lamar/combine_results.py index 6368d0f..131e452 100644 --- a/lamar/combine_results.py +++ b/lamar/combine_results.py @@ -5,6 +5,10 @@ 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() @@ -12,11 +16,13 @@ def assert_valid_txt_path(path: Path): 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,