[UPDATE] Saved to different TTA config folders depends on button press

app.py

@@ -187,15 +187,19 @@ def process(
 
     # ------------------------------------------------------------------
 
+    # Define save gif dir
+    planner.execute_tta = with_tta
+    gifs_save_dir = os.path.join(gifs_path, "no_tta") if not with_tta else os.path.join(gifs_path, "with_tta")
+    planner.gifs_path = gifs_save_dir
+
     # Empty gifs_path folder
-    if os.path.exists(gifs_path):
-        for file in os.listdir(gifs_path):
-            os.remove(os.path.join(gifs_path, file))
+    if os.path.exists(gifs_save_dir):
+        for file in os.listdir(gifs_save_dir):
+            os.remove(os.path.join(gifs_save_dir, file))
 
     # Optionally you may want to reset episode index or make it configurable.
     # For now we hard-code episode 0, mirroring the snippet.
     # Set execute_tta flag depending on button pressed
-    planner.execute_tta = with_tta
 
     t = threading.Thread(target=planner.run_episode, args=(0,), daemon=True)
     t.start()
@@ -205,7 +209,7 @@ def process(
     try:
         while t.is_alive():
             # discover any new pngs written by TestWorker
-            pngs = glob.glob(os.path.join(gifs_path, "*.png"))
+            pngs = glob.glob(os.path.join(gifs_save_dir, "*.png"))
             pngs.sort(key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
             for fp in pngs:
                 if fp not in sent:

test_multi_robot_worker.py

@@ -56,6 +56,9 @@ class TestWorker:
         self.perf_metrics = dict()
         self.bad_mask_init = False
 
+        # NOTE: Option to override gifs_path to interface with app.py
+        self.gifs_path = gifs_path
+
         # # TEMP - EXPORT START POSES FOR BASELINES
         # json_path = "eval_start_positions.json"
         # sat_to_start_pose_dict = {}
@@ -152,7 +155,7 @@ class TestWorker:
                     n_smooth_iter=2,    # smoothing parameter
                     ignore_label=-1,
                     plot=False,     # NOTE: Set to false since using app.py
-                    gifs_dir = gifs_path      
+                    gifs_dir = self.gifs_path   # NOTE: Set to self.gifs_path since using app.py      
                 )    
                 # Fit & predict (this will also plot the clusters before & after smoothing)
                 map_shape = (int(np.sqrt(self.clip_seg_tta.patch_embeds.shape[0])), int(np.sqrt(self.clip_seg_tta.patch_embeds.shape[0])))
@@ -284,13 +287,14 @@ class TestWorker:
                 robots_route = []
                 for robot in self.robot_list:
                     robots_route.append([robot.xPoints, robot.yPoints])
-                if not os.path.exists(gifs_path):
-                    os.makedirs(gifs_path)
+                # NOTE: Set to self.gifs_path since using app.py
+                if not os.path.exists(self.gifs_path):  
+                    os.makedirs(self.gifs_path)
                 sound_id_override = None if self.clip_seg_tta.sound_ids == [] else self.clip_seg_tta.sound_ids[0]
 
                 ## NOTE: Replaced since using app.py
                 if TAXABIND_TTA and USE_CLIP_PREDS:
-                    self.env.plot_heatmap(gifs_path, step, max(travel_dist_list), robots_route)
+                    self.env.plot_heatmap(self.gifs_path, step, max(travel_dist_list), robots_route)
                 # if TAXABIND_TTA and USE_CLIP_PREDS:
                 #     self.env.plot_env(
                 #         self.global_step, 
@@ -363,7 +367,7 @@ class TestWorker:
 
         # save gif
         if self.save_image:
-            path = gifs_path
+            path = self.gifs_path # NOTE: Set to self.gifs_path since using app.py
             self.make_gif(path, curr_episode)
 
         print(YELLOW, f"[Eps {curr_episode} Completed] Time Taken: {time()-eps_start:.2f}s, Steps: {step+1}", NC)