app.py

import gradio as gr
from gradio_bbox_annotator import BBoxAnnotator
from PIL import Image
import numpy as np
import torch
import os
import shutil
import time
import json
import uuid
from pathlib import Path
import tempfile
import zipfile
from skimage import measure
from matplotlib import cm
from glob import glob
from natsort import natsorted
from huggingface_hub import HfApi, upload_file
import spaces

from inference_seg import load_model as load_seg_model, run as run_seg
from inference_count import load_model as load_count_model, run as run_count
from inference_track import load_model as load_track_model, run as run_track
from _utils.image_io import (
    standardize_image, inspect_image, image_size, array_nbytes, pixel_stats,
    RECOMMENDED_SIZE, WARN_SIZE, MAX_SIZE, MIN_SIZE, MAX_READ_BYTES, TIFF_EXTENSIONS,
)

HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = "VisionLanguageGroup/feedback"


print("===== clearing cache =====")
cache_path = os.path.expanduser("~/.cache/huggingface/gradio")
if os.path.exists(cache_path):
    try:
        shutil.rmtree(cache_path)
        print("✅ Deleted ~/.cache/huggingface/gradio")
    except:
        pass

SEG_MODEL = None
SEG_DEVICE = torch.device("cpu")

COUNT_MODEL = None
COUNT_DEVICE = torch.device("cpu")

TRACK_MODEL = None
TRACK_DEVICE = torch.device("cpu")

def load_all_models():
    global SEG_MODEL, SEG_DEVICE
    global COUNT_MODEL, COUNT_DEVICE
    global TRACK_MODEL, TRACK_DEVICE
    
    print("\n" + "="*60)
    print("📦 Loading Segmentation Model")
    print("="*60)
    SEG_MODEL, SEG_DEVICE = load_seg_model(use_box=False)
    
    print("\n" + "="*60)
    print("📦 Loading Counting Model")
    print("="*60)
    COUNT_MODEL, COUNT_DEVICE = load_count_model(use_box=False)
    
    print("\n" + "="*60)
    print("📦 Loading Tracking Model")
    print("="*60)
    TRACK_MODEL, TRACK_DEVICE = load_track_model(use_box=False)
    
    print("\n" + "="*60)
    print("✅ All Models Loaded Successfully")
    print("="*60)

load_all_models()

DATASET_DIR = Path("solver_cache")
DATASET_DIR.mkdir(parents=True, exist_ok=True)

def save_feedback_to_hf(query_id, feedback_type, feedback_text=None, img_path=None, bboxes=None):
    """Save feedback to Hugging Face Dataset"""
    
    if not HF_TOKEN:
        print("⚠️ No HF_TOKEN found, using local storage")
        save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)
        return
    
    # Shared by the record and its image so the two pair up; query_id alone is not unique (a session can submit more than once).
    stem = f"feedback_{query_id}_{int(time.time())}"

    try:
        api = HfApi()

        image_in_repo = None
        if img_path and os.path.exists(img_path):
            try:
                image_in_repo = f"images/{stem}.png"
                api.upload_file(
                    path_or_fileobj=img_path,
                    path_in_repo=image_in_repo,
                    repo_id=DATASET_REPO,
                    repo_type="dataset",
                    token=HF_TOKEN
                )
            except Exception as e:
                print(f"⚠️ Failed to upload image: {e}")
                image_in_repo = None  

        feedback_data = {
            "query_id": query_id,
            "feedback_type": feedback_type,
            "feedback_text": feedback_text,
            "image_path": image_in_repo, 
            "bboxes": str(bboxes), 
            "datetime": time.strftime("%Y-%m-%d %H:%M:%S"),
            "timestamp": time.time()
        }

        filename = f"{stem}.json"

        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(feedback_data, f, indent=2, ensure_ascii=False)

        api.upload_file(
            path_or_fileobj=filename,
            path_in_repo=f"data/{filename}",
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )

        os.remove(filename)

        print(f"✅ Feedback saved to HF Dataset: {DATASET_REPO} ({stem})")

    except Exception as e:
        print(f"⚠️ Failed to save to HF Dataset: {e}")
        save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)


def save_feedback(query_id, feedback_type, feedback_text=None, img_path=None, bboxes=None):
    """Save feedback to local JSON file"""
    feedback_data = {
        "query_id": query_id,
        "feedback_type": feedback_type,
        "feedback_text": feedback_text,
        "image": img_path,
        "bboxes": bboxes,
        "datetime": time.strftime("%Y%m%d_%H%M%S")
    }
    feedback_file = DATASET_DIR / query_id / "feedback.json"
    feedback_file.parent.mkdir(parents=True, exist_ok=True)
    
    if feedback_file.exists():
        with feedback_file.open("r") as f:
            existing = json.load(f)
            if not isinstance(existing, list):
                existing = [existing]
            existing.append(feedback_data)
            feedback_data = existing
    else:
        feedback_data = [feedback_data]
    
    with feedback_file.open("w") as f:
        json.dump(feedback_data, f, indent=4, ensure_ascii=False)

def parse_first_bbox(bboxes):
    """Parse the first bounding box from the annotation input, supports dict or list format"""
    if not bboxes:
        return None
    b = bboxes[0]
    if isinstance(b, dict):
        x, y = float(b.get("x", 0)), float(b.get("y", 0))
        w, h = float(b.get("width", 0)), float(b.get("height", 0))
        return x, y, x + w, y + h
    if isinstance(b, (list, tuple)) and len(b) >= 4:
        return float(b[0]), float(b[1]), float(b[2]), float(b[3])
    return None

def parse_bboxes(bboxes):
    """Parse all bounding boxes from the annotation input"""
    if not bboxes:
        return None
    
    result = []
    for b in bboxes:
        if isinstance(b, dict):
            x, y = float(b.get("x", 0)), float(b.get("y", 0))
            w, h = float(b.get("width", 0)), float(b.get("height", 0))
            result.append([x, y, x + w, y + h])
        elif isinstance(b, (list, tuple)) and len(b) >= 4:
            result.append([float(b[0]), float(b[1]), float(b[2]), float(b[3])])
    
    return result

def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
    """Convert a 2D mask of instance IDs to a color image for visualization."""
    def hsv_to_rgb(h, s, v):
        i = int(h * 6.0)
        f = h * 6.0 - i
        i = i % 6
        p = v * (1 - s)
        q = v * (1 - f * s)
        t = v * (1 - (1 - f) * s)
        if i == 0: r, g, b = v, t, p
        elif i == 1: r, g, b = q, v, p
        elif i == 2: r, g, b = p, v, t
        elif i == 3: r, g, b = p, q, v
        elif i == 4: r, g, b = t, p, v
        else: r, g, b = v, p, q
        return int(r * 255), int(g * 255), int(b * 255)

    palette = [(0, 0, 0)]
    for i in range(1, num_colors):
        h = (i % num_colors) / float(num_colors)
        palette.append(hsv_to_rgb(h, 1.0, 0.95))

    palette_arr = np.array(palette, dtype=np.uint8)
    color_idx = mask % num_colors
    return palette_arr[color_idx]


def render_seg_overlay(img_np, inst_mask, overlay_alpha):
    """Render segmentation overlay from cached image/mask."""
    if img_np is None or inst_mask is None:
        return None

    overlay = img_np.copy()
    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))

    for inst_id in np.unique(inst_mask):
        if inst_id == 0:
            continue
        binary_mask = (inst_mask == inst_id).astype(np.uint8)
        color = get_well_spaced_color(inst_id)
        overlay[binary_mask == 1] = (1 - alpha) * overlay[binary_mask == 1] + alpha * color

        contours = measure.find_contours(binary_mask, 0.5)
        for contour in contours:
            contour = contour.astype(np.int32)
            valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
            valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
            overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]

    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
    return Image.fromarray(overlay)


def render_count_overlay(img_np, density_normalized, overlay_alpha):
    """Render counting heatmap overlay from cached image/density."""
    if img_np is None or density_normalized is None:
        return None

    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))
    cmap = cm.get_cmap("jet")
    density_colored = cmap(density_normalized)[:, :, :3]

    overlay = img_np.copy()
    threshold = 0.01
    significant_mask = density_normalized > threshold
    overlay[significant_mask] = (1 - alpha) * overlay[significant_mask] + alpha * density_colored[significant_mask]
    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
    return Image.fromarray(overlay)


def update_seg_overlay_alpha(overlay_alpha, seg_vis_cache):
    """Live update segmentation visualization without rerunning inference."""
    if not seg_vis_cache:
        return None
    return render_seg_overlay(seg_vis_cache.get("img_np"), seg_vis_cache.get("inst_mask"), overlay_alpha)


def update_count_overlay_alpha(overlay_alpha, count_vis_cache):
    """Live update counting visualization without rerunning inference."""
    if not count_vis_cache:
        return None
    return render_count_overlay(count_vis_cache.get("img_np"), count_vis_cache.get("density_normalized"), overlay_alpha)


def update_tracking_overlay_alpha(overlay_alpha, track_vis_cache):
    """Regenerate tracking visualization at new opacity using cached outputs."""
    if not track_vis_cache:
        return None

    tif_dir = track_vis_cache.get("tif_dir")
    output_dir = track_vis_cache.get("output_dir")
    valid_tif_files = track_vis_cache.get("valid_tif_files")
    if not tif_dir or not output_dir or not valid_tif_files:
        return None

    try:
        return create_tracking_visualization(
            tif_dir=tif_dir,
            output_dir=output_dir,
            valid_tif_files=valid_tif_files,
            overlay_alpha=overlay_alpha
        )
    except Exception as e:
        print(f"⚠️ Failed to update tracking opacity: {e}")
        return None


def cleanup_tracking_cache(track_vis_cache):
    """Delete cached tracking temp directories from the previous run."""
    if not track_vis_cache:
        return
    for key in ["input_temp_dir", "output_dir"]:
        path = track_vis_cache.get(key)
        if path and os.path.isdir(path):
            try:
                shutil.rmtree(path)
            except Exception:
                pass


# Per session. Refused rather than evicting the oldest, since there is no way to
# remove a single entry.
MAX_GALLERY_UPLOADS = 10

_GALLERY_RAW_SOURCE = {}
_RAW_SOURCE_CAP = 500


def _remember_gallery_source(thumb_path, raw_path):
    """Map a gallery thumbnail back to the original file it was made from.
    """
    _GALLERY_RAW_SOURCE[thumb_path] = raw_path
    while len(_GALLERY_RAW_SOURCE) > _RAW_SOURCE_CAP:  # dicts keep insertion order
        _GALLERY_RAW_SOURCE.pop(next(iter(_GALLERY_RAW_SOURCE)))


def _annot_path(annot_value):
    """Extract the image path from a BBoxAnnotator value (path or (path, boxes))."""
    if not annot_value:
        return None
    if isinstance(annot_value, (list, tuple)):
        return annot_value[0] if len(annot_value) > 0 else None
    return annot_value


def _human_bytes(n):
    """Format a byte count as MB or GB, whichever reads better."""
    gb = n / 1024 ** 3
    return f"{gb:.1f} GB" if gb >= 1 else f"{n / 1024 ** 2:.0f} MB"


def check_image(img_path):
    """Validate an uploaded file, cheapest checks first.

    Returns a rejection message if the file cannot be used, otherwise None.
    Reject: empty file; unreadable / unsupported format; dimensions over
    MAX_SIZE; full array over MAX_READ_BYTES (a small time-lapse can still be
    huge); corrupt pixels; non-finite (NaN/inf) pixels. Advisory ``gr.Warning``
    (does not reject): larger than WARN_SIZE; smaller than MIN_SIZE; blank
    (uniform) image; very narrow dynamic range.

    Only the header is touched until the size guards pass; pixel checks decode
    image data afterwards, when doing so is bounded by the memory guard.
    """
    if os.path.exists(img_path) and os.path.getsize(img_path) == 0:
        return "This file is empty (0 bytes). Please upload a valid image file."

    w, h = image_size(img_path)
    if w <= 0 or h <= 0:
        # Neither reader could parse the header, so say so rather than failing
        # silently further down.
        ext = os.path.splitext(img_path)[1].lower() or "(no extension)"
        return (f"Could not read {ext} as an image. Supported formats: "
                f"TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, and other formats supported by tifffile. "
                f"Please export to TIFF/PNG/JPG (e.g. from ImageJ/Fiji) first.")

    if w > MAX_SIZE or h > MAX_SIZE:
        return (f"Image is {w}×{h}, larger than the {MAX_SIZE}×{MAX_SIZE} limit. "
                f"Please crop or downsample it before uploading.")

    nbytes = array_nbytes(img_path)
    if nbytes > MAX_READ_BYTES:
        info = inspect_image(img_path)
        detail = f"{w}×{h}"
        if info.frames > 1:
            detail += f" × {info.frames} frames"
        if info.channels > 1:
            detail += f" × {info.channels} channels"
        return (f"This file needs {_human_bytes(nbytes)} of memory to open "
                f"({detail}), over the {_human_bytes(MAX_READ_BYTES)} limit. "
                f"Please crop it, or split out the frames you need.")

    if w > WARN_SIZE or h > WARN_SIZE:
        gr.Warning(f"Image is {w}×{h} and will be resized to "
                   f"{RECOMMENDED_SIZE}×{RECOMMENDED_SIZE}, so fine detail could be lost. "
                   f"You can try cropping the image for better results.",
                   duration=None, title="⚠️ Large image")

    if 0 < min(w, h) < MIN_SIZE:
        gr.Warning(f"Image is only {w}×{h} and will be upscaled to "
                   f"{RECOMMENDED_SIZE}×{RECOMMENDED_SIZE}, so results may be unreliable. "
                   f"A larger image may work better.",
                   duration=None, title="⚠️ Very small image")

    # Size guards passed, so pixel decoding is bounded. Check problems that
    # the header cannot reveal.
    rep = pixel_stats(img_path)
    if not rep.decoded:
        return ("This image appears to be corrupted or truncated and could not be read. "
                "Please re-export it (e.g. from ImageJ/Fiji) and try again.")
    if not rep.finite:
        return ("This image contains invalid pixel values (NaN or infinity). "
                "Please clean or re-export it before uploading.")
    if rep.vmin == rep.vmax:
        if rep.vmax == 0:
            kind = "completely black"
        elif rep.dtype_max and rep.vmax >= rep.dtype_max:
            kind = "completely white"
        else:
            kind = "a single uniform value"
        gr.Warning(f"This image is {kind}, it has no visible content, so results will not be meaningful.",
                   duration=None, title="⚠️ Blank image")
    elif rep.low_range:
        gr.Warning("This image may use a very narrow intensity range (low contrast). Signal may be too weak for reliable results; consider adjusting acquisition or contrast before uploading.",
                   duration=None, title="⚠️ Low dynamic range")
    return None


def _describe_read(info):
    """Plain-language summary of how a file's dimensions were interpreted.

    Deliberately avoids array jargon ("axis", "size-4"): the reader thinks in
    channels / z-slices / timepoints, not numpy axes.
    """
    # lines = [f"Detected shape: {info.shape}"]
    if info.frames >1:
        lines = [f"Detected {info.channels}-channel {info.frames}-frame image stack of size {info.width}×{info.height}."]
    else:
        lines = [f"Detected {info.channels}-channel image of size {info.width}×{info.height}."]
    if info.channels >= 3:
        lines.append(f"Loaded as 3-channel RGB image.")
    elif info.channels == 2:
        lines.append(f"Loaded as 2-channel image (red, green).")
    else:
        lines.append(f"Loaded as single-channel grayscale image.")

    if info.frames > 1:
        lines.append(f"Showing first frame by default (you can use the stack frame slider to pick another).")
        if info.sub_frames > 1:
            lines.append(f"Showing the first of {info.sub_frames} {info.sub_label}s by default (use the {info.sub_label} slider to pick another).")

    lines.append(
        "If wrong, please convert your image to a 8-bit RGB/Grayscale image file before uploading (there are available tools such as ImageJ/Fiji)."
    )
    return "<br>".join(lines)


def prepare_uploaded_image(annot_value):
    """On annotator upload: standardize frame 0 for preview and configure the
    stack sliders.

    Browsers cannot render 16-bit / float / multi-page TIFFs, so we standardize
    to an 8-bit RGB PNG. If the file is a stack, notify the user and reveal a
    frame slider (default frame 1); a file with a second stack axis (e.g. Z in a
    time+Z series) also gets a z-plane slider (0 = max-intensity projection).
    Sliders that don't apply stay hidden.

    Returns (annotator_value, raw_path, frame_slider_update, zplane_slider_update).
    """
    hidden = gr.update(visible=False, value=1)
    img_path = _annot_path(annot_value)
    if not img_path:
        return annot_value, None, hidden, hidden

    # Clear the annotator on refusal, so a rejected file cannot reach inference.
    rejected = check_image(img_path)
    if rejected:
        gr.Warning(rejected, duration=None, title="❌ Cannot use this file")
        return None, None, hidden, hidden

    info = inspect_image(img_path)
    display = standardize_image(img_path, frame=0,
                                sub_frame=0 if info.sub_frames > 1 else None)

    # Stay quiet unless something non-obvious happened; a plain RGB or grayscale
    # image needs no explanation.
    if info.guessed or info.frames > 1 or info.channels > 3:
        gr.Info(_describe_read(info), duration=None, title="📚 Image Loading Info")

    slider = (gr.update(visible=True, maximum=info.frames, value=1) if info.frames > 1
              else hidden)
    zslider = (gr.update(visible=True, maximum=info.sub_frames, value=1,
                         label=f"🔬 {info.sub_label.capitalize()} (choose plane to use)")
               if info.sub_frames > 1 else hidden)
    return display, img_path, slider, zslider


def select_frame(raw_path, frame_num, zplane=1):
    """Re-render the annotator preview for the chosen frame and z-plane (both
    1-based; z-plane is ignored for files with only one stack axis)."""
    if not raw_path:
        return gr.update()
    return standardize_image(raw_path, frame=int(frame_num) - 1, sub_frame=int(zplane) - 1)


@spaces.GPU
def segment_with_choice(use_box_choice, annot_value, overlay_alpha):
    """Segmentation handler - supports bounding box, returns colorized overlay and original mask path"""
    if annot_value is None or len(annot_value) < 1:
        print("❌ No annotation input")
        return None, None, {}

    img_path = annot_value[0]
    bboxes = annot_value[1] if len(annot_value) > 1 else []

    print(f"🖼️ Image path: {img_path}")
    # One standardized image for both the model and the overlay (WYSIWYG).
    img_path = standardize_image(img_path)
    print(f"🧪 Standardized image: {img_path}")
    box_array = None
    if use_box_choice == "Yes" and bboxes:
        box = parse_bboxes(bboxes)
        if box:
            box_array = box
            print(f"📦 Using bounding boxes: {box_array}")


    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        mask = run_seg(SEG_MODEL, img_path, box=box_array, device=device)
        print("📏 mask shape:", mask.shape, "dtype:", mask.dtype)
    except Exception as e:
        print(f"❌ Inference failed: {str(e)}")
        return None, None, {}

    temp_mask_file = tempfile.NamedTemporaryFile(delete=False, suffix=".tif")
    mask_img = Image.fromarray(mask.astype(np.uint16))
    mask_img.save(temp_mask_file.name)
    print(f"💾 Original mask saved to: {temp_mask_file.name}")

    try:
        img = Image.open(img_path)
        print("📷 Image mode:", img.mode, "size:", img.size)
    except Exception as e:
        print(f"❌ Failed to open image: {e}")
        return None, None, {}

    try:
        img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
        img_np = np.array(img_rgb, dtype=np.float32)
        if img_np.max() > 1.5:
            img_np = img_np / 255.0
    except Exception as e:
        print(f"❌ Error in image conversion/resizing: {e}")
        return None, None, {}

    mask_np = np.array(mask)
    inst_mask = mask_np.astype(np.int32)
    unique_ids = np.unique(inst_mask)
    num_instances = len(unique_ids[unique_ids != 0])
    if num_instances == 0:
        print("⚠️ No instance found, returning dummy red image")
        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None, {}

    overlay_img = render_seg_overlay(img_np, inst_mask, overlay_alpha)
    seg_vis_cache = {"img_np": img_np, "inst_mask": inst_mask}
    return overlay_img, temp_mask_file.name, seg_vis_cache


@spaces.GPU
def count_cells_handler(use_box_choice, annot_value, overlay_alpha):
    """Counting handler - supports bounding box, returns only density map"""
    if annot_value is None or len(annot_value) < 1:
        return None, None, "⚠️ Please provide an image.", {}

    image_path = annot_value[0]
    bboxes = annot_value[1] if len(annot_value) > 1 else []

    print(f"🖼️ Image path: {image_path}")
    # One standardized image for both the model and the overlay (WYSIWYG).
    image_path = standardize_image(image_path)
    print(f"🧪 Standardized image: {image_path}")
    box_array = None
    if use_box_choice == "Yes" and bboxes:
        box = parse_bboxes(bboxes)
        if box:
            box_array = box
            print(f"📦 Using bounding boxes: {box_array}")

    try:
        print(f"🔢 Counting - Image: {image_path}")

        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        result = run_count(
            COUNT_MODEL,
            image_path,
            box=box_array,
            device=device,
            visualize=True
        )
        
        if 'error' in result:
            return None, None, f"❌ Counting failed: {result['error']}", {}
        
        count = result['count']
        density_map = result['density_map']
        temp_density_file = tempfile.NamedTemporaryFile(delete=False, suffix=".npy")
        np.save(temp_density_file.name, density_map)
        print(f"💾 Density map saved to {temp_density_file.name}")
        

        try:
            img = Image.open(image_path)
            print("📷 Image mode:", img.mode, "size:", img.size)
        except Exception as e:
            print(f"❌ Failed to open image: {e}")
            return None, None, f"❌ Failed to open image: {str(e)}", {}

        try:
            img_rgb = img.convert("RGB").resize(density_map.shape[::-1], resample=Image.BILINEAR)
            img_np = np.array(img_rgb, dtype=np.float32)
            img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min() + 1e-8)
            if img_np.max() > 1.5:
                img_np = img_np / 255.0
        except Exception as e:
            print(f"❌ Error in image conversion/resizing: {e}")
            return None, None, f"❌ Error in image conversion/resizing: {str(e)}", {}

        
        density_normalized = density_map.copy()
        if density_normalized.max() > 0:
            density_normalized = (density_normalized - density_normalized.min()) / (density_normalized.max() - density_normalized.min())
        
        overlay_img = render_count_overlay(img_np, density_normalized, overlay_alpha)
        result_text = f"✅ Detected {round(count)} objects"
        if use_box_choice == "Yes" and box_array:
            result_text += f"\n📦 Using bounding box: {box_array}"
        

        print(f"✅ Counting done - Count: {count:.1f}")

        count_vis_cache = {"img_np": img_np, "density_normalized": density_normalized}
        return overlay_img, temp_density_file.name, result_text, count_vis_cache
        
        
    except Exception as e:
        print(f"❌ Counting error: {e}")
        import traceback
        traceback.print_exc()
        return None, None, f"❌ Counting failed: {str(e)}", {}


def find_tif_dir(root_dir):
    """Recursively find the first directory containing .tif files"""
    for dirpath, _, filenames in os.walk(root_dir):
        if '__MACOSX' in dirpath:
            continue
        if any(f.lower().endswith('.tif') for f in filenames):
            return dirpath
    return None

def is_valid_tiff(filepath):
    """Check if a file is a valid TIFF image"""
    try:
        with Image.open(filepath) as img:
            img.verify()
            return True
    except Exception as e:
        return False

def find_valid_tif_dir(root_dir):
    """Recursively find the first directory containing valid .tif files"""
    for dirpath, dirnames, filenames in os.walk(root_dir):
        if '__MACOSX' in dirpath:
            continue
        
        potential_tifs = [
            os.path.join(dirpath, f) 
            for f in filenames 
            if f.lower().endswith(('.tif', '.tiff')) and not f.startswith('._')
        ]
        
        if not potential_tifs:
            continue
        
        valid_tifs = [f for f in potential_tifs if is_valid_tiff(f)]
        
        if valid_tifs:
            print(f"✅ Found {len(valid_tifs)} valid TIFF files in: {dirpath}")
            return dirpath
    
    return None

def create_ctc_results_zip(output_dir):
    """
    Create a ZIP file with CTC format results
    
    Parameters:
    -----------
    output_dir : str
        Directory containing tracking results (res_track.txt, etc.)
    
    Returns:
    --------
    zip_path : str
        Path to created ZIP file
    """
    # Create temp directory for ZIP
    temp_zip_dir = tempfile.mkdtemp()
    zip_filename = f"tracking_results_{time.strftime('%Y%m%d_%H%M%S')}.zip"
    zip_path = os.path.join(temp_zip_dir, zip_filename)
    
    print(f"📦 Creating results ZIP: {zip_path}")
    
    # Create ZIP with all tracking results
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        # Add all files from output directory
        for root, dirs, files in os.walk(output_dir):
            for file in files:
                file_path = os.path.join(root, file)
                arcname = os.path.relpath(file_path, output_dir)
                zipf.write(file_path, arcname)
                print(f"  📄 Added: {arcname}")
        
        # Add a README with summary
        readme_content = f"""Tracking Results Summary
            ========================

            Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}

            Files:
            ------
            - res_track.txt: CTC format tracking data
            Format: track_id start_frame end_frame parent_id
            
            - Segmentation masks

            For more information on CTC format:
            http://celltrackingchallenge.net/
            """
        zipf.writestr("README.txt", readme_content)
    
    print(f"✅ ZIP created: {zip_path} ({os.path.getsize(zip_path) / 1024:.1f} KB)")
    return zip_path


def get_well_spaced_color(track_id, num_colors=256):
    """Generate well-spaced colors, using contrasting colors for adjacent IDs"""

    golden_ratio = 0.618033988749895
    hue = (track_id * golden_ratio) % 1.0

    import colorsys
    rgb = colorsys.hsv_to_rgb(hue, 0.9, 0.95)
    return np.array(rgb)


def extract_first_frame(tif_dir):
    """
    Extract the first frame from a directory of TIF files
    
    Returns:
    --------
    first_frame_path : str
        Path to the first TIF frame
    """
    tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 
                      glob(os.path.join(tif_dir, "*.tiff")))
    valid_tif_files = [f for f in tif_files 
                      if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]
    
    if valid_tif_files:
        return valid_tif_files[0]
    return None

def create_tracking_visualization(tif_dir, output_dir, valid_tif_files, overlay_alpha=0.3):
    """
    Create an animated GIF/video showing tracked objects with consistent colors
    
    Parameters:
    -----------
    tif_dir : str
        Directory containing input TIF frames
    output_dir : str
        Directory containing tracking results (masks)
    valid_tif_files : list
        List of valid TIF file paths
    
    Returns:
    --------
    video_path : str
        Path to generated visualization (GIF or first frame)
    """
    import numpy as np
    from matplotlib import colormaps
    from skimage import measure
    import tifffile
    
    # Look for tracking mask files in output directory
    # Common CTC formats: man_track*.tif, mask*.tif, or numbered masks
    mask_files = natsorted(glob(os.path.join(output_dir, "mask*.tif")) + 
                       glob(os.path.join(output_dir, "man_track*.tif")) +
                       glob(os.path.join(output_dir, "*.tif")))
    
    if not mask_files:
        print("⚠️  No mask files found in output directory")
        # Return first frame as fallback
        return valid_tif_files[0]
    
    print(f"📊 Found {len(mask_files)} mask files")

    
    frames = []
    alpha = float(np.clip(overlay_alpha, 0.0, 1.0))  # Transparency for overlay
    
    # Process each frame
    num_frames = min(len(valid_tif_files), len(mask_files))
    for i in range(num_frames):
        try:
            # Load original image using tifffile (handles ZSTD compression)
            try:
                img_np = tifffile.imread(valid_tif_files[i])

                # Normalize to [0, 1] range based on actual data type and values
                if img_np.dtype == np.uint8:
                    img_np = img_np.astype(np.float32) / 255.0
                elif img_np.dtype == np.uint16:
                    # Normalize uint16 to [0, 1] using actual min/max
                    img_min, img_max = img_np.min(), img_np.max()
                    if img_max > img_min:
                        img_np = (img_np.astype(np.float32) - img_min) / (img_max - img_min)
                    else:
                        img_np = img_np.astype(np.float32) / 65535.0
                else:
                    # For float or other types, normalize based on actual range
                    img_np = img_np.astype(np.float32)
                    img_min, img_max = img_np.min(), img_np.max()
                    if img_max > img_min:
                        img_np = (img_np - img_min) / (img_max - img_min)
                    else:
                        img_np = np.clip(img_np, 0, 1)

                # Convert to RGB if grayscale
                if img_np.ndim == 2:
                    img_np = np.stack([img_np]*3, axis=-1)
                img_np = img_np.astype(np.float32)
                if img_np.max() > 1.5:
                    img_np = img_np / 255.0
            except Exception as e:
                print(f"⚠️  Error loading image frame {i}: {e}")
                # Fallback to PIL
                img = Image.open(valid_tif_files[i]).convert("RGB")
                img_np = np.array(img, dtype=np.float32) / 255.0
            
            # Load tracking mask using tifffile (handles ZSTD compression)
            try:
                mask = tifffile.imread(mask_files[i])
            except Exception as e:
                print(f"⚠️  Error loading mask frame {i}: {e}")
                # Fallback to PIL
                mask = np.array(Image.open(mask_files[i]))
            
            # Resize mask to match image if needed
            if mask.shape[:2] != img_np.shape[:2]:
                from scipy.ndimage import zoom
                zoom_factors = [img_np.shape[0] / mask.shape[0], img_np.shape[1] / mask.shape[1]]
                mask = zoom(mask, zoom_factors, order=0).astype(mask.dtype)
            
            # Create overlay
            overlay = img_np.copy()
            
            # Get unique track IDs (excluding background 0)
            track_ids = np.unique(mask)
            track_ids = track_ids[track_ids != 0]
            
            # Color each tracked object
            for track_id in track_ids:
                # Create binary mask for this track
                binary_mask = (mask == track_id)
                
                # Get consistent color for this track ID
                # color = np.array(cmap(int(track_id) % 256)[:3])
                color = get_well_spaced_color(int(track_id))
                
                # Blend color onto image
                overlay[binary_mask] = (1 - alpha) * overlay[binary_mask] + alpha * color
                
                # Draw contours (optional, adds yellow boundaries)
                try:
                    contours = measure.find_contours(binary_mask.astype(np.uint8), 0.5)
                    for contour in contours:
                        contour = contour.astype(np.int32)
                        valid_y = np.clip(contour[:, 0], 0, overlay.shape[0] - 1)
                        valid_x = np.clip(contour[:, 1], 0, overlay.shape[1] - 1)
                        overlay[valid_y, valid_x] = [1.0, 1.0, 0.0]  # Yellow contour
                except:
                    pass  # Skip contours if they fail
            
            # Convert to uint8
            overlay_uint8 = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
            frames.append(Image.fromarray(overlay_uint8))
            
            if i % 10 == 0 or i == num_frames - 1:
                print(f"  📸 Processed frame {i+1}/{num_frames}")
        
        except Exception as e:
            print(f"⚠️  Error processing frame {i}: {e}")
            import traceback
            traceback.print_exc()
            continue
    
    if not frames:
        print("⚠️  No frames were processed successfully")
        return valid_tif_files[0]
    
    # Save as animated GIF
    try:
        temp_gif = tempfile.NamedTemporaryFile(delete=False, suffix=".gif")
        frames[0].save(
            temp_gif.name,
            save_all=True,
            append_images=frames[1:],
            duration=200,  # 200ms per frame = 5fps
            loop=0
        )
        temp_gif.close()  # Close the file handle
        print(f"✅ Created tracking visualization GIF: {temp_gif.name}")
        print(f"   Size: {os.path.getsize(temp_gif.name)} bytes, Frames: {len(frames)}")
        return temp_gif.name
    except Exception as e:
        print(f"⚠️  Failed to create GIF: {e}")
        import traceback
        traceback.print_exc()
        # Return first frame as static image fallback
        try:
            temp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            frames[0].save(temp_img.name)
            temp_img.close()
            return temp_img.name
        except:
            return valid_tif_files[0]

@spaces.GPU
def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj, overlay_alpha, prev_track_vis_cache):
    """
    Tracking handler - processes a ZIP of TIF frames, supports bounding box, returns visualization and results ZIP
    
    Parameters:
    -----------
    use_box_choice : str
        "Yes" or "No" - whether to use bounding box annotation for tracking
    first_frame_annot : tuple or None
        (image_path, bboxes) from BBoxAnnotator, only used if user annotated first frame
    zip_file_obj : File
        Uploaded ZIP file containing TIF sequence
    """
    if zip_file_obj is None:
        return None, "⚠️  Please upload a ZIP file containing video frames (.zip)", None, None, {}
    
    cleanup_tracking_cache(prev_track_vis_cache)
    temp_dir = None
    output_temp_dir = None
    
    try:
        # Parse bounding box if provided
        box_array = None
        if use_box_choice == "Yes" and first_frame_annot is not None:
            if isinstance(first_frame_annot, (list, tuple)) and len(first_frame_annot) > 1:
                bboxes = first_frame_annot[1]
                if bboxes:
                    box = parse_bboxes(bboxes)
                    if box:
                        box_array = box
                        print(f"📦 Using bounding boxes: {box_array}")
        
        # Extract input ZIP
        temp_dir = tempfile.mkdtemp()
        print(f"\n📦 Extracting to temporary directory: {temp_dir}")

        with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:
            extracted_count = 0
            skipped_count = 0
            
            for member in zip_ref.namelist():
                basename = os.path.basename(member)
                
                if ('__MACOSX' in member or 
                    basename.startswith('._') or 
                    basename.startswith('.DS_Store') or
                    member.endswith('/')):
                    skipped_count += 1
                    continue
                
                try:
                    zip_ref.extract(member, temp_dir)
                    extracted_count += 1
                    if basename.lower().endswith(('.tif', '.tiff')):
                        print(f"📄 Extracted TIFF: {basename}")
                except Exception as e:
                    print(f"⚠️  Failed to extract {member}: {e}")

            print(f"\n📊 Extracted: {extracted_count} files, Skipped: {skipped_count} files")

        # Find valid TIFF directory
        tif_dir = find_valid_tif_dir(temp_dir)
        
        if tif_dir is None:
            return None, "❌ Did not find valid TIF directory", None, None, {}
        
        # Validate TIFF files
        tif_files = natsorted(glob(os.path.join(tif_dir, "*.tif")) + 
                          glob(os.path.join(tif_dir, "*.tiff")))
        valid_tif_files = [f for f in tif_files 
                          if not os.path.basename(f).startswith('._') and is_valid_tiff(f)]
        
        if len(valid_tif_files) == 0:
            return None, "❌ Did not find valid TIF files", None, None, {}

        print(f"📈 Using {len(valid_tif_files)} TIF files")

        # Store paths for later visualization
        first_frame_path = valid_tif_files[0]

        # Create temporary output directory for CTC results
        output_temp_dir = tempfile.mkdtemp()
        print(f"💾 CTC-format results will be saved to: {output_temp_dir}")

        # Run tracking with optional bounding box
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        result = run_track(
            TRACK_MODEL,
            video_dir=tif_dir,
            box=box_array,  # Pass bounding box if specified
            device=device,
            output_dir=output_temp_dir
        )
        
        if 'error' in result:
            return None, f"❌ Tracking failed: {result['error']}", None, None, {}
        
        # Create visualization video of tracked objects
        print("\n🎬 Creating tracking visualization...")
        try:
            tracking_video = create_tracking_visualization(
                tif_dir, 
                output_temp_dir, 
                valid_tif_files,
                overlay_alpha=overlay_alpha
            )
        except Exception as e:
            print(f"⚠️  Failed to create visualization: {e}")
            import traceback
            traceback.print_exc()
            # Fallback to first frame if visualization fails
            try:
                tracking_video = Image.open(first_frame_path)
            except:
                tracking_video = None
        
        # Create downloadable ZIP with results
        try:
            results_zip = create_ctc_results_zip(output_temp_dir)
        except Exception as e:
            print(f"⚠️  Failed to create ZIP: {e}")
            results_zip = None
        
        bbox_info = ""
        if box_array:
            bbox_info = f"\n🔲 Using bounding box: [{box_array[0][0]}, {box_array[0][1]}, {box_array[0][2]}, {box_array[0][3]}]"

        result_text = f"""✅ Tracking completed!

            🖼️  Processed frames: {len(valid_tif_files)}{bbox_info}

            📥 Click the button below to download CTC-format results
            The results include:
            - res_track.txt (CTC-format tracking data)
            - Other tracking-related files
            - README.txt (Results description)
            """

        if use_box_choice == "Yes" and box_array:
            result_text += f"\n📦 Using bounding box: {box_array}"

        print(f"\n✅ Tracking completed")

        track_vis_cache = {
            "tif_dir": tif_dir,
            "valid_tif_files": valid_tif_files,
            "output_dir": output_temp_dir,
            "input_temp_dir": temp_dir,
        }

        return results_zip, result_text, gr.update(visible=True), tracking_video, track_vis_cache

    except zipfile.BadZipFile:
        return None, "❌ Not a valid ZIP file", None, None, {}
    except Exception as e:
        import traceback
        traceback.print_exc()
        
        # Clean up on error
        for d in [temp_dir, output_temp_dir]:
            if d:
                try:
                    shutil.rmtree(d)
                except:
                    pass
        return None, f"❌ Tracking failed: {str(e)}", None, None, {}



# ===== Example Images =====
example_images_seg = [f for f in glob("example_imgs/seg/*")]
example_images_cnt = [f for f in glob("example_imgs/cnt/*")]
example_tracking_zips = [f for f in glob("example_imgs/tra/*.zip")]

# ===== Gradio UI =====
CSS = """
/* ── Layout ──────────────────────────────────────────── */
.gradio-container {
    max-width: 1380px !important;
    margin: 0 auto !important;
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
}

/* ── Header markdown polish ───────────────────────────── */
.gradio-container .prose h1 {
    font-size: 2rem !important;
    font-weight: 700 !important;
    color: #1e293b !important;
    letter-spacing: -0.5px !important;
    margin-bottom: 10px !important;
}
.gradio-container .prose h3 {
    font-size: 1rem !important;
    font-weight: 600 !important;
    color: #0284c7 !important;
    margin-top: 14px !important;
    margin-bottom: 4px !important;
}
.gradio-container .prose p {
    margin-top: 4px !important;
    margin-bottom: 6px !important;
    color: #475569 !important;
    line-height: 1.7 !important;
}
.gradio-container .prose ul,
.gradio-container .prose ol {
    margin-top: 4px !important;
    margin-bottom: 6px !important;
}
.gradio-container .prose li {
    color: #475569 !important;
    line-height: 1.7 !important;
}

/* ── Top-level header section ─────────────────────────── */
.gradio-container > .gap > .prose:first-child {
    background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 50%, #f0fdf4 100%) !important;
    border: 1px solid #bae6fd !important;
    border-radius: 16px !important;
    padding: 28px 36px !important;
    margin-bottom: 20px !important;
    box-shadow: 0 4px 20px rgba(14,165,233,0.08) !important;
}

/* ── Tabs ────────────────────────────────────────────── */
.tabs > .tab-nav {
    border-bottom: 2px solid #e2e8f0 !important;
    margin-bottom: 20px !important;
    gap: 4px !important;
}
.tabs button {
    font-size: 15px !important;
    font-weight: 600 !important;
    padding: 11px 24px !important;
    border-radius: 8px 8px 0 0 !important;
    color: #64748b !important;
    transition: color 0.15s, background 0.15s !important;
}
.tabs button:hover {
    color: #0ea5e9 !important;
    background: #f0f9ff !important;
}
.tabs button.selected {
    color: #0284c7 !important;
    border-bottom: 3px solid #0284c7 !important;
    background: transparent !important;
}

/* ── Buttons ─────────────────────────────────────────── */
button.primary {
    background: linear-gradient(135deg, #0284c7 0%, #0ea5e9 100%) !important;
    border: none !important;
    border-radius: 10px !important;
    color: #fff !important;
    font-weight: 600 !important;
    font-size: 15px !important;
    box-shadow: 0 3px 12px rgba(14,165,233,0.35) !important;
    transition: transform 0.12s ease, box-shadow 0.15s ease !important;
}
button.primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(14,165,233,0.45) !important;
}
button.secondary {
    border-radius: 10px !important;
    font-weight: 500 !important;
    border: 1.5px solid #cbd5e1 !important;
    color: #475569 !important;
    transition: border-color 0.12s, color 0.12s, background 0.12s !important;
}
button.secondary:hover {
    border-color: #94a3b8 !important;
    color: #1e293b !important;
    background: #f8fafc !important;
}

/* ── Blocks and panels ───────────────────────────────── */
.gradio-container .block { border-radius: 14px !important; }
.gradio-container .gr-form,
.gradio-container .gr-box,
.gradio-container .gr-panel {
    border-radius: 14px !important;
    border-color: #e2e8f0 !important;
}

/* ── Labels ──────────────────────────────────────────── */
label { font-weight: 500 !important; color: #374151 !important; }

/* File/media title pills use .block-label on a <label>, so the rule above
   forces their text dark. Restore the theme color/weight so they match the
   other native labels (e.g. "Overlay Opacity"). */
.block-label,
.block-label span {
    color: var(--block-label-text-color) !important;
    font-weight: var(--block-label-text-weight) !important;
}

/* ── Collapsible "Image requirements" accordion ───────── */

/* The label is a <button class="label-wrap"> whose text sits in an inner <span>.
   Gradio styles that span directly:
       span.svelte-xxx { font-weight: var(--section-header-text-weight);  // = 400
                         font-size:   var(--section-header-text-size); }
   so font-size/weight set on the button is overridden and has no effect - the
   span itself must be targeted. Colour is not set on the span, so it inherits. */
.req-accordion .label-wrap {
    color: var(--block-label-text-color) !important;
}
.req-accordion .label-wrap span:not(.icon) {
    font-size: 0.9rem !important;
    font-weight: 700 !important;
    color: inherit !important;
}

/* ── Frame headings (labels above media frames) ──────── */
/* Neutralize the Gradio block wrapper so it doesn't clip/round the pill */
.frame-heading {
    /* The column has a 16px flex gap; pull the pill back down so it sits
       close to the frame below instead of floating far above it. */
    margin-bottom: -8px !important;
    padding: 0 !important;
    min-width: 0 !important;
    min-height: 0 !important;
    background: transparent !important;
    border: none !important;
    box-shadow: none !important;
    border-radius: 0 !important;
    overflow: visible !important;
}
/* Match the theme's native field labels (e.g. "Overlay Opacity",
   "Download Original Prediction") exactly: same colors, borderless, fully
   rounded — driven by theme variables so it adapts to light/dark mode. */
.frame-heading p {
    display: inline-block !important;
    width: fit-content !important;
    background: var(--block-label-background-fill) !important;
    color: var(--block-label-text-color) !important;
    border: none !important;
    border-radius: var(--block-label-radius) !important;
    font-weight: var(--block-label-text-weight) !important;
    font-size: var(--block-label-text-size) !important;
    line-height: 1.3 !important;
    padding: 5px 12px !important;
    margin: 0 !important;
}

/* ── Image output ────────────────────────────────────── */
.uniform-height {
    height: 480px !important;
    display: flex !important;
    align-items: center !important;
    justify-content: center !important;
    border-radius: 12px !important;
    background: #f8fafc !important;
}
.uniform-height img, .uniform-height canvas {
    max-height: 480px !important;
    object-fit: contain !important;
}

/* ── Density map output ──────────────────────────────── */
#density_map_output { height: 480px !important; }
#density_map_output .image-container { height: 480px !important; }
#density_map_output img {
    height: 460px !important;
    width: auto !important;
    max-width: 95% !important;
    object-fit: contain !important;
}

/* ── Tab content description markdown ───────────────── */
.tabitem .prose h2 {
    font-size: 1.3rem !important;
    font-weight: 700 !important;
    color: #1e293b !important;
    margin-top: 0 !important;
    margin-bottom: 10px !important;
    padding-bottom: 8px !important;
    border-bottom: 2px solid #e0f2fe !important;
}
.tabitem .prose:nth-child(2) {
    background: #f8fafc !important;
    border: 1px solid #e2e8f0 !important;
    border-radius: 10px !important;
    padding: 12px 18px !important;
    margin-bottom: 16px !important;
}
.tabitem .prose:nth-child(2) p,
.tabitem .prose:nth-child(2) li {
    font-size: 0.91rem !important;
    color: #64748b !important;
}
.tabitem .prose:nth-child(2) strong {
    color: #0f172a !important;
}

/* ════════════════════════════════════════════════════════
   DARK MODE  (.dark is added to <html> by Gradio)
   ════════════════════════════════════════════════════════ */

/* ── Header text ─────────────────────────────────────── */
.dark .gradio-container .prose h1 {
    color: #e2e8f0 !important;
}
.dark .gradio-container .prose h3 {
    color: #38bdf8 !important;
}
.dark .gradio-container .prose p,
.dark .gradio-container .prose li {
    color: #94a3b8 !important;
}

/* ── Top-level header card ───────────────────────────── */
.dark .gradio-container > .gap > .prose:first-child {
    background: linear-gradient(135deg, #0c1a2e 0%, #0f2942 50%, #0d1f12 100%) !important;
    border-color: #1e3a5f !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.4) !important;
}

/* ── Tabs ────────────────────────────────────────────── */
.dark .tabs > .tab-nav {
    border-bottom-color: #334155 !important;
}
.dark .tabs button {
    color: #94a3b8 !important;
}
.dark .tabs button:hover {
    color: #38bdf8 !important;
    background: rgba(56,189,248,0.08) !important;
}
.dark .tabs button.selected {
    color: #38bdf8 !important;
    border-bottom-color: #38bdf8 !important;
}

/* ── Buttons ─────────────────────────────────────────── */
.dark button.secondary {
    border-color: #475569 !important;
    color: #94a3b8 !important;
    background: transparent !important;
}
.dark button.secondary:hover {
    border-color: #64748b !important;
    color: #e2e8f0 !important;
    background: rgba(255,255,255,0.05) !important;
}

/* ── Blocks / panels ─────────────────────────────────── */
.dark .gradio-container .gr-form,
.dark .gradio-container .gr-box,
.dark .gradio-container .gr-panel {
    border-color: #334155 !important;
}

/* ── Labels ──────────────────────────────────────────── */
.dark label {
    color: #cbd5e1 !important;
}
.dark .block-label,
.dark .block-label span {
    color: var(--block-label-text-color) !important;
    font-weight: var(--block-label-text-weight) !important;
}

/* ── Image output area ───────────────────────────────── */
.dark .uniform-height {
    background: #1e293b !important;
}

/* ── Tab content markdown ────────────────────────────── */
.dark .tabitem .prose h2 {
    color: #e2e8f0 !important;
    border-bottom-color: #1e3a5f !important;
}
.dark .tabitem .prose:nth-child(2) {
    background: #1e293b !important;
    border-color: #334155 !important;
}
.dark .tabitem .prose:nth-child(2) p,
.dark .tabitem .prose:nth-child(2) li {
    color: #94a3b8 !important;
}
.dark .tabitem .prose:nth-child(2) strong {
    color: #e2e8f0 !important;
}
"""

with gr.Blocks(
    title="Microscopy Analysis Suite",
    theme=gr.themes.Soft(
        primary_hue=gr.themes.colors.sky,
        secondary_hue=gr.themes.colors.slate,
        neutral_hue=gr.themes.colors.slate,
        # A weight must be listed here to be usable; the default is (400, 600).
        font=gr.themes.GoogleFont("Inter", weights=(400, 600, 700, 800)),
    ),
    css=CSS,
) as demo:
    gr.Markdown(
        """
        # 🔬 MicroscopyMatching: Microscopy Image Analysis Suite
        
        ### Supporting three key tasks:
        - 🎨 **Segmentation**: Instance segmentation of microscopic objects
        - 🔢 **Counting**: Counting microscopic objects based on density maps
        - 🎬 **Tracking**: Tracking microscopic objects in video sequences

        ### 💡 Technical Details:

        **MicroscopyMatching** - A general-purpose microscopy image analysis toolkit based on pre-trained Latent Diffusion Model
        
        ### 📒 Note:

        This project is currently available with usage limits for research trial use and feedback collection. Please note that response speed may vary depending on GPU allocation queues. We plan to release a free public version in the future. We are actively improving the toolkit and greatly appreciate your feedback!

        """
    )
    

    # Must stay a callable: a plain value is deepcopied, so every session would
    # share one id.
    current_query_id = gr.State(lambda: str(uuid.uuid4()))
    user_uploaded_examples = gr.State(example_images_seg.copy())
    seg_vis_state = gr.State({})
    count_vis_state = gr.State({})
    track_vis_state = gr.State({})
    # The annotator only holds the flattened preview, so keep the original for
    # re-extracting other stack frames.
    seg_raw_state = gr.State(None)
    count_raw_state = gr.State(None)
    
    with gr.Tabs():
        # ===== Tab 1: Segmentation =====
        with gr.Tab("🎨 Segmentation"):
            gr.Markdown("## Instance Segmentation of Microscopic Objects")
            gr.Markdown(
                        """
                        **Instructions:**
                        1. Select an example image from the Example Image Gallery or upload your own image (**Please see the "Image Upload Requirements" section below before uploading**)
                        2. (Optional) Specify a target object with a bounding box and select "Yes", or click "Run Segmentation" directly
                        3. Click "Run Segmentation"
                        4. View the segmentation results (you can adjust the overlay opacity by sliding the opacity bar below the visualization), download the original predicted mask (.tif format); if needed, click "Clear Selection" to choose a new image

                        💡 Want to reuse an image? Upload it under "➕ Upload New Example Image to Gallery" and click "➕ Add to Gallery". The uploaded image will appear at the front of the gallery, and can be loaded by clicking its thumbnail. Gallery additions last for the current session only and are cleared when you reload the page.

                        🤘 Tell us about your experience by rating and submitting feedback, which would greatly help us improve the framework!
                        """
                    )
            # Must stay after the instructions markdown: the CSS targets it with
            # .prose:nth-child(2).
            with gr.Accordion("📋 Image Upload Requirements (click to expand or fold)", open=False, elem_classes="req-accordion"):
                gr.Markdown(
                    """
                    **Please do not upload images or data containing protected health information (PHI) or personally identifiable information (PII).**

                    - **Formats:** Supports TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, as well as various other microscopy image formats supported by [tifffiles](https://github.com/cgohlke/tifffile)
                    - **Maximum image size:** 4096 × 4096 pixels
                    - **Channels:** Accepts RGB (3-channel), Grayscale (single-channel); for images with more than 3 channels, only the first three channels are used as RGB channels.
                    - **For image stacks:** First frame is loaded by default; you can pick another with the stack frame slider that appears after upload
                    
                    **Alternatively, you can use existing tools (e.g., ImageJ/Fiji) to convert your image to a 8-bit RGB/Grayscale PNG/JPG/TIFF file before uploading.**
                    
                    """
                )
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("🖼️ Upload Image (Optional: Provide a Bounding Box)", elem_classes="frame-heading")
                    annotator = BBoxAnnotator(
                        show_label=False,
                        categories=["cell"],
                    )
                    seg_frame_slider = gr.Slider(
                        minimum=1, maximum=1, step=1, value=1,
                        label="📚 Stack frame (choose frame to use)",
                        visible=False,
                    )
                    seg_zplane_slider = gr.Slider(
                        minimum=1, maximum=1, step=1, value=1,
                        label="🔬 Z-plane (choose plane to use)",
                        visible=False,
                    )

                    # Example Images Gallery
                    gr.Markdown("📁 Example Image Gallery", elem_classes="frame-heading")
                    example_gallery = gr.Gallery(
                        show_label=True,
                        columns=len(example_images_seg),
                        rows=1,
                        height=120,
                        object_fit="cover",
                        show_download_button=False
                    )
                    
                    
                    with gr.Row():
                        use_box_radio = gr.Radio(
                            choices=["Yes", "No"],
                            value="No",
                            label="🔲 Specify Bounding Box?"
                        )
                    with gr.Row():
                        run_seg_btn = gr.Button("▶️ Run Segmentation", variant="primary", size="lg")
                        clear_btn = gr.Button("🔄 Clear Selection", variant="secondary")

                    # Upload Example Image
                    gr.Markdown("➕ Upload New Example Image to Gallery", elem_classes="frame-heading")
                    image_uploader = gr.File(
                        show_label=False,
                        file_types=["image"] + list(TIFF_EXTENSIONS),
                        type="filepath"
                    )
                    add_to_gallery_btn = gr.Button("➕ Add to Gallery", variant="secondary", size="sm")
                    add_gallery_status = gr.Markdown(visible=False)


                with gr.Column(scale=2):
                    gr.Markdown("📸 Segmentation Result", elem_classes="frame-heading")
                    seg_output = gr.Image(
                        type="pil",
                        show_label=False,
                        elem_classes="uniform-height"
                    )
                    seg_alpha_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.5,
                        label="🪄 Overlay Opacity"
                    )
                    
                    # Download Original Prediction
                    gr.Markdown("📥 Download Original Prediction (.tif format)", elem_classes="frame-heading")
                    download_mask_btn = gr.File(
                        show_label=False,
                        visible=True,
                        height=40,
                    )

                    # Satisfaction Rating
                    score_slider = gr.Slider(
                        minimum=1,
                        maximum=5,
                        step=1,
                        value=5,
                        label="🌟 Satisfaction Rating (1-5)"
                    )

                    # Feedback Textbox
                    feedback_box = gr.Textbox(
                        placeholder="Please enter your feedback...",
                        lines=2,
                        label="💬 Feedback"
                    )

                    # Submit Button
                    submit_feedback_btn = gr.Button("💾 Submit Feedback", variant="secondary")

                    feedback_status = gr.Textbox(
                        label="✅ Submission Status",
                        lines=1,
                        visible=False
                    )
            
            annotator.upload(
                fn=prepare_uploaded_image,
                inputs=annotator,
                outputs=[annotator, seg_raw_state, seg_frame_slider, seg_zplane_slider]
            )
            for _slider in (seg_frame_slider, seg_zplane_slider):
                _slider.release(
                    fn=select_frame,
                    inputs=[seg_raw_state, seg_frame_slider, seg_zplane_slider],
                    outputs=annotator
                )

            # click event for segmentation
            run_seg_btn.click(
                fn=segment_with_choice,
                inputs=[use_box_radio, annotator, seg_alpha_slider],
                outputs=[seg_output, download_mask_btn, seg_vis_state]
            )
            seg_alpha_slider.input(
                fn=update_seg_overlay_alpha,
                inputs=[seg_alpha_slider, seg_vis_state],
                outputs=seg_output
            )

            # click event for clear button
            clear_btn.click(
                fn=lambda: (None, {}, None, gr.update(visible=False, value=1),
                            gr.update(visible=False, value=1)),
                inputs=None,
                outputs=[annotator, seg_vis_state, seg_raw_state, seg_frame_slider, seg_zplane_slider]
            )
            
            # init Gallery with example images
            demo.load(
                fn=lambda: example_images_seg.copy(),
                outputs=example_gallery
            )
            
            # Add the uploaded image to the gallery on button click, confirm with
            # a status message, and clear the uploader so the upload box returns.
            def add_to_gallery(img_path, current_imgs):
                if not img_path:
                    return (current_imgs, current_imgs,
                            gr.update(value="⚠️ Please upload an image first.", visible=True), None)
                rejected = check_image(img_path)
                if rejected:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"❌ {rejected}", visible=True), None)
                if len(current_imgs) - len(example_images_seg) >= MAX_GALLERY_UPLOADS:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"⚠️ Gallery upload limit reached "
                                            f"({MAX_GALLERY_UPLOADS} images). Reload the page "
                                            f"to start over.", visible=True), None)
                std_path = standardize_image(img_path)
                _remember_gallery_source(std_path, img_path)
                if std_path not in current_imgs:
                    current_imgs.insert(0, std_path)
                return (current_imgs, current_imgs,
                        gr.update(value="✅ Added to the Example Gallery below.", visible=True), None)

            add_to_gallery_btn.click(
                fn=add_to_gallery,
                inputs=[image_uploader, user_uploaded_examples],
                outputs=[user_uploaded_examples, example_gallery, add_gallery_status, image_uploader]
            )
            
            # Reuse the upload path so the toast and frame slider behave the same.
            def load_from_gallery(evt: gr.SelectData, all_imgs):
                if evt.index is not None and evt.index < len(all_imgs):
                    item = all_imgs[evt.index]
                    return prepare_uploaded_image(_GALLERY_RAW_SOURCE.get(item, item))
                return None, None, gr.update(visible=False, value=1), gr.update(visible=False, value=1)

            example_gallery.select(
                fn=load_from_gallery,
                inputs=user_uploaded_examples,
                outputs=[annotator, seg_raw_state, seg_frame_slider, seg_zplane_slider]
            )
            
            # click event for submitting feedback
            def submit_user_feedback(query_id, score, comment, annot_val):
                try:
                    img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
                    bboxes = annot_val[1] if annot_val and len(annot_val) > 1 else []
                    
                    # save_feedback(
                    #     query_id=query_id,
                    #     feedback_type=f"score_{int(score)}",
                    #     feedback_text=comment,
                    #     img_path=img_path,
                    #     bboxes=bboxes
                    # )

                    save_feedback_to_hf(
                        query_id=query_id,
                        feedback_type=f"score_{int(score)}",
                        feedback_text=comment,
                        img_path=img_path,
                        bboxes=bboxes
                    )
                    return "✅ Feedback submitted, thank you!", gr.update(visible=True)
                except Exception as e:
                    return f"❌ Submission failed: {str(e)}", gr.update(visible=True)

            submit_feedback_btn.click(
                fn=submit_user_feedback,
                inputs=[current_query_id, score_slider, feedback_box, annotator],
                outputs=[feedback_status, feedback_status]
            )
        
        # ===== Tab 2: Counting =====
        with gr.Tab("🔢 Counting"):
            gr.Markdown("## Microscopy Object Counting Analysis")
            gr.Markdown(
                """
                **Usage Instructions:**
                1. Select an example image from the Example Image Gallery or upload your own image (**Please see the "Image Upload Requirements" section below before uploading**)
                2. (Optional) Specify a target object with a bounding box and select "Yes", or click "Run Counting" directly
                3. Click "Run Counting"
                4. View the density map (you can adjust the density opacity by sliding the opacity bar below the visualization), download the original prediction (.npy format); if needed, click "Clear Selection" to choose a new image to run

                💡 Want to reuse an image? Upload it under "➕ Upload New Example Image to Gallery" and click "➕ Add to Gallery". The uploaded image will appear at the front of the gallery, and can be loaded by clicking its thumbnail. Gallery additions last for the current session only and are cleared when you reload the page.

                🤘 Tell us about your experience by rating and submitting feedback, which would greatly help us improve the framework!
                """
            )
            # Must stay after the instructions markdown: the CSS targets it with
            # .prose:nth-child(2).
            with gr.Accordion("📋 Image requirements", open=False, elem_classes="req-accordion"):
                gr.Markdown(
                    """
                    **Please do not upload images or data containing protected health information (PHI) or personally identifiable information (PII).**

                    - **Formats:** Supports TIFF / OME-TIFF (8/16/32-bit, stacks, multi-channel), PNG, JPG, as well as various other microscopy image formats supported by [tifffiles](https://github.com/cgohlke/tifffile)
                    - **Maximum image size:** 4096 × 4096 pixels
                    - **Channels:** Accepts RGB (3-channel), Grayscale (single-channel); for images with more than 3 channels, only the first three channels are used as RGB channels.
                    - **For image stacks:** First frame is loaded by default; you can pick another with the stack frame slider that appears after upload
                    
                    **Alternatively, you can use existing tools (e.g., ImageJ/Fiji) to convert your image to a 8-bit RGB/Grayscale PNG/JPG/TIFF file before uploading.**
                    """
                )
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("🖼️ Upload Image (Optional: Provide a Bounding Box)", elem_classes="frame-heading")
                    count_annotator = BBoxAnnotator(
                        show_label=False,
                        categories=["cell"],
                    )
                    count_frame_slider = gr.Slider(
                        minimum=1, maximum=1, step=1, value=1,
                        label="📚 Stack frame (choose frame to use)",
                        visible=False,
                    )
                    count_zplane_slider = gr.Slider(
                        minimum=1, maximum=1, step=1, value=1,
                        label="🔬 Z-plane (choose plane to use)",
                        visible=False,
                    )

                    # Example gallery with "add" functionality
                    gr.Markdown("📁 Example Image Gallery", elem_classes="frame-heading")
                    with gr.Row():
                        count_example_gallery = gr.Gallery(
                            show_label=False,
                            columns=len(example_images_cnt),
                            rows=1,
                            object_fit="cover",
                            height=120,
                            value=example_images_cnt.copy(),  # Initialize with examples
                            show_download_button=False
                        )
                    
                    
                    with gr.Row():
                        count_use_box_radio = gr.Radio(
                            choices=["Yes", "No"],
                            value="No",
                            label="🔲 Specify Bounding Box?"
                        )

                    with gr.Row():
                        count_btn = gr.Button("▶️ Run Counting", variant="primary", size="lg")
                        clear_btn = gr.Button("🔄 Clear Selection", variant="secondary")
                    
                    # Add button to upload new examples
                    gr.Markdown("➕ Add Example Image to Gallery", elem_classes="frame-heading")
                    count_image_uploader = gr.File(
                        show_label=False,
                        file_types=["image"] + list(TIFF_EXTENSIONS),
                        type="filepath"
                    )
                    add_to_count_gallery_btn = gr.Button("➕ Add to Gallery", variant="secondary", size="sm")
                    count_add_status = gr.Markdown(visible=False)

                
                with gr.Column(scale=2):
                    gr.Markdown("📸 Density Map", elem_classes="frame-heading")
                    count_output = gr.Image(
                        show_label=False,
                        type="filepath",
                        elem_id="density_map_output"

                    )
                    count_alpha_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.3,
                        label="🪄 Density Opacity"
                    )
                    count_status = gr.Textbox(
                        label="📊 Statistics",
                        lines=2
                    )
                    gr.Markdown("📥 Download Original Prediction (.npy format)", elem_classes="frame-heading")
                    download_density_btn = gr.File(
                        show_label=False,
                        visible=True
                    )

                    # Satisfaction rating
                    score_slider = gr.Slider(
                        minimum=1,
                        maximum=5,
                        step=1,
                        value=5,
                        label="🌟 Satisfaction Rating (1-5)"
                    )

                    # Feedback textbox
                    feedback_box = gr.Textbox(
                        placeholder="Please enter your feedback...",
                        lines=2,
                        label="💬 Feedback"
                    )

                    # Submit button
                    submit_feedback_btn = gr.Button("💾 Submit Feedback", variant="secondary")

                    feedback_status = gr.Textbox(
                        label="✅ Submission Status",
                        lines=1,
                        visible=False
                    )
            
            # State for managing gallery images
            count_user_examples = gr.State(example_images_cnt.copy())
            
            # Add the uploaded image to the gallery on button click, confirm with
            # a status message, and clear the uploader afterward.
            def add_to_count_gallery(new_img_file, current_imgs):
                """Add uploaded image to gallery"""
                if new_img_file is None:
                    return (current_imgs, current_imgs,
                            gr.update(value="⚠️ Please upload an image first.", visible=True), None)
                rejected = check_image(new_img_file)
                if rejected:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"❌ {rejected}", visible=True), None)
                if len(current_imgs) - len(example_images_cnt) >= MAX_GALLERY_UPLOADS:
                    return (current_imgs, current_imgs,
                            gr.update(value=f"⚠️ Gallery upload limit reached "
                                            f"({MAX_GALLERY_UPLOADS} images). Reload the page "
                                            f"to start over.", visible=True), None)
                std_path = standardize_image(new_img_file)
                _remember_gallery_source(std_path, new_img_file)
                if std_path not in current_imgs:
                    current_imgs.insert(0, std_path)
                    print(f"✅ Added image to gallery: {std_path}")
                return (current_imgs, current_imgs,
                        gr.update(value="✅ Added to the Example Gallery above.", visible=True), None)

            # When user clicks "Add to Gallery"
            add_to_count_gallery_btn.click(
                fn=add_to_count_gallery,
                inputs=[count_image_uploader, count_user_examples],
                outputs=[count_user_examples, count_example_gallery, count_add_status, count_image_uploader]
            )
            
            count_annotator.upload(
                fn=prepare_uploaded_image,
                inputs=count_annotator,
                outputs=[count_annotator, count_raw_state, count_frame_slider, count_zplane_slider]
            )
            for _slider in (count_frame_slider, count_zplane_slider):
                _slider.release(
                    fn=select_frame,
                    inputs=[count_raw_state, count_frame_slider, count_zplane_slider],
                    outputs=count_annotator
                )

            # When user selects from gallery, load into annotator
            def load_from_count_gallery(evt: gr.SelectData, all_imgs):
                """Load a gallery image, reusing the upload path."""
                if evt.index is not None and evt.index < len(all_imgs):
                    selected_img = all_imgs[evt.index]
                    print(f"📸 Loading image from gallery: {selected_img}")
                    return prepare_uploaded_image(_GALLERY_RAW_SOURCE.get(selected_img, selected_img))
                return None, None, gr.update(visible=False, value=1), gr.update(visible=False, value=1)

            count_example_gallery.select(
                fn=load_from_count_gallery,
                inputs=count_user_examples,
                outputs=[count_annotator, count_raw_state, count_frame_slider, count_zplane_slider]
            )

            # Run counting
            count_btn.click(
                fn=count_cells_handler,
                inputs=[count_use_box_radio, count_annotator, count_alpha_slider],
                outputs=[count_output, download_density_btn, count_status, count_vis_state]
            )
            count_alpha_slider.input(
                fn=update_count_overlay_alpha,
                inputs=[count_alpha_slider, count_vis_state],
                outputs=count_output
            )

            # Clear selection
            clear_btn.click(
                fn=lambda: (None, {}, None, gr.update(visible=False, value=1),
                            gr.update(visible=False, value=1)),
                inputs=None,
                outputs=[count_annotator, count_vis_state, count_raw_state, count_frame_slider, count_zplane_slider]
            )

            # Submit feedback
            def submit_user_feedback(query_id, score, comment, annot_val):
                try:
                    img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
                    bboxes = annot_val[1] if annot_val and len(annot_val) > 1 else []


                    save_feedback_to_hf(
                        query_id=query_id,
                        feedback_type=f"score_{int(score)}",
                        feedback_text=comment,
                        img_path=img_path,
                        bboxes=bboxes
                    )
                    return "✅ Feedback submitted successfully, thank you!", gr.update(visible=True)
                except Exception as e:
                    return f"❌ Submission failed: {str(e)}", gr.update(visible=True)

            submit_feedback_btn.click(
                fn=submit_user_feedback,
                inputs=[current_query_id, score_slider, feedback_box, annotator],
                outputs=[feedback_status, feedback_status]
            )
        
        # ===== Tab 3: Tracking =====
        with gr.Tab("🎬 Tracking"):
            gr.Markdown("## Microscopy Object Video Tracking - Supports ZIP Upload")
            gr.Markdown(
                        """
                        **Instructions:**
                        1. Select a video from the example library or Upload your own video ZIP file. The ZIP should contain a sequence of TIF images named in chronological order (e.g., t000.tif, t001.tif...)
                        2. (Optional) Specify a target object with a bounding box on the first frame and select "Yes", or click "Run Tracking" directly
                        3. Click "Run Tracking"
                        4. View the tracking results (you can adjust the overlay opacity by sliding the opacity bar below the visualization), download the CTC format results; if needed, click "Clear Selection" to choose a new ZIP file to run

                        💡 Want to reuse a video? Upload the ZIP under "➕ Upload New Example Image to Gallery" and click "➕ Add to Gallery". The uploaded video will appear at the front of the gallery, and can be loaded by clicking its thumbnail. Gallery additions last for the current session only and are cleared when you reload the page.
                        
                        🤘 Tell us about your experience by rating and submitting feedback, which would greatly help us improve the framework!

                        """
                    )
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("📦 Upload Image Sequence in ZIP File", elem_classes="frame-heading")
                    track_zip_upload = gr.File(
                        show_label=False,
                        file_types=[".zip"]
                    )

                    # First frame annotation for bounding box (heading + annotator
                    # share one column so they show/hide together)
                    with gr.Column(visible=False) as track_annot_group:
                        gr.Markdown("🖼️ (Optional) First Frame Bounding Box Annotation", elem_classes="frame-heading")
                        track_first_frame_annotator = BBoxAnnotator(
                            show_label=False,
                            categories=["cell"],
                        )

                    # Example ZIP gallery
                    gr.Markdown("📁 Example Video Gallery (Click to Select)", elem_classes="frame-heading")
                    track_example_gallery = gr.Gallery(
                        show_label=False,
                        columns=10,
                        rows=1,
                        height=120,
                        object_fit="contain",
                        show_download_button=False
                    )
                    
                    with gr.Row():
                        track_use_box_radio = gr.Radio(
                            choices=["Yes", "No"],
                            value="No",
                            label="🔲 Specify Bounding Box?"
                        )

                    with gr.Row():
                        track_btn = gr.Button("▶️ Run Tracking", variant="primary", size="lg")
                        clear_btn = gr.Button("🔄 Clear Selection", variant="secondary")
                    
                    # Add to gallery button
                    gr.Markdown("➕ Add ZIP to Example Gallery", elem_classes="frame-heading")
                    track_gallery_upload = gr.File(
                        show_label=False,
                        file_types=[".zip"],
                        type="filepath"
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("📸 Tracking Visualization", elem_classes="frame-heading")
                    track_first_frame_preview = gr.Image(
                        show_label=False,
                        type="filepath",
                        # height=400,
                        elem_classes="uniform-height",
                        interactive=False
                    )
                    track_alpha_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        step=0.05,
                        value=0.3,
                        label="🪄 Overlay Opacity"
                    )
                    
                    track_output = gr.Textbox(
                        label="📊 Tracking Information",
                        lines=8,
                        interactive=False
                    )
                    
                    # heading + download share one column so they show/hide together
                    with gr.Column(visible=False) as track_dl_group:
                        gr.Markdown("📥 Download Tracking Results (CTC Format)", elem_classes="frame-heading")
                        track_download = gr.File(
                            show_label=False,
                        )

                    # Satisfaction rating
                    score_slider = gr.Slider(
                        minimum=1,
                        maximum=5,
                        step=1,
                        value=5,
                        label="🌟 Satisfaction Rating (1-5)"
                    )

                    # Feedback textbox
                    feedback_box = gr.Textbox(
                        placeholder="Please enter your feedback...",
                        lines=2,
                        label="💬 Feedback"
                    )

                    # Submit button
                    submit_feedback_btn = gr.Button("💾 Submit Feedback", variant="secondary")

                    feedback_status = gr.Textbox(
                        label="✅ Submission Status",
                        lines=1,
                        visible=False
                    )
            
            # State for tracking examples
            track_user_examples = gr.State(example_tracking_zips.copy())
            
            # Function to get preview image from ZIP
            def get_zip_preview(zip_path):
                """Extract first frame from ZIP for gallery preview"""
                try:
                    temp_dir = tempfile.mkdtemp()
                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                        for member in zip_ref.namelist():
                            basename = os.path.basename(member)
                            if ('__MACOSX' not in member and 
                                not basename.startswith('._') and
                                basename.lower().endswith(('.tif', '.tiff', '.png', '.jpg'))):
                                zip_ref.extract(member, temp_dir)
                                extracted_path = os.path.join(temp_dir, member)
                                
                                # Load and normalize for preview
                                import tifffile
                                import numpy as np
                                
                                img_np = tifffile.imread(extracted_path)
                                if img_np.dtype == np.uint16:
                                    img_min, img_max = img_np.min(), img_np.max()
                                    if img_max > img_min:
                                        img_np = ((img_np.astype(np.float32) - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                                
                                if img_np.ndim == 2:
                                    img_np = np.stack([img_np]*3, axis=-1)
                                
                                # Save preview
                                preview_path = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
                                Image.fromarray(img_np).save(preview_path.name)
                                return preview_path.name
                except:
                    pass
                return None
            
            # Initialize gallery with previews
            def init_tracking_gallery():
                """Create preview images for ZIP examples"""
                previews = []
                for zip_path in example_tracking_zips:
                    if os.path.exists(zip_path):
                        preview = get_zip_preview(zip_path)
                        if preview:
                            previews.append(preview)
                return previews
            
            # Load gallery on startup
            demo.load(
                fn=init_tracking_gallery,
                outputs=track_example_gallery
            )
            
            # Add ZIP to gallery
            def add_zip_to_gallery(zip_path, current_zips):
                if not zip_path:
                    return current_zips, track_example_gallery
                try:
                    if zip_path not in current_zips:
                        current_zips.append(zip_path)
                        print(f"✅ Added ZIP to gallery: {zip_path}")
                    # Regenerate previews
                    previews = []
                    for zp in current_zips:
                        preview = get_zip_preview(zp)
                        if preview:
                            previews.append(preview)
                    return current_zips, previews
                except Exception as e:
                    print(f"⚠️ Error: {e}")
                    return current_zips, []
            
            track_gallery_upload.upload(
                fn=add_zip_to_gallery,
                inputs=[track_gallery_upload, track_user_examples],
                outputs=[track_user_examples, track_example_gallery]
            )
            
            # Select ZIP from gallery
            def load_zip_from_gallery(evt: gr.SelectData, all_zips):
                if evt.index is not None and evt.index < len(all_zips):
                    selected_zip = all_zips[evt.index]
                    print(f"📁 Selected ZIP from gallery: {selected_zip}")
                    return selected_zip
                return None
            
            track_example_gallery.select(
                fn=load_zip_from_gallery,
                inputs=track_user_examples,
                outputs=track_zip_upload
            )

            # Load first frame when ZIP is uploaded
            def load_first_frame_for_annotation(zip_file_obj):
                '''Load and normalize first frame from ZIP for annotation'''
                if zip_file_obj is None:
                    return None, gr.update(visible=False)
                
                import tifffile
                import numpy as np
                
                try:
                    temp_dir = tempfile.mkdtemp()
                    with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:
                        for member in zip_ref.namelist():
                            basename = os.path.basename(member)
                            if ('__MACOSX' not in member and 
                                not basename.startswith('._') and
                                basename.lower().endswith(('.tif', '.tiff'))):
                                zip_ref.extract(member, temp_dir)
                    
                    tif_dir = find_valid_tif_dir(temp_dir)
                    if tif_dir:
                        first_frame = extract_first_frame(tif_dir)
                        if first_frame:
                            # Load and normalize the first frame
                            try:
                                img_np = tifffile.imread(first_frame)
                                
                                # Normalize to [0, 255] uint8 range for display
                                if img_np.dtype == np.uint8:
                                    pass  # Already uint8
                                elif img_np.dtype == np.uint16:
                                    # Normalize uint16 using actual min/max
                                    img_min, img_max = img_np.min(), img_np.max()
                                    if img_max > img_min:
                                        img_np = ((img_np.astype(np.float32) - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                                    else:
                                        img_np = (img_np.astype(np.float32) / 65535.0 * 255).astype(np.uint8)
                                else:
                                    # Float or other types
                                    img_np = img_np.astype(np.float32)
                                    img_min, img_max = img_np.min(), img_np.max()
                                    if img_max > img_min:
                                        img_np = ((img_np - img_min) / (img_max - img_min) * 255).astype(np.uint8)
                                    else:
                                        img_np = np.clip(img_np * 255, 0, 255).astype(np.uint8)
                                
                                # Convert to RGB if grayscale
                                if img_np.ndim == 2:
                                    img_np = np.stack([img_np]*3, axis=-1)
                                elif img_np.ndim == 3 and img_np.shape[2] > 3:
                                    img_np = img_np[:, :, :3]
                                
                                # Save normalized image to temp file
                                temp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
                                Image.fromarray(img_np).save(temp_img.name)
                                
                                print(f"✅ Loaded and normalized first frame: {first_frame}")
                                print(f"   Original dtype: {tifffile.imread(first_frame).dtype}")
                                print(f"   Normalized to uint8 RGB for annotation")
                                
                                return temp_img.name, gr.update(visible=True)
                            except Exception as e:
                                print(f"⚠️  Error normalizing first frame: {e}")
                                import traceback
                                traceback.print_exc()
                                # Fallback to original file
                                return first_frame, gr.update(visible=True)
                except Exception as e:
                    print(f"⚠️  Error loading first frame: {e}")
                    import traceback
                    traceback.print_exc()
                return None, gr.update(visible=False)
            
            # Load first frame when ZIP is uploaded
            track_zip_upload.change(
                fn=load_first_frame_for_annotation,
                inputs=track_zip_upload,
                outputs=[track_first_frame_annotator, track_annot_group]
            )
            
            # Run tracking
            track_btn.click(
                fn=track_video_handler,
                inputs=[track_use_box_radio, track_first_frame_annotator, track_zip_upload, track_alpha_slider, track_vis_state],
                outputs=[track_download, track_output, track_dl_group, track_first_frame_preview, track_vis_state]
            )
            track_alpha_slider.change(
                fn=update_tracking_overlay_alpha,
                inputs=[track_alpha_slider, track_vis_state],
                outputs=track_first_frame_preview
            )

            # Clear selection
            clear_btn.click(
                fn=lambda: (None, {}),
                inputs=None,
                outputs=[track_first_frame_annotator, track_vis_state]
            )

            # Submit feedback
            def submit_user_feedback(query_id, score, comment, annot_val):
                try:
                    img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
                    bboxes = annot_val[1] if annot_val and len(annot_val) > 1 else []

                    save_feedback_to_hf(
                        query_id=query_id,
                        feedback_type=f"score_{int(score)}",
                        feedback_text=comment,
                        img_path=img_path,
                        bboxes=bboxes
                    )
                    return "✅ Feedback submitted successfully, thank you!", gr.update(visible=True)
                except Exception as e:
                    return f"❌ Submission failed: {str(e)}", gr.update(visible=True)

            submit_feedback_btn.click(
                fn=submit_user_feedback,
                inputs=[current_query_id, score_slider, feedback_box, annotator],
                outputs=[feedback_status, feedback_status]
            )
    
    

if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        ssr_mode=False,
        show_error=True,
    )