[NEW] Able to interface with TestWorker with minimal changes to VLM-Search related files. Able to visualize robot movement gif on gradio

.vscode/launch.json

@@ -0,0 +1,15 @@
+{
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Debug app.py",
+            "type": "debugpy",
+            "request": "launch",
+            "program": "${workspaceFolder}/app.py",
+            "cwd": "${workspaceFolder}",
+            "console": "integratedTerminal",
+            "justMyCode": false,
+            "python": "/home/user/anaconda3/envs/vlm-search/bin/python3"
+        }
+    ]
+}

Taxabind/TaxaBind/SatBind/clip_seg_tta.py

@@ -116,7 +116,9 @@ class ClipSegTTA:
                 self.sound_model.eval()
 
         # Params
-        self.reset(sample_idx=self.sample_index)
+        # NOTE: ADDED CONDITIONAL
+        if sample_index >= 0:
+            self.reset(sample_idx=self.sample_index)
 
         # for idx, related_img_path in enumerate(self.related_imgs_paths):
         #     self.visualize_heatmap(
@@ -153,8 +155,10 @@ class ClipSegTTA:
         if self.load_model:
             self.reset_local_model()    # Reset to global weights as init
 
-        self.img_paths, self.imo_path, self.imgs, self.imo, self.sounds, self.sound_ids, self.species_name, self.target_positions, self.gt_mask_name = self.dataset.get_search_ds_data(sample_idx)
-        self.imgs = self.imgs.to(self.device)
+        # NOTE: Added 'sample_idx' conditional
+        if sample_idx >= 0:
+            self.img_paths, self.imo_path, self.imgs, self.imo, self.sounds, self.sound_ids, self.species_name, self.target_positions, self.gt_mask_name = self.dataset.get_search_ds_data(sample_idx)
+            self.imgs = self.imgs.to(self.device)
         # self.img_path = self.img_paths[0]       # Select 1st img as query
         # self.img = self.imgs[0]                 # Select 1st img as query
         # self.img = self.img.to(self.device)     

Taxabind/TaxaBind/SatBind/dataset.py

@@ -32,42 +32,43 @@ class SatNatDataset(Dataset):
         # ADDED
         self.current_epoch = 0
 
-        self.json = json.load(open(json_path, 'r'))
-        # self.sat_filt_json = json.load(open(sat_filtered_json_path, 'r'))
-        self.sat_to_img_ids_json_path = json.load(open(sat_to_img_ids_json_path, 'r'))
-        self.images = self.json['images']
-        self.annot = self.json['annotations']
-        for i in range(len(self.images)):
-            assert self.images[i]['id'] == self.annot[i]['id']
-            self.images[i]['label'] = self.annot[i]['category_id']
-        self.filtered_json = [d for d in self.images if d['latitude'] is not None and d['longitude'] is not None]
-        self.species_text = list(set([" ".join(d['file_name'].split("/")[1].split("_")[1:]) for d in self.filtered_json]))
-
-        # self.sat_filtered_json = [d for d in self.sat_filt_json['images'] if d['latitude'] is not None and d['longitude'] is not None]
-        # self.sat_paths = {d['id']: str(d['id'])+'_'+str(d['latitude'])+'_'+str(d['longitude'])+'.jpg' for d in self.sat_filtered_json}
-        self.inat_json_dict = {
-            "images": {img["id"]: img for img in self.images},
-            "annotations": {ann["id"]: ann for ann in self.annot},
-        }
-
-        # Expand dict
-        self.sat_to_img_ids_tuples = []
-        if self.sat_to_img_ids_json_is_train_dict:
-            # for i in range(len(self.sat_filtered_json)):
-            #     id = self.sat_filtered_json[i]['id']
-            #     sat_id = Path(self.sat_paths[id]).stem      # remove file extension
-            #     img_ids = self.sat_to_img_ids_json_path[sat_id]["img_ids"]
-            for sat_key, sat_sample in self.sat_to_img_ids_json_path.items():
-                id = sat_sample["id"] # int(sat_key.split("_")[0]) # sat_sample['id']
-                sat_path = sat_sample["sat_path"]
-                img_ids = sat_sample["img_ids"]
-                # img_locs = sat_sample["target_positions"] # NOTE: Not accurate after resize
-                for img_id in img_ids:
-                    self.sat_to_img_ids_tuples.append((id, sat_path, img_id))
-            print("len(self.sat_to_img_ids_json_path): ", len(self.sat_to_img_ids_json_path))
-            print("len(self.sat_to_img_ids_tuples): ", len(self.sat_to_img_ids_tuples))
-        else:
-            self.filtered_val_ds_by_tax = [d for d in self.sat_to_img_ids_json_path if self.tax_to_filter_val in d['taxonomy']] 
+        ## NOTE: Removed as unnecessary for app.py
+        # self.json = json.load(open(json_path, 'r'))
+        # # self.sat_filt_json = json.load(open(sat_filtered_json_path, 'r'))
+        # self.sat_to_img_ids_json_path = json.load(open(sat_to_img_ids_json_path, 'r'))
+        # self.images = self.json['images']
+        # self.annot = self.json['annotations']
+        # for i in range(len(self.images)):
+        #     assert self.images[i]['id'] == self.annot[i]['id']
+        #     self.images[i]['label'] = self.annot[i]['category_id']
+        # self.filtered_json = [d for d in self.images if d['latitude'] is not None and d['longitude'] is not None]
+        # self.species_text = list(set([" ".join(d['file_name'].split("/")[1].split("_")[1:]) for d in self.filtered_json]))
+
+        # # self.sat_filtered_json = [d for d in self.sat_filt_json['images'] if d['latitude'] is not None and d['longitude'] is not None]
+        # # self.sat_paths = {d['id']: str(d['id'])+'_'+str(d['latitude'])+'_'+str(d['longitude'])+'.jpg' for d in self.sat_filtered_json}
+        # self.inat_json_dict = {
+        #     "images": {img["id"]: img for img in self.images},
+        #     "annotations": {ann["id"]: ann for ann in self.annot},
+        # }
+
+        # # Expand dict
+        # self.sat_to_img_ids_tuples = []
+        # if self.sat_to_img_ids_json_is_train_dict:
+        #     # for i in range(len(self.sat_filtered_json)):
+        #     #     id = self.sat_filtered_json[i]['id']
+        #     #     sat_id = Path(self.sat_paths[id]).stem      # remove file extension
+        #     #     img_ids = self.sat_to_img_ids_json_path[sat_id]["img_ids"]
+        #     for sat_key, sat_sample in self.sat_to_img_ids_json_path.items():
+        #         id = sat_sample["id"] # int(sat_key.split("_")[0]) # sat_sample['id']
+        #         sat_path = sat_sample["sat_path"]
+        #         img_ids = sat_sample["img_ids"]
+        #         # img_locs = sat_sample["target_positions"] # NOTE: Not accurate after resize
+        #         for img_id in img_ids:
+        #             self.sat_to_img_ids_tuples.append((id, sat_path, img_id))
+        #     print("len(self.sat_to_img_ids_json_path): ", len(self.sat_to_img_ids_json_path))
+        #     print("len(self.sat_to_img_ids_tuples): ", len(self.sat_to_img_ids_tuples))
+        # else:
+        #     self.filtered_val_ds_by_tax = [d for d in self.sat_to_img_ids_json_path if self.tax_to_filter_val in d['taxonomy']] 
 
         if mode == 'train':
             self.img_transform = transforms.Compose([

app.py

@@ -11,7 +11,12 @@ This version mirrors the layout of `app_BACKUP.py` but:
 # ────────────────────────── imports ───────────────────────────────────
 from pathlib import Path
 
+# Use non-GUI backend to avoid Tkinter errors in background threads
+import matplotlib
+matplotlib.use("Agg", force=True)
+
 import gradio as gr
+import os, glob, threading, time
 import torch
 from PIL import Image
 
@@ -26,7 +31,7 @@ from Taxabind.TaxaBind.SatBind.clip_seg_tta import ClipSegTTA   # noqa: E402
 
 
 # CHANGE ME!
-currEpisode = 0
+POLL_INTERVAL = 0.1 # For visualization
 
 # Prepare the model
 # device = torch.device('cpu') #if USE_GPU_TRAINING else torch.device('cpu')
@@ -51,7 +56,7 @@ if TAXABIND_TTA:
         sat_to_img_ids_json_path=TAXABIND_SAT_TO_IMG_IDS_JSON_PATH,
         patch_size=TAXABIND_PATCH_SIZE,
         sat_checkpoint_path=TAXABIND_SAT_CHECKPOINT_PATH,
-        sample_index = 0,   # Set using 'reset' in worker
+        sample_index = -1,   # Set using 'reset' in worker
         blur_kernel = TAXABIND_GAUSSIAN_BLUR_KERNEL,
         device=device,
         sat_to_img_ids_json_is_train_dict=False, # for search ds val
@@ -66,23 +71,12 @@ if TAXABIND_TTA:
 else:
     clip_seg_tta = None
 
-# Define TestWorker
-planner = TestWorker(
-    meta_agent_id=0, 
-    n_agent=1, 
-    policy_net=policy_net, 
-    global_step=3, 
-    device='cuda', 
-    greedy=True, 
-    save_image=SAVE_GIFS, 
-    clip_seg_tta=clip_seg_tta
-)
 
 # ────────────────────────── Gradio process fn ─────────────────────────
 
 def process(
-    sat_img: Image.Image | None,
-    ground_img: Image.Image | None,
+    sat_path: str | None,
+    ground_path: str | None,
     taxonomy: str | None = None,
 ):
     """Callback executed when the user presses **Run** in the UI.
@@ -93,17 +87,67 @@ def process(
     retained to conform to the requested interface.
     """
     # If no satellite image is provided we bail out early.
-    if sat_img is None:
-        return {"error": "Please provide a satellite image."}
+    if sat_path is None:
+        return None
+
+    # ------------------------------------------------------------------
+    # Load images from paths and configure ClipSegTTA inputs
+    sat_img = Image.open(sat_path).convert("RGB")
+    ground_img_pil = Image.open(ground_path).convert("RGB") if ground_path else None
+
+    clip_seg_tta.img_paths = [ground_path] if ground_path else []
+    clip_seg_tta.imo_path = sat_path
+    clip_seg_tta.imgs = ([clip_seg_tta.dataset.img_transform(ground_img_pil).to(device)]
+                         if ground_img_pil else [])
+    clip_seg_tta.imo = clip_seg_tta.dataset.imo_transform(sat_img).to(device)
+    clip_seg_tta.sounds = []
+    clip_seg_tta.sound_ids = [] # None
+    clip_seg_tta.species_name = taxonomy or ""
+    clip_seg_tta.target_positions = [(0,0)]  # PLACEHOLDERS
+    clip_seg_tta.gt_mask_name = taxonomy.replace(" ", "_") # None
+
+    # Define TestWorker
+    planner = TestWorker(
+        meta_agent_id=0, 
+        n_agent=1, 
+        policy_net=policy_net, 
+        global_step=-1, 
+        device='cuda', 
+        greedy=True, 
+        save_image=SAVE_GIFS, 
+        clip_seg_tta=clip_seg_tta
+    )
+
+    # ------------------------------------------------------------------
+
+    # 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))
 
     # Optionally you may want to reset episode index or make it configurable.
     # For now we hard-code episode 0, mirroring the snippet.
-    planner.run_episode(currEpisode)
+    t = threading.Thread(target=planner.run_episode, args=(0,), daemon=True)
+    t.start()
+    # planner.run_episode(0)
 
-    print("planner.perf_metrics: ", planner.perf_metrics)
+    sent = set()
+    last_img = None
+    while t.is_alive():
+        # discover any new pngs written by TestWorker
+        pngs = glob.glob(os.path.join(gifs_path, "*.png"))
+        pngs.sort(key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+        for fp in pngs:
+            if fp not in sent:
+                sent.add(fp)
+                last_img = fp
+                yield fp        # stream update
+        time.sleep(POLL_INTERVAL)
 
-    # Return the collected performance metrics so they can be inspected.
-    return planner.perf_metrics
+    # one final yield after the loop finishes
+    yield last_img
+
+    print("planner.perf_metrics: ", planner.perf_metrics)
 
 
 # ────────────────────────── Gradio UI ─────────────────────────────────
@@ -121,7 +165,7 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
             sat_input = gr.Image(
                 label="Satellite Image",
                 sources=["upload"],
-                type="pil",
+                type="filepath",
                 height=320,
             )
             taxonomy_input = gr.Textbox(
@@ -131,14 +175,14 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
             ground_input = gr.Image(
                 label="Ground-level Image (optional)",
                 sources=["upload"],
-                type="pil",
+                type="filepath",
                 height=320,
             )
             run_btn = gr.Button("Run", variant="primary")
 
         with gr.Column():
-            gr.Markdown("### Episode Metrics")
-            metrics_out = gr.JSON(label="Performance Metrics")
+            gr.Markdown("### Live Heatmap")
+            display_img = gr.Image(label="Current Heatmap", type="filepath", height=512)
 
     # Bind callback
 
@@ -175,8 +219,8 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
                 ],
             ],
             inputs=[sat_input, ground_input, taxonomy_input],
-            outputs=[metrics_out],
-            fn=lambda sat, grd, tax: process(sat, grd),
+            outputs=[display_img],
+            fn=lambda sat, grd, tax: process(sat, grd, tax),
             cache_examples=False,
         )
 
@@ -207,8 +251,8 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
                 ],
             ],
             inputs=[sat_input, ground_input, taxonomy_input],
-            outputs=[metrics_out],
-            fn=lambda sat, grd, tax: process(sat, grd),
+            outputs=[display_img],
+            fn=lambda sat, grd, tax: process(sat, grd, tax),
             cache_examples=False,
         )
 
@@ -216,7 +260,7 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
     run_btn.click(
         fn=process,
         inputs=[sat_input, ground_input, taxonomy_input],
-        outputs=metrics_out,
+        outputs=display_img,
     )
 
 # ────────────────────────── unchanged worker initialisation ───────────

env.py

@@ -739,6 +739,78 @@ class Env():
         self.frame_files.append(frame)
         plt.close()
 
+####################
+# ADDED: For app.py
+####################
+
+    def plot_heatmap(self, save_dir, step, travel_dist, robots_route=None):
+        """Plot only the segmentation heatmap and save it as ``{step}.png`` in
+        ``save_dir``. This lightweight helper is meant for asynchronous
+        streaming in the Gradio demo when full `plot_env` is too heavy.
+
+        Parameters
+        ----------
+        save_dir : str
+            Directory to save the generated PNG file.
+        step : int
+            Current timestep; becomes the filename ``{step}.png``.
+        robots_route : list | None
+            Optional list of routes (xPoints, yPoints) to overlay.
+        Returns
+        -------
+        str
+            Full path to the generated PNG file.
+        """
+        import os
+        plt.switch_backend('agg')
+        plt.cla()
+
+        color_list = ["r", "g", "c", "m", "y", "k"]
+        fig, ax = plt.subplots(1, 1, figsize=(6, 6))
+
+        # Select the mask to visualise
+        if TAXABIND_TTA and USE_CLIP_PREDS:
+            side_dim = int(np.sqrt(self.segmentation_info_mask.shape[0]))
+            mask_viz = self.segmentation_info_mask.squeeze().reshape((side_dim, side_dim)).T
+            scale_y = math.ceil(self.ground_truth_size[1] / side_dim)
+            scale_x = math.ceil(self.ground_truth_size[0] / side_dim)
+            upscaled_mask_viz = np.kron(mask_viz, np.ones((scale_y, scale_x)))
+            upscaled_mask_viz = upscaled_mask_viz[:self.ground_truth_size[1], :self.ground_truth_size[0]]
+            im = ax.imshow(upscaled_mask_viz, cmap="viridis")
+            ax.axis("off")
+        else:
+            im = ax.imshow(self.segmentation_mask.mean(axis=-1), cmap='viridis', vmin=0, vmax=100)
+            ax.axis((0, self.ground_truth_size[1], self.ground_truth_size[0], 0))
+
+        # Optionally overlay robot paths
+        if robots_route is not None:
+            for i, route in enumerate(robots_route):
+                robot_marker_color = color_list[i % len(color_list)]
+                xPoints, yPoints = route
+                ax.plot(xPoints, yPoints, c=robot_marker_color, linewidth=2)
+                ax.plot(xPoints[-1], yPoints[-1], markersize=12, zorder=99, marker="^", ls="-", c=robot_marker_color, mec="black")
+
+        # Plot target positions
+        for target in self.target_positions:
+            if self.coverage_belief[target[1], target[0]] == 255:
+                # ax.plot(target[0], target[1], 'go', markersize=8, zorder=99)
+                ax.plot(target[0], target[1], color='g', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99)
+            else:
+                # ax.plot(target[0], target[1], 'ro', markersize=8, zorder=99)
+                ax.plot(target[0], target[1], color='r', marker='x', linestyle='-', markersize=12, markeredgewidth=4, zorder=99)
+
+        # Color bar
+        cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
+        cbar.set_label("Normalized Probs")
+
+        plt.suptitle('Targets Found: {}/{}  Coverage ratio: {:.4g}  Travel Dist: {:.4g}'.format(self.num_targets_found, \
+                                                                                    len(self.target_positions), self.explored_rate, travel_dist))
+
+        os.makedirs(save_dir, exist_ok=True)
+        out_path = os.path.join(save_dir, f"{step}.png")
+        fig.savefig(out_path, dpi=100)
+        plt.close(fig)
+        return out_path
 
 ####################
 

test_multi_robot_worker.py

@@ -142,7 +142,7 @@ class TestWorker:
                     min_patch_size=5,   # smoothing parameter
                     n_smooth_iter=2,    # smoothing parameter
                     ignore_label=-1,
-                    plot=self.save_image,
+                    plot=False,     # NOTE: Set to false since using app.py
                     gifs_dir = gifs_path      
                 )    
                 # Fit & predict (this will also plot the clusters before & after smoothing)
@@ -277,31 +277,35 @@ class TestWorker:
                 if not os.path.exists(gifs_path):
                     os.makedirs(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_env(
-                        self.global_step, 
-                        gifs_path, 
-                        step, 
-                        max(travel_dist_list),
-                        robots_route, 
-                        img_path_override=self.clip_seg_tta.img_paths[0],    # Viz 1st
-                        sat_path_override=self.clip_seg_tta.imo_path,
-                        msk_name_override=self.clip_seg_tta.species_name, 
-                        sound_id_override=sound_id_override,
-                        colormap_mid_val=np.max(self.clip_seg_tta.heatmap_unnormalized_initial)
-                    )
-                else:
-                    self.env.plot_env(
-                        self.global_step, 
-                        gifs_path, 
-                        step, 
-                        max(travel_dist_list),
-                        robots_route,
-                        img_path_override=self.clip_seg_tta.img_paths[0],    # Viz 1st
-                        sat_path_override=self.clip_seg_tta.imo_path,
-                        msk_name_override=self.clip_seg_tta.species_name,
-                        sound_id_override=sound_id_override,
-                    )
+                    self.env.plot_heatmap(gifs_path, step, max(travel_dist_list), robots_route)
+                # if TAXABIND_TTA and USE_CLIP_PREDS:
+                #     self.env.plot_env(
+                #         self.global_step, 
+                #         gifs_path, 
+                #         step, 
+                #         max(travel_dist_list),
+                #         robots_route, 
+                #         img_path_override=self.clip_seg_tta.img_paths[0],    # Viz 1st
+                #         sat_path_override=self.clip_seg_tta.imo_path,
+                #         msk_name_override=self.clip_seg_tta.species_name, 
+                #         sound_id_override=sound_id_override,
+                #         colormap_mid_val=np.max(self.clip_seg_tta.heatmap_unnormalized_initial)
+                #     )
+                # else:
+                #     self.env.plot_env(
+                #         self.global_step, 
+                #         gifs_path, 
+                #         step, 
+                #         max(travel_dist_list),
+                #         robots_route,
+                #         img_path_override=self.clip_seg_tta.img_paths[0],    # Viz 1st
+                #         sat_path_override=self.clip_seg_tta.imo_path,
+                #         msk_name_override=self.clip_seg_tta.species_name,
+                #         sound_id_override=sound_id_override,
+                #     )
 
             if done:
                 break