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

HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = "phoebe777777/celltool_feedback"


# ===== 清理缓存目录 =====
print("===== clearing cache =====")
# cache_path = os.path.expanduser("~/.cache/")
cache_path = os.path.expanduser("~/.cache/huggingface/gradio")
if os.path.exists(cache_path):
    try:
        shutil.rmtree(cache_path)
        # print("✅ Deleted ~/.cache/")
        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):
    """保存反馈到 Hugging Face Dataset"""
    
    # 如果没有 token，回退到本地存储
    if not HF_TOKEN:
        print("⚠️ No HF_TOKEN found, using local storage")
        save_feedback(query_id, feedback_type, feedback_text, img_path, bboxes)
        return
    
    feedback_data = {
        "query_id": query_id,
        "feedback_type": feedback_type,
        "feedback_text": feedback_text,
        "image_path": img_path,
        "bboxes": str(bboxes),  # 转为字符串
        "datetime": time.strftime("%Y-%m-%d %H:%M:%S"),
        "timestamp": time.time()
    }
    
    try:
        api = HfApi()
        
        # 创建临时文件
        filename = f"feedback_{query_id}_{int(time.time())}.json"
        
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump(feedback_data, f, indent=2, ensure_ascii=False)
        
        # 上传到 dataset
        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}")
        
    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):
    """保存用户反馈到JSON文件"""
    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):
    """解析第一个边界框"""
    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):
    """解析所有边界框"""
    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:
    """将实例掩码转换为彩色图像"""
    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]

# ===== 分割功能 =====
# @spaces.GPU
def segment_with_choice(use_box_choice, annot_value):
    """分割主函数 - 每个实例不同颜色+轮廓"""
    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}")
    box_array = None
    if use_box_choice == "Yes" and bboxes:
        # box = parse_first_bbox(bboxes)
        # if box:
        #     xmin, ymin, xmax, ymax = map(int, box)
        #     box_array = [[xmin, ymin, xmax, ymax]]
        #     print(f"📦 Using bounding box: {box_array}")
        box = parse_bboxes(bboxes)
        if box:
            box_array = box
            print(f"📦 Using bounding boxes: {box_array}")


    # 运行分割模型
    try:
        mask = run_seg(SEG_MODEL, img_path, box=box_array, device=SEG_DEVICE)
        print("📏 mask shape:", mask.shape, "dtype:", mask.dtype, "unique:", np.unique(mask))
    except Exception as e:
        print(f"❌ Inference failed: {str(e)}")
        return None, None

    # 保存原始mask为TIF文件
    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])
    print(f"✅ Instance IDs found: {unique_ids}, Total instances: {num_instances}")

    if num_instances == 0:
        print("⚠️ No instance found, returning dummy red image")
        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), None

    # ==== Color Overlay (每个实例一个颜色) ====
    overlay = img_np.copy()
    alpha = 0.5
    # cmap = cm.get_cmap("hsv", num_instances + 1)

    for inst_id in np.unique(inst_mask):
        if inst_id == 0:
            continue
        binary_mask = (inst_mask == inst_id).astype(np.uint8)
        # color = np.array(cmap(inst_id / (num_instances + 1))[:3])  # RGB only, ignore alpha
        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), temp_mask_file.name

# ===== 计数功能 =====
# @spaces.GPU
def count_cells_handler(use_box_choice, annot_value):
    """Counting handler - supports bounding box, returns only density map"""
    if annot_value is None or len(annot_value) < 1:
        return 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}")
    box_array = None
    if use_box_choice == "Yes" and bboxes:
        # box = parse_first_bbox(bboxes)
        # if box:
        #     xmin, ymin, xmax, ymax = map(int, box)
        #     box_array = [[xmin, ymin, xmax, ymax]]
        #     print(f"📦 Using bounding box: {box_array}")
        box = parse_bboxes(bboxes)
        if box:
            box_array = box
            print(f"📦 Using bounding boxes: {box_array}")
    
    try:
        print(f"🔢 Counting - Image: {image_path}")
        
        result = run_count(
            COUNT_MODEL,
            image_path,
            box=box_array,
            device=COUNT_DEVICE,
            visualize=True
        )
        
        if 'error' in result:
            return None, f"❌ Counting failed: {result['error']}"
        
        count = result['count']
        density_map = result['density_map']
        # save density map as temp file
        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

        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

        
        # Normalize density map to [0, 1]
        density_normalized = density_map.copy()
        if density_normalized.max() > 0:
            density_normalized = (density_normalized - density_normalized.min()) / (density_normalized.max() - density_normalized.min())
        
        # Apply colormap
        cmap = cm.get_cmap("jet")
        alpha = 0.3
        density_colored = cmap(density_normalized)[:, :, :3]  # RGB only, ignore alpha
        
        # Create overlay
        overlay = img_np.copy()
        
        # Blend only where density is significant (optional: threshold)
        threshold = 0.01  # Only overlay where density > 1% of max
        significant_mask = density_normalized > threshold
        
        overlay[significant_mask] = (1 - alpha) * overlay[significant_mask] + alpha * density_colored[significant_mask]
        
        # Clip and convert to uint8
        overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
        
        



        result_text = f"✅ Detected {round(count)} objects"
        if use_box_choice == "Yes" and box:
            result_text += f"\n📦 Using bounding box: {box_array}"
        

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

        return Image.fromarray(overlay), temp_density_file.name, result_text
        
        # return density_path, result_text
        
    except Exception as e:
        print(f"❌ Counting error: {e}")
        import traceback
        traceback.print_exc()
        return None, f"❌ Counting failed: {str(e)}"

# ===== Tracking Functionality =====
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

# 使用更智能的颜色分配 - 让相邻的ID颜色差异更大
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):
    """
    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")
    
    # Create color map for consistent track IDs
    # Use a colormap with many distinct colors
    # try:
    #     cmap = colormaps.get_cmap("hsv")
    # except:
    #     from matplotlib import cm
    #     cmap = cm.get_cmap("hsv")
    
    frames = []
    alpha = 0.3  # 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):
    """
    支持 ZIP 压缩包上传的 Tracking 处理函数 - 支持首帧边界框
    
    Parameters:
    -----------
    use_box_choice : str
        "Yes" or "No" - 是否使用边界框
    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, "⚠️ 请上传包含视频帧的压缩包 (.zip)", None, None
    
    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_first_bbox(bboxes)
                    # if box:
                    #     xmin, ymin, xmax, ymax = map(int, box)
                    #     box_array = [[xmin, ymin, xmax, ymax]]
                    #     print(f"📦 Using bounding box: {box_array}")
                    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
        result = run_track(
            TRACK_MODEL,
            video_dir=tif_dir,
            box=box_array,  # Pass bounding box if specified
            device=TRACK_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
            )
        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:
            result_text += f"\n📦 Using bounding box: {box_array}"

        print(f"\n✅ Tracking completed")

        # Clean up input temp directory (keep output temp for download)
        if temp_dir:
            try:
                shutil.rmtree(temp_dir)
                print(f"🗑️  Cleared input temp directory")
            except:
                pass
        
        return results_zip, result_text, gr.update(visible=True), tracking_video

    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_seg = [f for f in glob("example_imgs/seg/*")]
# ["example_imgs/seg/003_img.png", "example_imgs/seg/1977_Well_F-5_Field_1.png"]
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 =====
with gr.Blocks(
    title="Microscopy Analysis Suite", 
    theme=gr.themes.Soft(),
    css="""
    .tabs button {
        font-size: 18px !important;
        font-weight: 600 !important;
        padding: 12px 20px !important;
    }
    .uniform-height {
        height: 500px !important;
        display: flex !important;
        align-items: center !important;
        justify-content: center !important;
    }

    .uniform-height img,
    .uniform-height canvas {
        max-height: 500px !important;
        object-fit: contain !important;
    }

    /* 强制密度图容器和图片高度 */
    #density_map_output {
        height: 500px !important;
    }
    
    #density_map_output .image-container {
        height: 500px !important;
    }
    
    #density_map_output img {
        height: 480px !important;
        width: auto !important;
        max-width: 90% !important;
        object-fit: contain !important;
    }
    """
    ) as demo:
    gr.Markdown(
        """
        # 🔬 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
        """
    )
    
    # 全局状态
    current_query_id = gr.State(str(uuid.uuid4()))
    user_uploaded_examples = gr.State(example_images_seg.copy())  # 初始化时包含原始示例
    
    with gr.Tabs():
        # ===== Tab 1: Segmentation =====
        with gr.Tab("🎨 Segmentation"):
            gr.Markdown("## Instance Segmentation of Microscopic Objects")
            gr.Markdown(
                        """
                        **Instructions:**
                        1. Upload an image or select an example image (supports various formats: .png, .jpg, .tif)
                        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, download the original predicted mask (.tif format); if needed, click "Clear Selection" to choose a new image
                        
                        🤘 Rate and submit feedback to help us improve the model!
                        """
                    )
            
            with gr.Row():
                with gr.Column(scale=1):
                    annotator = BBoxAnnotator(
                        label="🖼️ Upload Image (Optional: Provide a Bounding Box)",
                        categories=["cell"],
                    )
                    
                    # Example Images Gallery
                    example_gallery = gr.Gallery(
                        label="📁 Example Image Gallery",
                        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
                    image_uploader = gr.Image(
                        label="➕ Upload New Example Image to Gallery",
                        type="filepath"
                    )


                with gr.Column(scale=2):
                    seg_output = gr.Image(
                        type="pil",
                        label="📸 Segmentation Result",
                        elem_classes="uniform-height"
                    )
                    
                    # Download Original Prediction
                    download_mask_btn = gr.File(
                        label="📥 Download Original Prediction (.tif format)",
                        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
                    )
            
            # 绑定事件: 运行分割
            run_seg_btn.click(
                fn=segment_with_choice,
                inputs=[use_box_radio, annotator],
                outputs=[seg_output, download_mask_btn]
            )

            # 清空按钮事件
            clear_btn.click(
                fn=lambda: None,
                inputs=None,
                outputs=annotator
            )
            
            # 初始化Gallery显示
            demo.load(
                fn=lambda: example_images_seg.copy(),
                outputs=example_gallery
            )
            
            # 绑定事件: 上传示例图片
            def add_to_gallery(img_path, current_imgs):
                if not img_path:
                    return current_imgs
                try:
                    if img_path not in current_imgs:
                        current_imgs.append(img_path)
                    return current_imgs
                except:
                    return current_imgs
            
            image_uploader.change(
                fn=add_to_gallery,
                inputs=[image_uploader, user_uploaded_examples],
                outputs=user_uploaded_examples
            ).then(
                fn=lambda imgs: imgs,
                inputs=user_uploaded_examples,
                outputs=example_gallery
            )
            
            # 绑定事件: 点击Gallery加载
            def load_from_gallery(evt: gr.SelectData, all_imgs):
                if evt.index is not None and evt.index < len(all_imgs):
                    return all_imgs[evt.index]
                return None
            
            example_gallery.select(
                fn=load_from_gallery,
                inputs=user_uploaded_examples,
                outputs=annotator
            )
            
            # 绑定事件: 提交反馈
            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
                    # )
                    # 使用 HF 存储
                    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. Upload an image or select an example image (supports multiple formats: .png, .jpg, .tif)
                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, download the original prediction (.npy format); if needed, click "Clear Selection" to choose a new image to run
                
                🤘 Rate and submit feedback to help us improve the model!
                """
            )
            
            with gr.Row():
                with gr.Column(scale=1):
                    count_annotator = BBoxAnnotator(
                        label="🖼️ Upload Image (Optional: Provide a Bounding Box)",
                        categories=["cell"],
                    )
                    
                    # Example gallery with "add" functionality
                    with gr.Row():
                        count_example_gallery = gr.Gallery(
                            label="📁 Example Image Gallery",
                            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
                    with gr.Row():
                        count_image_uploader = gr.File(
                            label="➕ Add Example Image to Gallery",
                            file_types=["image"],
                            type="filepath"
                        )

                
                with gr.Column(scale=2):
                    count_output = gr.Image(
                        label="📸 Density Map",
                        type="filepath",
                        elem_id="density_map_output"
                        
                    )
                    count_status = gr.Textbox(
                        label="📊 Statistics",
                        lines=2
                    )
                    download_density_btn = gr.File(
                        label="📥 Download Original Prediction (.npy format)",
                        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())
            
            # Function to add image to gallery
            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
                
                try:
                    # Add new image path to list
                    if new_img_file not in current_imgs:
                        current_imgs.append(new_img_file)
                        print(f"✅ Added image to gallery: {new_img_file}")
                except Exception as e:
                    print(f"⚠️  Failed to add image: {e}")
                
                return current_imgs, current_imgs
            
            # When user uploads a new image file
            count_image_uploader.upload(
                fn=add_to_count_gallery,
                inputs=[count_image_uploader, count_user_examples],
                outputs=[count_user_examples, count_example_gallery]
            )
            
            # When user selects from gallery, load into annotator
            def load_from_count_gallery(evt: gr.SelectData, all_imgs):
                """Load selected image from gallery into annotator"""
                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 selected_img
                return None
            
            count_example_gallery.select(
                fn=load_from_count_gallery,
                inputs=count_user_examples,
                outputs=count_annotator
            )
            
            # Run counting
            count_btn.click(
                fn=count_cells_handler,
                inputs=[count_use_box_radio, count_annotator],
                outputs=[count_output, download_density_btn, count_status]
            )

            # 清空按钮事件
            clear_btn.click(
                fn=lambda: None,
                inputs=None,
                outputs=count_annotator
            )

            # 绑定事件: 提交反馈
            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
                    # )
                    # 使用 HF 存储
                    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. Upload a ZIP file or select from the example library. 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. Download the CTC format results; if needed, click "Clear Selection" to choose a new ZIP file to run
                        
                        🤘 Rate and submit feedback to help us improve the model!

                        """
                    )
            
            with gr.Row():
                with gr.Column(scale=1):
                    track_zip_upload = gr.File(
                        label="📦 Upload Image Sequence in ZIP File",
                        file_types=[".zip"]
                    )

                    # First frame annotation for bounding box
                    track_first_frame_annotator = BBoxAnnotator(
                        label="🖼️ (Optional) First Frame Bounding Box Annotation",
                        categories=["cell"],
                        visible=False,  # Hidden initially
                    )

                    # Example ZIP gallery
                    track_example_gallery = gr.Gallery(
                        label="📁 Example Video Gallery (Click to Select)",
                        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
                    track_gallery_upload = gr.File(
                        label="➕ Add ZIP to Example Gallery",
                        file_types=[".zip"],
                        type="filepath"
                    )
                
                with gr.Column(scale=2):
                    track_first_frame_preview = gr.Image(
                        label="📸 Tracking Visualization",
                        type="filepath",
                        # height=400,
                        elem_classes="uniform-height",
                        interactive=False
                    )
                    
                    track_output = gr.Textbox(
                        label="📊 Tracking Information",
                        lines=8,
                        interactive=False
                    )
                    
                    track_download = gr.File(
                        label="📥 Download Tracking Results (CTC Format)",
                        visible=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_first_frame_annotator]
            )
            
            # Run tracking
            track_btn.click(
                fn=track_video_handler,
                inputs=[track_use_box_radio, track_first_frame_annotator, track_zip_upload],
                outputs=[track_download, track_output, track_download, track_first_frame_preview]
            )

            # 清空按钮事件
            clear_btn.click(
                fn=lambda: None,
                inputs=None,
                outputs=track_first_frame_annotator
            )

            # 绑定事件: 提交反馈
            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
                    # )
                    # 使用 HF 存储
                    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]
            )
    
    gr.Markdown(
        """
        ---
        ### 💡 Technical Details

        **MicroscopyMatching** - A general-purpose microscopy image analysis toolkit based on Stable Diffusion
        """
    )

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