cosypose inference example: add a main scope, ref agimus-project#159

nim65s · Apr 17, 2024 · a8746e7 · a8746e7
1 parent c412cb3
commit a8746e7
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py b/happypose/pose_estimators/cosypose/cosypose/scripts/run_inference_on_example.py
@@ -60,7 +60,7 @@ def run_inference(
     return data_TCO.cpu()
 
 
-if __name__ == "__main__":
+def main():
     set_logging_level("info")
     parser = argparse.ArgumentParser()
     parser.add_argument("example_name")
@@ -77,7 +77,7 @@ def run_inference(
     assert (
         example_dir.exists()
     ), "Example {args.example_name} not available, follow download instructions"
-    dataset_to_use = args.dataset  # hope/tless/ycbv
+    # dataset_to_use = args.dataset  # hope/tless/ycbv
 
     # Load data
     detections = load_detections(example_dir).to(device)
@@ -118,3 +118,7 @@ def run_inference(
         make_poses_visualization(
             rgb, object_dataset, object_datas, camera_data, example_dir
         )
+
+
+if __name__ == "__main__":
+    main()