Provide indicator that TTA is being run

Taxabind/TaxaBind/SatBind/clip_seg_tta.py

@@ -150,6 +150,9 @@ class ClipSegTTA:
 
         self.clip_inference_time = 0.0
         self.tta_time = 0.0
+        # NOTE: integration with app.py on hf
+        self.executing_tta = False
+
 
 
     def load_data(self):
@@ -291,6 +294,8 @@ class ClipSegTTA:
         :param viz_heatmap: If True, perform visualization. If False, skip plotting.
         """
 
+        # NOTE: integration with app.py on hf
+        self.executing_tta = True
         # with enable_grad():
 
         ### Option 1: SAMPLE FROM DATASET
@@ -444,6 +449,9 @@ class ClipSegTTA:
         self.model_local.imo_encoder.to(self.device)
         self.model_local.bio_model.to(self.device)
 
+        # NOTE: integration with app.py on hf
+        self.executing_tta = False
+
         return self.heatmap
 
 

app.py

@@ -54,6 +54,9 @@ def _stop_thread(thread: threading.Thread):
 _running_threads: list[threading.Thread] = []
 _running_threads_lock = threading.Lock()
 
+# Map worker threads to their ClipSegTTA instance so UI can read executing_tta flag
+_thread_clip_map: dict[threading.Thread, ClipSegTTA] = {}
+
 def _register_thread(th: threading.Thread):
     """Record a newly started worker thread so we can later cancel it."""
     with _running_threads_lock:
@@ -277,6 +280,7 @@ def process_search_tta(
         """Prepare directory, build planner, run an episode, record errors."""
         try:
             planner = build_planner(enable_tta, save_dir, clip_obj)
+            _thread_clip_map[threading.current_thread()] = planner.clip_seg_tta
             planner.run_episode(0)
         except Exception as exc:
             # Mark that this planner crashed so UI can show an error status
@@ -352,15 +356,23 @@ def process_search_tta(
 
             # Determine status based on whether we already have a frame and whether
             # the corresponding thread is still alive.
-            def _mk_status(last_frame, thread_alive, errored: bool):
+            def _mk_status(last_frame, thread_alive, errored: bool, running_tta: bool=False):
                 if errored:
                     return "Error!"
                 if last_frame is None:
                     return "Initializing model…"
-                return "Running…" if thread_alive else "Done."
-
-            status_tta = _mk_status(last_tta, thread_tta.is_alive(), error_flags["tta"])
-            status_no  = _mk_status(last_no,  thread_no.is_alive(), error_flags["no"])
+                if not thread_alive:
+                    return "Done."
+                return "Running TTA…" if running_tta else "Running Planner…"
+
+            exec_tta_flag = False
+            if thread_tta.is_alive():
+                clip_obj = _thread_clip_map.get(thread_tta)
+                if clip_obj is not None and getattr(clip_obj, "executing_tta", False):
+                    exec_tta_flag = True
+
+            status_tta = _mk_status(last_tta, thread_tta.is_alive(), error_flags["tta"], exec_tta_flag)
+            status_no  = _mk_status(last_no,  thread_no.is_alive(), error_flags["no"], False)
 
             # Determine if we should reveal sliders (once corresponding thread has finished)
             show_slider_tta = (not thread_tta.is_alive()) and (last_tta is not None)