make saving perception results to be separate thread

bdaiinstitute · Sep 13, 2024 · 8a4f0a4 · 8a4f0a4
1 parent 09fa53f
commit 8a4f0a4
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 13 deletions.
diff --git a/predicators/envs/spot_env.py b/predicators/envs/spot_env.py
@@ -10,6 +10,7 @@
 from pathlib import Path
 from typing import Callable, ClassVar, Collection, Dict, Iterator, List, \
     Optional, Sequence, Set, Tuple
+import threading
 
 import matplotlib
 import numpy as np
@@ -646,14 +647,6 @@ def _build_realworld_observation(
         all_detections, all_artifacts = detect_objects(
             all_object_detection_ids, rgbds, self._allowed_regions)
 
-        if CFG.spot_render_perception_outputs:
-            outdir = Path(CFG.spot_perception_outdir)
-            time_str = time.strftime("%Y%m%d-%H%M%S")
-            detections_outfile = outdir / f"detections_{time_str}.png"
-            no_detections_outfile = outdir / f"no_detections_{time_str}.png"
-            visualize_all_artifacts(all_artifacts, detections_outfile,
-                                    no_detections_outfile)
-
         # Separately, get detections for the hand in particular.
         hand_rgbd = {
             k: v
@@ -663,10 +656,14 @@ def _build_realworld_observation(
             all_object_detection_ids, hand_rgbd, self._allowed_regions)
 
         if CFG.spot_render_perception_outputs:
-            detections_outfile = outdir / f"hand_detections_{time_str}.png"
-            no_detect_outfile = outdir / f"hand_no_detections_{time_str}.png"
-            visualize_all_artifacts(hand_artifacts, detections_outfile,
-                                    no_detect_outfile)
+            outdir = Path(CFG.spot_perception_outdir)
+            time_str = time.strftime("%Y%m%d-%H%M%S")
+            detections_outfile = outdir / f"detections_{time_str}.png"
+            no_detections_outfile = outdir / f"no_detections_{time_str}.png"
+            hand_detections_outfile = outdir / f"hand_detections_{time_str}.png"
+            hand_no_detections_outfile = outdir / f"hand_no_detections_{time_str}.png"
+            threading.Thread(target=visualize_all_artifacts, args=(all_artifacts, detections_outfile, no_detections_outfile)).start()
+            threading.Thread(target=visualize_all_artifacts, args=(hand_artifacts, hand_detections_outfile, hand_no_detections_outfile)).start()
 
         # Also, get detections that every camera except the back camera can
         # see. This is important for our 'InView' predicate.

diff --git a/predicators/perception/spot_perceiver.py b/predicators/perception/spot_perceiver.py
@@ -567,7 +567,6 @@ def _create_goal(self, state: State,
                 GroundAtom(ContainingWater, [cup]),
             }
         if goal_description == "place empty cup into the box":
-            # TODO
             cup = Object("cup", _container_type)
             cardboard_box = Object("cardboard_box", _container_type)
             ContainingWaterKnown = pred_name_to_pred["ContainingWaterKnown"]