[BUGFIX] Forces process thread to stop if page is refreshed while operation is done

app.py

@@ -16,6 +16,7 @@ import matplotlib
 matplotlib.use("Agg", force=True)
 
 import gradio as gr
+import ctypes  # for safely stopping background threads
 import os, glob, threading, time
 import torch
 from PIL import Image
@@ -31,6 +32,21 @@ from test_multi_robot_worker import TestWorker  # noqa: E402
 from Taxabind.TaxaBind.SatBind.clip_seg_tta import ClipSegTTA   # noqa: E402
 
 
+# Helper to kill a Python thread by injecting SystemExit
+
+def _stop_thread(thread: threading.Thread):
+    """Forcefully raise SystemExit in the given thread (best-effort)."""
+    if thread is None or not thread.is_alive():
+        return
+    tid = thread.ident
+    if tid is None:
+        return
+    # Ask CPython to raise SystemExit in the thread context
+    res = ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), ctypes.py_object(SystemExit))
+    if res > 1:
+        # If it returned >1, cleanup and fail safe
+        ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_long(tid), None)
+
 # CHANGE ME!
 POLL_INTERVAL = 0.1 # For visualization
 
@@ -182,23 +198,29 @@ def process(
 
     t = threading.Thread(target=planner.run_episode, args=(0,), daemon=True)
     t.start()
-    # planner.run_episode(0)
 
-    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)
-
-    # one final yield after the loop finishes
-    yield last_img
+    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_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)
+    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
 
     print("planner.perf_metrics: ", planner.perf_metrics)