"""
Simplified Gradio demo for Search-TTA evaluation.
This version mirrors the layout of `app_BACKUP.py` but:
1. Loads no OpenCLIP / CLAP / Satellite encoders at import-time.
2. Keeps only the Satellite and Ground-level image inputs.
3. Exposes the high-level wrapper classes `ClipSegTTA` and
`TestWorker` and calls `TestWorker.run_episode` inside the
`process` callback.
"""
# ────────────────────────── 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 ctypes # for safely stopping background threads
import os, glob, threading, time
import torch
from PIL import Image
import json
import copy
import shutil
import spaces # integration with ZeroGPU on hf
# Import configuration & RL / TTA utilities -------------------------------------------------
# NOTE: we import * so that the global names (e.g. USE_GPU, MODEL_NAME, etc.)
# are available exactly as referenced later in the unchanged snippet.
from test_parameter import * # noqa: F403, F401 (wild-import is intentional here)
from model import PolicyNet # noqa: E402 – after wild import on purpose
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)
# ──────────── Thread Registry for Cleanup on Tab Switch ─────────────
_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:
_running_threads.append(th)
def _kill_running_threads():
"""Stop all worker threads that are still alive."""
with _running_threads_lock:
for t in list(_running_threads):
_stop_thread(t)
# Clear list regardless of alive status
_running_threads.clear()
# ──────────── Run directory rotation ─────────────
RUN_HISTORY_LIMIT = 30 # keep at most this many timestamped run directories per instance
def _prune_old_run_dirs(base_dir: str, limit: int = RUN_HISTORY_LIMIT):
"""Delete oldest timestamp-named run directories leaving only *limit* of the newest ones."""
try:
dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))]
# Timestamp format YYYYmmdd_HHMMSS ensures lexicographic order == chronological order
dirs.sort()
if len(dirs) > limit:
for obsolete in dirs[:-limit]:
shutil.rmtree(os.path.join(base_dir, obsolete), ignore_errors=True)
except Exception:
# Best-effort; ignore cleanup errors
pass
# CHANGE ME!
POLL_INTERVAL = 1.0 # For visualization
# Prepare the model
# device = torch.device('cpu') #if USE_GPU_TRAINING else torch.device('cpu')
device = torch.device('cuda') if USE_GPU and torch.cuda.is_available() else torch.device('cpu')
policy_net = PolicyNet(INPUT_DIM, EMBEDDING_DIM).to(device)
# script_dir = os.path.dirname(os.path.abspath(__file__))
script_dir = Path(__file__).resolve().parent
print("real_script_dir: ", script_dir)
# checkpoint = torch.load(f'{script_dir}/modules/vlm_search/{model_path}/{MODEL_NAME}')
checkpoint = torch.load(f'{model_path}/{MODEL_NAME}')
policy_net.load_state_dict(checkpoint['policy_model'])
print('Model loaded!')
# print(next(policy_net.parameters()).device)
# # (ClipSegTTA will now be instantiated lazily inside each planner thread)
# clip_seg_tta_1 = clip_seg_tta_2 = None # placeholder; real instances created per thread
# if False and TAXABIND_TTA:
# # 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,
# 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, # device,
# sat_to_img_ids_json_is_train_dict=False, # for search ds val
# 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,
# # sat_filtered_json_path=TAXABIND_FILTERED_INAT_JSON_PATH,
# )
# 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,
# )
# Load metadata json
tgts_metadata_json_path = os.path.join(script_dir, "examples/metadata.json")
tgts_metadata = json.load(open(tgts_metadata_json_path))
# ────────────────────────── Gradio process fn ─────────────────────────
# Helper wrappers so that Gradio recognises streaming (generator) functions
# NOTE: A lambda that *returns* a generator is NOT itself a generator *function*,
# hence Gradio fails to detect streaming and treats the return value as a plain
# object. By defining explicit generator functions (with `yield from`) we ensure
# `inspect.isgeneratorfunction` evaluates to True and Gradio streams correctly.
# # # integration with ZeroGPU on hf
# @spaces.GPU
def process_search_tta(
sat_path: str | None,
ground_path: str | None,
taxonomy: str | None = None,
session_threads: list[threading.Thread] | None = None,
):
"""Run both TTA and non-TTA search episodes concurrently and stream both heat-maps."""
if session_threads is None:
session_threads = []
# Disable Run button and clear image/status outputs, hide sliders, clear frame states
yield (
gr.update(interactive=False),
gr.update(value=None),
gr.update(value=None),
gr.update(value="Initializing model…", visible=True),
gr.update(value="Initializing model…", visible=True),
gr.update(visible=False),
gr.update(visible=False),
[],
[],
session_threads,
)
# Bail early if satellite image missing
if sat_path is None:
yield (
gr.update(interactive=True),
gr.update(value=None),
gr.update(value=None),
gr.update(value="No satellite image provided.", visible=True),
gr.update(value="", visible=True),
gr.update(visible=False),
gr.update(visible=False),
[],
[],
session_threads,
)
return
# 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
# 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):
# Lazily (re)create a ClipSegTTA instance per thread if not provided
local_clip = clip_obj
if TAXABIND_TTA and local_clip is None:
local_clip = 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,
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,
)
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
# ────────────── Per-run output directories ──────────────
# Ensure base directory exists
os.makedirs(gifs_path, exist_ok=True)
run_id = time.strftime("%Y%m%d_%H%M%S") # unique timestamp
run_root = os.path.join(gifs_path, run_id)
gifs_dir_tta = os.path.join(run_root, "with_tta")
gifs_dir_no = os.path.join(run_root, "no_tta")
os.makedirs(gifs_dir_tta, exist_ok=True)
os.makedirs(gifs_dir_no, exist_ok=True)
# House-keep old runs so we never keep more than RUN_HISTORY_LIMIT
_prune_old_run_dirs(gifs_path, RUN_HISTORY_LIMIT)
# Shared dict to record if a thread hit an exception
error_flags = {"tta": False, "no": False}
def _planner_thread(enable_tta: bool, save_dir: str, clip_obj, key: str):
"""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
error_flags[key] = True
# Log full traceback so developers can debug via console logs
import traceback, sys
traceback.print_exc()
# Still exit the thread
return
# Launch both planners in background threads – preparation included
thread_tta = threading.Thread(
target=_planner_thread,
args=(True, gifs_dir_tta, None, "tta"),
daemon=True,
)
thread_no = threading.Thread(
target=_planner_thread,
args=(False, gifs_dir_no, None, "no"),
daemon=True,
)
# Track threads for this user session
session_threads.extend([thread_tta, thread_no])
thread_tta.start()
thread_no.start()
sent_tta: set[str] = set()
sent_no: set[str] = set()
last_tta = None
last_no = None
# Track previous status strings so we can emit updates when only the
# status (Running…/Done.) changes even if no new frame was produced.
# Previous status values so we can detect changes and yield updates
prev_status_tta = "Initializing model…"
prev_status_no = "Initializing model…"
try:
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:
# Ensure file is fully written (non-empty & readable)
try:
if os.path.getsize(fp) == 0:
continue
with open(fp, "rb") as fh:
fh.read(1)
except Exception:
# Skip this round; we'll retry next poll
continue
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_no:
try:
if os.path.getsize(fp) == 0:
continue
with open(fp, "rb") as fh:
fh.read(1)
except Exception:
continue
sent_no.add(fp)
last_no = fp
updated = True
# 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, running_tta: bool=False):
if errored:
return "Error!"
if last_frame is None:
return "Initializing model…"
if not thread_alive:
return "Done."
return "Executing TTA (Scheduling GPUs)…" if running_tta else "Executing 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)
show_slider_no = (not thread_no.is_alive()) and (last_no is not None)
# Build slider updates
slider_tta_upd = gr.update()
slider_no_upd = gr.update()
frames_tta_upd = gr.update()
frames_no_upd = gr.update()
if show_slider_tta:
n_tta_frames = max(len(sent_tta), 1)
slider_tta_upd = gr.update(visible=True, minimum=1, maximum=n_tta_frames, value=n_tta_frames)
frames_tta_upd = sorted(sent_tta, key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
if show_slider_no:
n_no_frames = max(len(sent_no), 1)
slider_no_upd = gr.update(visible=True, minimum=1, maximum=n_no_frames, value=n_no_frames)
frames_no_upd = sorted(sent_no, key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
# Emit update if we have a new frame OR status changed OR slider visibility changed
if (
updated
or status_tta != prev_status_tta
or status_no != prev_status_no
or show_slider_tta
or show_slider_no
):
yield (
gr.update(interactive=False),
last_tta,
last_no,
gr.update(value=status_tta, visible=True),
gr.update(value=status_no, visible=True),
slider_tta_upd,
slider_no_upd,
frames_tta_upd,
frames_no_upd,
session_threads,
)
prev_status_tta = status_tta
prev_status_no = status_no
time.sleep(POLL_INTERVAL)
finally:
# 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)
# Remove finished threads from global registry
with _running_threads_lock:
# Clear session thread list
session_threads.clear()
# Small delay to ensure last frame files are fully flushed
time.sleep(0.2)
# One last scan after both threads have finished to catch any frame
# that may have been written just before termination but after the last
# polling iteration.
for fp in sorted(glob.glob(os.path.join(gifs_dir_tta, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0])):
if fp not in sent_tta:
sent_tta.add(fp)
last_tta = fp
for fp in sorted(glob.glob(os.path.join(gifs_dir_no, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0])):
if fp not in sent_no:
sent_no.add(fp)
last_no = fp
# Prepare frames list and slider configs
frames_tta = sorted(glob.glob(os.path.join(gifs_dir_tta, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
frames_no = sorted(glob.glob(os.path.join(gifs_dir_no, "*.png")), key=lambda p: int(os.path.splitext(os.path.basename(p))[0]))
if last_tta is None and frames_tta:
last_tta = frames_tta[-1]
if last_no is None and frames_no:
last_no = frames_no[-1]
n_tta = len(frames_tta) or 1 # prevent zero-range slider
n_no = len(frames_no) or 1
# Final emit: re-enable button, hide statuses, show sliders set to last frame
yield (
gr.update(interactive=True),
last_tta,
last_no,
gr.update(visible=False),
gr.update(visible=False),
gr.update(visible=True, minimum=1, maximum=n_tta, value=n_tta),
gr.update(visible=True, minimum=1, maximum=n_no, value=n_no),
frames_tta,
frames_no,
session_threads,
)
# ────────────────────────── Gradio UI ─────────────────────────────────
with gr.Blocks(title="Search-TTA (Simplified)", theme=gr.themes.Base()) as demo:
gr.Markdown(
"""
# Search-TTA: A Multimodal Test-Time Adaptation Framework for Visual Search in the Wild Demo
Click on any of the examples below and run the TTA demo. Check out the multimodal heatmap generation feature by switching to the next tab above.
Note that the model initialization, RL planner, and TTA updates are not fully optimized on GPU for this huggingface demo, and hence may experience some lag during execution.
If you encounter an 'Error' status, refresh the browser and rerun the demo, or try again the next day. We will improve this in the future.
Project Website
"""
)
# gr.Markdown(
# """
#