few fixes, run_inference return predictions

agimus-project · Sep 7, 2023 · 5a02066 · 5a02066
1 parent cbacd8f
commit 5a02066
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 10 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
@@ -169,23 +169,33 @@ def run_inference(
     example_dir: Path,
     model_name: str,
     dataset_to_use: str,
+    example_name: str,
+    render: bool = True,
 ) -> None:
-    observation = load_observation_tensor(example_dir)
-    CosyPose = CosyPoseWrapper(dataset_name=dataset_to_use, n_workers=8)
-    predictions = CosyPose.inference(observation)
-    renderings = rendering(predictions, example_dir)
-    save_predictions(example_dir, renderings)
+    try:
+        observation = load_observation_tensor(example_dir)
+        CosyPose = CosyPoseWrapper(dataset_name=dataset_to_use, n_workers=8)
+        predictions = CosyPose.inference(observation)
+        if render:
+            renderings = rendering(predictions, example_dir, example_name)
+            save_predictions(example_dir, renderings)
+
+        return predictions.poses, predictions.infos
+    except AttributeError as err:
+        print("found nothing")
+        return [],[]
 
 
 if __name__ == "__main__":
     set_logging_level("info")
     parser = argparse.ArgumentParser()
     parser.add_argument("example_name")
     parser.add_argument("--model", type=str, default="megapose-1.0-RGB-multi-hypothesis")
-    parser.add_argument("--dataset", type=str, default="ycbv")
+    parser.add_argument("--dataset", type=str, default="tless")
     #parser.add_argument("--vis-detections", action="store_true")
     parser.add_argument("--run-inference", action="store_true", default=True)
     #parser.add_argument("--vis-outputs", action="store_true")
+    parser.add_argument("--render", action="store_true", default=False)
     args = parser.parse_args()
 
     data_dir = os.getenv("MEGAPOSE_DATA_DIR")
@@ -197,7 +207,7 @@ def run_inference(
     #    make_detections_visualization(example_dir)
 
     if args.run_inference:
-        run_inference(example_dir, args.model, dataset_to_use)
+        run_inference(example_dir, args.model, dataset_to_use, args.example_name, args.render)
 
     #if args.vis_outputs:
     #    make_output_visualization(example_dir)
diff --git a/happypose/pose_estimators/cosypose/cosypose/simulator/base_scene.py b/happypose/pose_estimators/cosypose/cosypose/simulator/base_scene.py
@@ -12,7 +12,7 @@ class BaseScene:
     _connected = False
     _simulation_step = 1/240.
 
-    def connect(self, gpu_renderer=True, gui=False):
+    def connect(self, gpu_renderer=False, gui=False):
         assert not self._connected, 'Already connected'
         if gui:
             self._client_id = pb.connect(pb.GUI, '--width=640 --height=480')

diff --git a/happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py b/happypose/pose_estimators/cosypose/cosypose/utils/cosypose_wrapper.py
@@ -110,7 +110,7 @@ def load_pose_models(coarse_run_id, refiner_run_id, n_workers):
         #renderer = BulletBatchRenderer(object_set=cfg.urdf_ds_name, n_workers=n_workers, gpu_renderer=gpu_renderer)
         #
 
-        object_dataset = make_object_dataset("ycbv")
+        object_dataset = make_object_dataset("tless.cad")
         mesh_db = MeshDataBase.from_object_ds(object_dataset)
         renderer = Panda3dBatchRenderer(object_dataset, n_workers=n_workers, preload_cache=False)
         mesh_db_batched = mesh_db.batched().to(device)
@@ -189,4 +189,4 @@ def inference(self, observation, coarse_guess=None):
                 n_coarse_iterations=0, n_refiner_iterations=4)
         print("inference successfully.")
         # result: this_batch_detections, final_preds
-        return final_preds.cpu()
+        return final_preds.cpu()