Update app.py

app.py

@@ -5,7 +5,6 @@ import numpy as np
 import torch
 import os
 import shutil
-import subprocess
 import time
 import json
 import uuid
@@ -26,9 +25,9 @@ cache_path = os.path.expanduser("~/.cache")
 if os.path.exists(cache_path):
     try:
         shutil.rmtree(cache_path)
-        print("✅ Deleted ~/.cache to free space.")
-    except Exception as e:
-        print(f"⚠️ Could not delete cache: {e}")
+        print("✅ Deleted ~/.cache")
+    except:
+        pass
 
 # ===== 全局模型变量 =====
 SEG_MODEL = None
@@ -46,19 +45,16 @@ def load_all_models():
     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)
@@ -68,22 +64,8 @@ def load_all_models():
     print("✅ All Models Loaded Successfully")
     print("="*60)
 
-# 启动时加载所有模型
 load_all_models()
 
-# ===== BBox 解析 =====
-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
-
 # ===== 保存用户反馈 =====
 DATASET_DIR = Path("solver_cache")
 DATASET_DIR.mkdir(parents=True, exist_ok=True)
@@ -100,6 +82,7 @@ def save_feedback(query_id, feedback_type, feedback_text=None, img_path=None, bb
     }
     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)
@@ -109,140 +92,168 @@ def save_feedback(query_id, feedback_type, feedback_text=None, img_path=None, bb
             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)
 
-# ===== 彩色 mask 可视化 =====
+# ===== 辅助函数 =====
+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 colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
-    mask = mask.astype(np.int32)
-
-    def hsv_to_rgb(hh, ss, vv):
-        i = int(hh * 6.0)
-        f = hh * 6.0 - i
-        p = vv * (1.0 - ss)
-        q = vv * (1.0 - f * ss)
-        t = vv * (1.0 - (1.0 - f) * ss)
+    """将实例掩码转换为彩色图像"""
+    def hsv_to_rgb(h, s, v):
+        i = int(h * 6.0)
+        f = h * 6.0 - i
         i = i % 6
-        if i == 0: r, g, b = vv, t, p
-        elif i == 1: r, g, b = q, vv, p
-        elif i == 2: r, g, b = p, vv, t
-        elif i == 3: r, g, b = p, q, vv
-        elif i == 4: r, g, b = t, p, vv
-        else: r, g, b = vv, p, q
-        return int(r*255), int(g*255), int(b*255)
+        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 k in range(1, num_colors):
-        hue = (k % num_colors) / float(num_colors)
-        palette.append(hsv_to_rgb(hue, 1.0, 0.95))
+    for i in range(1, num_colors):
+        h = (i % num_colors) / float(num_colors)
+        palette.append(hsv_to_rgb(h, 1.0, 0.95))
 
-    color_idx = mask % num_colors
     palette_arr = np.array(palette, dtype=np.uint8)
+    color_idx = mask % num_colors
     return palette_arr[color_idx]
 
-# ===== 推理 + 实例彩色可视化 =====
+def overlay_instances(img, mask, alpha=0.5, cmap_name="tab20"):
+    """叠加实例颜色"""
+    img = img.astype(np.float32)
+    if len(img.shape) == 2:
+        img = np.stack([img]*3, axis=-1)
+    if img.max() > 1.5:
+        img = img / 255.0
+
+    overlay = img.copy()
+    cmap = cm.get_cmap(cmap_name, np.max(mask) + 1)
+
+    for inst_id in np.unique(mask):
+        if inst_id == 0:
+            continue
+        color = np.array(cmap(inst_id)[:3])
+        overlay[mask == inst_id] = (1 - alpha) * overlay[mask == inst_id] + alpha * color
+
+    return overlay
+
+# ===== 分割功能 =====
 def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
+    """分割主函数"""
     if annot_value is None or len(annot_value) < 1:
-        print("❌ No annotation input")
-        return None, "❌ 没有输入图像"
+        return None, "⚠️ 请上传图像"
 
     img_path = annot_value[0]
     bboxes = annot_value[1] if len(annot_value) > 1 else []
 
-    print(f"🖼️ Image path: {img_path}")
+    print(f"🖼️ 图像路径: {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 box: {box_array}")
+            print(f"📦 使用边界框: {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))
+        print("📏 mask shape:", mask.shape, "unique ids:", np.unique(mask))
     except Exception as e:
-        print(f"❌ Error during inference: {e}")
         return None, f"❌ 推理失败: {str(e)}"
 
+    # 读取原图
     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, f"❌ 无法打开图像: {str(e)}"
-
-    try:
-        img_rgb = img.convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
-        img_np = np.array(img_rgb, dtype=np.float32)
+        img = Image.open(img_path).convert("RGB").resize(mask.shape[::-1], resample=Image.BILINEAR)
+        img_np = np.array(img).astype(np.float32)
         if img_np.max() > 1.5:
-            img_np = img_np / 255.0
+            img_np /= 255.0
     except Exception as e:
-        print(f"❌ Error in image conversion/resizing: {e}")
-        return None, f"❌ 图像转换失败: {str(e)}"
+        return None, f"❌ 图像读取失败: {str(e)}"
 
-    mask_np = np.array(mask)
-    inst_mask = mask_np.astype(np.int32)
+    inst_mask = mask.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}")
+    print(f"✅ 实例数量: {num_instances}")
 
     if num_instances == 0:
-        print("⚠️ No instance found, returning dummy red image")
-        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), "⚠️ 未检测到任何实例"
-
-    # ==== Color Overlay (每个实例一个颜色) ====
-    overlay = img_np.copy()
-    alpha = 0.5
-    cmap = cm.get_cmap("nipy_spectral", 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
-        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)
-            overlay[contour[:, 0], contour[:, 1]] = [1.0, 1.0, 0.0]  # 黄色轮廓
-
-    overlay = np.clip(overlay * 255.0, 0, 255).astype(np.uint8)
-
-    status_msg = f"✅ 分割完成! 检测到 {num_instances} 个实例"
-    
-    if mode == "Instance Mask Only":
-        return Image.fromarray(colorize_mask(inst_mask, num_colors=512)), status_msg
-
-    return Image.fromarray(overlay), status_msg
-
-# ===== Count Handler =====
-def count_cells_handler(input_image):
-    if input_image is None:
-        return None, "❌ 请先上传图像"
+        return Image.new("RGB", mask.shape[::-1], (255, 0, 0)), "⚠️ 未检测到实例"
+
+    # 可视化
+    if mode == "Overlay":
+        overlay = overlay_instances(img_np, inst_mask, alpha=0.5, cmap_name="tab20")
+        overlay_img = Image.fromarray((overlay * 255).astype(np.uint8))
+        return overlay_img, f"✅ 检测到 {num_instances} 个细胞"
+    elif mode == "Instance Mask Only":
+        color_mask = colorize_mask(inst_mask, num_colors=512)
+        return Image.fromarray(color_mask), f"✅ 检测到 {num_instances} 个细胞"
+
+    return None, "❓ 无效显示模式"
+
+# ===== 计数功能 =====
+def count_cells_handler(image_path):
+    """计数处理函数"""
+    if image_path is None:
+        return None, "⚠️ 请先上传图像"
     
     try:
-        result_image, cell_count = run_count(COUNT_MODEL, input_image, device=COUNT_DEVICE)
-        status = f"✅ 计数完成! 检测到 {cell_count} 个细胞"
-        return result_image, status
+        print(f"🔢 Counting - Image: {image_path}")
+        
+        result = run_count(
+            COUNT_MODEL,
+            image_path,
+            box=None,
+            device=COUNT_DEVICE,
+            visualize=True
+        )
+        
+        if 'error' in result:
+            return None, f"❌ 计数失败: {result['error']}"
+        
+        count = result['count']
+        viz_path = result['visualized_path']
+        result_text = f"✅ 检测到 {count:.1f} 个细胞"
+        
+        print(f"✅ Counting done - Count: {count:.1f}")
+        
+        return viz_path, result_text
+        
     except Exception as e:
+        print(f"❌ Counting error: {e}")
         import traceback
         traceback.print_exc()
         return None, f"❌ 计数失败: {str(e)}"
 
-# ===== Track Handler =====
+# ===== 跟踪功能 =====
 def find_tif_dir(root_dir):
-    for dirpath, dirnames, filenames in os.walk(root_dir):
-        tif_files = [f for f in filenames if f.lower().endswith(('.tif', '.tiff'))]
-        if tif_files:
+    """递归查找第一个包含 .tif 文件的目录"""
+    for dirpath, _, filenames in os.walk(root_dir):
+        if any(f.lower().endswith('.tif') for f in filenames):
             return dirpath
     return None
 
 def track_video_handler(zip_file_obj):
+    """支持 ZIP 压缩包上传的 Tracking 处理函数"""
     if zip_file_obj is None:
-        return None, "❌ 请先上传 ZIP 文件"
+        return None, "⚠️ 请上传包含视频帧的压缩包 (.zip)"
     
     try:
         temp_dir = tempfile.mkdtemp()
@@ -253,9 +264,9 @@ def track_video_handler(zip_file_obj):
 
         tif_dir = find_tif_dir(temp_dir)
         if tif_dir is None:
-            return None, f"❌ 跟踪失败: 解压后未找到任何 .tif 图像"
+            return None, "❌ 解压后未找到任何 .tif 图像"
         
-        print(f"🎬 Tracking - Found .tif in folder: {tif_dir}")
+        print(f"🎬 Tracking - Found .tif in: {tif_dir}")
 
         result = run_track(
             TRACK_MODEL,
@@ -282,23 +293,17 @@ def track_video_handler(zip_file_obj):
 """
 
         print(f"✅ Tracking done - {num_tracks} tracks")
-
         return None, result_text
 
     except zipfile.BadZipFile:
         return None, "❌ 上传的文件不是有效的 ZIP 压缩包"
-
     except Exception as e:
         import traceback
         traceback.print_exc()
         return None, f"❌ 跟踪失败: {str(e)}"
 
-# ===== 示例图像数据 =====
-example_data = [
-    ("003_img.png", [(50, 60, 120, 150, "cell")]),
-    ("1977_Well_F-5_Field_1.png", [(30, 40, 100, 130, "cell")]),
-]
-gallery_images = [p for p, _ in example_data]
+# ===== 示例图像 =====
+example_images = ["003_img.png", "1977_Well_F-5_Field_1.png"]
 
 # ===== Gradio UI =====
 with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as demo:
@@ -313,9 +318,9 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
         """
     )
     
-    # 全局状态: 用于存储当前query_id和用户上传的示例图片
+    # 全局状态
     current_query_id = gr.State(str(uuid.uuid4()))
-    user_uploaded_images = gr.State([])
+    user_uploaded_examples = gr.State(example_images.copy())  # 初始化时包含原始示例
     
     with gr.Tabs():
         # ===== Tab 1: Segmentation =====
@@ -329,18 +334,17 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         categories=["cell"]
                     )
                     
-                    # 示例图片展示
+                    # 示例图片Gallery
                     example_gallery = gr.Gallery(
-                        value=gallery_images,
                         label="📁 示例图片",
-                        columns=[3],
+                        columns=3,
                         object_fit="cover",
-                        height=128
+                        height=150
                     )
                     
-                    # 用户上传示例图片
+                    # 上传示例图片
                     image_uploader = gr.Image(
-                        label="➕ 上传新示例图片到Gallery",
+                        label="➕ 上传新示例到Gallery",
                         type="filepath"
                     )
                     
@@ -361,7 +365,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                     gr.Markdown(
                         """
                         **使用说明:**
-                        1. 上传图像或选择示例图片
+                        1. 上传图像或从Gallery选择示例
                         2. (可选) 标注边界框并选择 "Yes"
                         3. 选择显示模式
                         4. 点击 "运行分割"
@@ -380,25 +384,26 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                     )
                     
                     # 满意度评分
-                    score = gr.Slider(
-                        1, 5,
+                    score_slider = gr.Slider(
+                        minimum=1,
+                        maximum=5,
                         step=1,
                         value=3,
                         label="🌟 满意度评分 (1-5)"
                     )
                     
                     # 反馈文本框
-                    comment_box = gr.Textbox(
+                    feedback_box = gr.Textbox(
                         placeholder="请输入您的反馈意见...",
                         lines=2,
                         label="💬 反馈意见"
                     )
                     
-                    # 提交评分按钮
-                    submit_score = gr.Button("💾 提交评分", variant="secondary")
+                    # 提交按钮
+                    submit_feedback_btn = gr.Button("💾 提交反馈", variant="secondary")
                     
                     feedback_status = gr.Textbox(
-                        label="✅ 反馈提交状态",
+                        label="✅ 提交状态",
                         lines=1,
                         visible=False
                     )
@@ -410,56 +415,55 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                 outputs=[seg_output, seg_status]
             )
             
-            # 绑定事件: 上传示例图片到Gallery
-            def add_uploaded_image(img_path, current_gallery):
+            # 初始化Gallery显示
+            demo.load(
+                fn=lambda: example_images.copy(),
+                outputs=example_gallery
+            )
+            
+            # 绑定事件: 上传示例图片
+            def add_to_gallery(img_path, current_imgs):
                 if not img_path:
-                    return current_gallery
+                    return current_imgs
                 try:
-                    img = Image.open(img_path)
-                    img.thumbnail((128, 128))
-                    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-                    img.save(temp_file.name, format="PNG")
-                    thumb_path = temp_file.name
-                    if thumb_path not in current_gallery:
-                        current_gallery.append(thumb_path)
-                    return current_gallery
-                except Exception as e:
-                    print(f"❌ Failed to add image to gallery: {e}")
-                    return current_gallery
+                    if img_path not in current_imgs:
+                        current_imgs.append(img_path)
+                    return current_imgs
+                except:
+                    return current_imgs
             
             image_uploader.change(
-                fn=add_uploaded_image,
-                inputs=[image_uploader, user_uploaded_images],
-                outputs=user_uploaded_images
+                fn=add_to_gallery,
+                inputs=[image_uploader, user_uploaded_examples],
+                outputs=user_uploaded_examples
             ).then(
                 fn=lambda imgs: imgs,
-                inputs=user_uploaded_images,
+                inputs=user_uploaded_examples,
                 outputs=example_gallery
             )
             
-            # 绑定事件: 点击Gallery图片加载到annotator
-            def load_example(evt: gr.SelectData, examples):
-                if evt.index is not None and evt.index < len(examples):
-                    img_path = examples[evt.index]
-                    return img_path
+            # 绑定事件: 点击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_example,
-                inputs=user_uploaded_images,
+                fn=load_from_gallery,
+                inputs=user_uploaded_examples,
                 outputs=annotator
             )
             
-            # 绑定事件: 提交评分
-            def submit_feedback(query_id, score_val, comment_text, annot_value):
+            # 绑定事件: 提交反馈
+            def submit_user_feedback(query_id, score, comment, annot_val):
                 try:
-                    img_path = annot_value[0] if annot_value and len(annot_value) > 0 else None
-                    bboxes = annot_value[1] if annot_value and len(annot_value) > 1 else []
+                    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_val)}",
-                        feedback_text=comment_text,
+                        feedback_type=f"score_{int(score)}",
+                        feedback_text=comment,
                         img_path=img_path,
                         bboxes=bboxes
                     )
@@ -467,9 +471,9 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                 except Exception as e:
                     return f"❌ 提交失败: {str(e)}", gr.update(visible=True)
             
-            submit_score.click(
-                fn=submit_feedback,
-                inputs=[current_query_id, score, comment_box, annotator],
+            submit_feedback_btn.click(
+                fn=submit_user_feedback,
+                inputs=[current_query_id, score_slider, feedback_box, annotator],
                 outputs=[feedback_status, feedback_status]
             )
         
@@ -491,17 +495,12 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         1. 上传细胞图像
                         2. 点击 "运行计数"
                         3. 查看密度图和计数结果
-                        
-                        **特点:**
-                        - 基于 Stable Diffusion 特征
-                        - 自动生成密度图
-                        - 无需手动标注
                         """
                     )
                 
                 with gr.Column(scale=2):
                     count_output = gr.Image(
-                        label="📸 计数结果 (左: 原图 | 右: 密度图)",
+                        label="📸 计数结果",
                         type="filepath",
                         height=500
                     )
@@ -510,7 +509,6 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         lines=2
                     )
             
-            # 绑定事件
             count_btn.click(
                 fn=count_cells_handler,
                 inputs=count_input,
@@ -532,21 +530,9 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                     gr.Markdown(
                         """
                         **使用说明:**
-                        1. 上传包含视频帧序列的压缩包 `.zip`
-                        2. 压缩包应直接包含 `.tif` 格式图像,如 t000.tif, t001.tif, ...
-                        3. 点击 "运行跟踪"
-                        4. 结果将保存到 `tracked_results/` 目录
-                        
-                        **压缩包示例结构:**
-                        ```
-                        frames.zip
-                        ├── t000.tif
-                        ├── t001.tif
-                        ├── t002.tif
-                        └── ...
-                        ```
-                        
-                        **跟踪模式:** Greedy (快速)
+                        1. 上传包含 `.tif` 图像的 ZIP 压缩包
+                        2. 点击 "运行跟踪"
+                        3. 结果保存到 `tracked_results/` 目录
                         """
                     )
                 
@@ -557,12 +543,11 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         interactive=False
                     )
             
-            # 绑定事件:上传zip → 解压 → Tracking
-            dummy_output = gr.Textbox(visible=False)
+            dummy = gr.Textbox(visible=False)
             track_btn.click(
                 fn=track_video_handler,
                 inputs=track_zip_upload,
-                outputs=[dummy_output, track_output]
+                outputs=[dummy, track_output]
             )
     
     gr.Markdown(
@@ -570,20 +555,9 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
         ---
         ### 💡 技术说明
         
-        **分割 (Segmentation)**
-        - 模型: 基于 Stable Diffusion 特征的实例分割
-        - 输出: 每个细胞一个唯一颜色的掩码
-        
-        **计数 (Counting)**
-        - 模型: 密度图估计
-        - 输出: 密度热力图 + 总计数
-        
-        **跟踪 (Tracking)**
-        - 模型: Trackastra 跟踪算法
-        - 输出: CTC 格式的轨迹文件
-        
-        ---
-        📧 问题反馈 | 🌟 GitHub
+        **分割 (Segmentation)** - 基于 Stable Diffusion 特征的实例分割  
+        **计数 (Counting)** - 密度图估计  
+        **跟踪 (Tracking)** - Trackastra 跟踪算法
         """
     )