[NEW] Able to launch both gifs concurrently

app.py

@@ -21,6 +21,7 @@ import os, glob, threading, time
 import torch
 from PIL import Image
 import json
+import copy
 import spaces   # integration with ZeroGPU on hf
 
 # Import configuration & RL / TTA utilities -------------------------------------------------
@@ -66,8 +67,8 @@ print('Model loaded!')
 
 # Init Taxabind here (only need to init once)
 if TAXABIND_TTA:
-    # self.clip_seg_tta = None
-    clip_seg_tta = ClipSegTTA(
+    # Instantiate TWO independent ClipSegTTA objects (one per concurrent run)
+    clip_seg_tta_1 = ClipSegTTA(
         img_dir=TAXABIND_IMG_DIR,
         imo_dir=TAXABIND_IMO_DIR,
         json_path=TAXABIND_INAT_JSON_PATH,
@@ -85,9 +86,28 @@ if TAXABIND_TTA:
         sound_checkpoint_path=TAXABIND_SOUND_CHECKPOINT_PATH,
         # sat_filtered_json_path=TAXABIND_FILTERED_INAT_JSON_PATH,
     )
-    print("ClipSegTTA Loaded!")
+    clip_seg_tta_2 = ClipSegTTA(
+        img_dir=TAXABIND_IMG_DIR,
+        imo_dir=TAXABIND_IMO_DIR,
+        json_path=TAXABIND_INAT_JSON_PATH,
+        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 = -1,   # Set using 'reset' in worker
+        blur_kernel = TAXABIND_GAUSSIAN_BLUR_KERNEL,
+        device=device,
+        sat_to_img_ids_json_is_train_dict=False,
+        tax_to_filter_val=QUERY_TAX,
+        load_model=USE_CLIP_PREDS,
+        initial_modality=INITIAL_MODALITY,
+        sound_data_path=TAXABIND_SOUND_DATA_PATH,
+        sound_checkpoint_path=TAXABIND_SOUND_CHECKPOINT_PATH,
+    )
+    print("ClipSegTTA instances loaded!")
+    # Keep original name for single-run mode compatibility
+    clip_seg_tta = clip_seg_tta_1
 else:
-    clip_seg_tta = None
+    clip_seg_tta_1 = clip_seg_tta_2 = clip_seg_tta = None
 
 # Load metadata json
 tgts_metadata_json_path = os.path.join(script_dir, "examples/metadata.json")
@@ -102,132 +122,120 @@ tgts_metadata = json.load(open(tgts_metadata_json_path))
 # object. By defining explicit generator functions (with `yield from`) we ensure
 # `inspect.isgeneratorfunction` evaluates to True and Gradio streams correctly.
 
-def process_with_tta(
-    sat_path: str | None,
-    ground_path: str | None,
-    taxonomy: str | None = None,
-):
-    """Stream search episode **with** TTA enabled while disabling buttons."""
-    # Disable buttons initially (image reset)
-    yield gr.update(interactive=False), gr.update(interactive=False), gr.update(value=None), gr.update(value="Initializing model…")
-
-    last_img = None
-    for img in process(sat_path, ground_path, taxonomy, True):
-        last_img = img
-        yield gr.update(interactive=False), gr.update(interactive=False), img, gr.update(value="Running…")
-
-    # Re-enable buttons at the end
-    yield gr.update(interactive=True), gr.update(interactive=True), last_img, gr.update(value="Done.")
-
-
-
-def process_no_tta(
+# # # integration with ZeroGPU on hf
+# @spaces.GPU
+def process_search_tta(
     sat_path: str | None,
     ground_path: str | None,
     taxonomy: str | None = None,
 ):
-    """Stream search episode **without** TTA enabled while disabling buttons."""
-    yield gr.update(interactive=False), gr.update(interactive=False), gr.update(value=None), gr.update(value="Initializing model…")
-
-    last_img = None
-    for img in process(sat_path, ground_path, taxonomy, False):
-        last_img = img
-        yield gr.update(interactive=False), gr.update(interactive=False), img, gr.update(value="Running…")
-
-    yield gr.update(interactive=True), gr.update(interactive=True), last_img, gr.update(value="Done.")
+    """Run both TTA and non-TTA search episodes concurrently and stream both heat-maps."""
 
+    # Disable Run button and clear image/status outputs
+    yield gr.update(interactive=False), gr.update(value=None), gr.update(value=None), gr.update(value="Initializing model…"), gr.update(value="Initializing model…")
 
-# # integration with ZeroGPU on hf
-# @spaces.GPU
-def process(
-    sat_path: str | None,
-    ground_path: str | None,
-    taxonomy: str | None = None,
-    with_tta: bool = True,
-):
-    """Callback executed when the user presses **Run** in the UI.
-
-    At test-time we simply trigger the RL search episode via
-    ``planner.run_episode`` and return its performance metrics.
-    The image inputs are currently *not* used directly here but are
-    retained to conform to the requested interface.
-    """
-    # If no satellite image is provided we bail out early.
+    # Bail early if satellite image missing
     if sat_path is None:
-        return None
+        yield gr.update(interactive=True), gr.update(value=None), gr.update(value=None), gr.update(value="No satellite image provided."), gr.update(value="")
+        return
 
-    # ------------------------------------------------------------------
-    # Load images from paths and configure ClipSegTTA inputs
+    # Prepare PIL images
     sat_img = Image.open(sat_path).convert("RGB")
     ground_img_pil = Image.open(ground_path).convert("RGB") if ground_path else None
-    tgts = [tuple(tgt) for tgt in tgts_metadata[taxonomy]["target_positions"]]
-
-    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.gt_mask_name = taxonomy.replace(" ", "_") # None
-    clip_seg_tta.target_positions = tgts if tgts != [] else [(0,0)]
-
-    # Define TestWorker
-    planner = TestWorker(
-        meta_agent_id=0, 
-        n_agent=1, 
-        policy_net=policy_net, 
-        global_step=-1, 
-        device=device, 
-        greedy=True, 
-        save_image=SAVE_GIFS, 
-        clip_seg_tta=clip_seg_tta
-    )
-
-    # ------------------------------------------------------------------
 
-    # 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_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
-
-    t = threading.Thread(target=planner.run_episode, args=(0,), daemon=True)
-    t.start()
+    # Lookup target positions metadata (may be empty)
+    tgt_positions = []
+    if taxonomy and taxonomy in tgts_metadata:
+        tgt_positions = [tuple(t) for t in tgts_metadata[taxonomy]["target_positions"]]
+
+    # Helper to build a TestWorker with/without TTA
+    def build_planner(enable_tta: bool, save_dir: str, clip_obj):
+        local_clip = clip_obj  # re-use the pre-instantiated ClipSegTTA
+        if local_clip is not None:
+            # Feed inputs to ClipSegTTA copy
+            local_clip.img_paths = [ground_path] if ground_path else []
+            local_clip.imo_path = sat_path
+            local_clip.imgs = ([local_clip.dataset.img_transform(ground_img_pil).to(device)] if ground_img_pil else [])
+            local_clip.imo = local_clip.dataset.imo_transform(sat_img).to(device)
+            local_clip.sounds = []
+            local_clip.sound_ids = []
+            local_clip.species_name = taxonomy or ""
+            local_clip.gt_mask_name = taxonomy.replace(" ", "_") if taxonomy else ""
+            local_clip.target_positions = tgt_positions if tgt_positions else [(0, 0)]
+
+        planner = TestWorker(
+            meta_agent_id=0,
+            n_agent=1,
+            policy_net=policy_net,
+            global_step=-1,
+            device=device,
+            greedy=True,
+            save_image=SAVE_GIFS,
+            clip_seg_tta=local_clip,
+        )
+        planner.execute_tta = enable_tta
+        planner.gifs_path = save_dir
+        return planner
+
+    # Prepare gif directories
+    gifs_dir_tta = os.path.join(gifs_path, "with_tta")
+    gifs_dir_no = os.path.join(gifs_path, "no_tta")
+    for d in (gifs_dir_tta, gifs_dir_no):
+        os.makedirs(d, exist_ok=True)
+        # Clean previous pngs
+        for f in os.listdir(d):
+            # if f.endswith(".png"):
+            os.remove(os.path.join(d, f))
+
+    planner_tta = build_planner(True, gifs_dir_tta, clip_seg_tta_1)
+    planner_no  = build_planner(False, gifs_dir_no, clip_seg_tta_2)
+
+    # Launch both planners in background threads
+    thread_tta = threading.Thread(target=planner_tta.run_episode, args=(0,), daemon=True)
+    thread_no  = threading.Thread(target=planner_no.run_episode,  args=(0,), daemon=True)
+    thread_tta.start()
+    thread_no.start()
+
+    sent_tta: set[str] = set()
+    sent_no:  set[str] = set()
+    last_tta = None
+    last_no  = None
 
-    sent: set[str] = set()
-    last_img: str | None = None
     try:
-        while t.is_alive():
-            # discover any new pngs written by TestWorker
-            pngs = glob.glob(os.path.join(gifs_save_dir, "*.png"))
+        while thread_tta.is_alive() or thread_no.is_alive():
+            updated = False
+            # Collect new frames from TTA dir
+            pngs = glob.glob(os.path.join(gifs_dir_tta, "*.png"))
+            pngs.sort(key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
+            for fp in pngs:
+                if fp not in sent_tta:
+                    sent_tta.add(fp)
+                    last_tta = fp
+                    updated = True
+            # Collect new frames from no-TTA dir
+            pngs = glob.glob(os.path.join(gifs_dir_no, "*.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
+                if fp not in sent_no:
+                    sent_no.add(fp)
+                    last_no = fp
+                    updated = True
+
+            if updated:
+                status_tta = "Running…" if thread_tta.is_alive() else "Done."
+                status_no  = "Running…" if thread_no.is_alive()  else "Done."
+                yield gr.update(interactive=False), last_tta, last_no, gr.update(value=status_tta), gr.update(value=status_no)
+
             time.sleep(POLL_INTERVAL)
     finally:
-        # This block runs when the generator is cancelled (e.g. page refresh)
-        if t.is_alive():
-            _stop_thread(t)
-            t.join(timeout=1)
-
-    # If the episode finished naturally, send the last frame once more
-    if last_img is not None:
-        yield last_img
+        # Ensure background threads are stopped on cancel
+        for th in (thread_tta, thread_no):
+            if th.is_alive():
+                _stop_thread(th)
+                th.join(timeout=1)
 
-    print("planner.perf_metrics: ", planner.perf_metrics)
+    # Final emit after both finish
+    yield gr.update(interactive=True), last_tta, last_no, gr.update(value="Done."), gr.update(value="Done.")
 
 
 # ────────────────────────── Gradio UI ─────────────────────────────────
@@ -280,15 +288,14 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
                 type="filepath",
                 height=320,
             )
-            run_tta_btn = gr.Button("Run (with TTA)", variant="primary")
-            run_no_tta_btn = gr.Button("Run (without TTA)", variant="secondary")
+            run_btn = gr.Button("Run Search-TTA", variant="primary")
 
         with gr.Column():
             gr.Markdown("### Live Heatmap (with TTA)")
-            display_img_tta = gr.Image(label="Heatmap (TTA)", type="filepath", height=512)
+            display_img_tta = gr.Image(label="Heatmap (TTA)", type="filepath", height=400)  # 512
             status_tta = gr.Markdown("")
             gr.Markdown("### Live Heatmap (without TTA)")
-            display_img_no_tta = gr.Image(label="Heatmap (no TTA)", type="filepath", height=512)
+            display_img_no_tta = gr.Image(label="Heatmap (no TTA)", type="filepath", height=400)  # 512
             status_no_tta = gr.Markdown("")
 
     # Bind callback
@@ -321,21 +328,16 @@ with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
                 ],
             ],
             inputs=[sat_input, ground_input, taxonomy_input],
-            outputs=[run_tta_btn, run_no_tta_btn, display_img_tta, status_tta],
-            fn=process_with_tta,
+            outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta],
+            fn=process_search_tta,
             cache_examples=False,
         )
 
 
-    run_tta_btn.click(
-        fn=process_with_tta,
-        inputs=[sat_input, ground_input, taxonomy_input],
-        outputs=[run_tta_btn, run_no_tta_btn, display_img_tta, status_tta],
-    )
-    run_no_tta_btn.click(
-        fn=process_no_tta,
+    run_btn.click(
+        fn=process_search_tta,
         inputs=[sat_input, ground_input, taxonomy_input],
-        outputs=[run_tta_btn, run_no_tta_btn, display_img_no_tta, status_no_tta],
+        outputs=[run_btn, display_img_tta, display_img_no_tta, status_tta, status_no_tta],
     )
 
     # Footer to point out to model and data from app page.

env.py

@@ -816,8 +816,9 @@ class Env():
         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))
+        # Change coverage to 1dp
+        plt.suptitle('Targets Found: {}/{}  Coverage: {:.1f}%  Steps: {}/{}'.format(self.num_targets_found, \
+                                                                                    len(self.target_positions), self.explored_rate*100, step, NUM_EPS_STEPS))
 
         os.makedirs(save_dir, exist_ok=True)
         out_path = os.path.join(save_dir, f"{step}.png")