Update app.py

app.py

@@ -289,6 +289,7 @@ import tempfile
 from skimage import measure
 from matplotlib import cm
 
+
 # ===== 导入三个推理模块 =====
 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
@@ -474,61 +475,32 @@ def segment_with_choice(use_box_choice, annot_value, mode="Overlay"):
     return Image.fromarray(overlay), result_text
 
 # ===== 计数功能 =====
-def count_cells_handler(image_path):
-    """计数处理函数"""
-    if image_path is None:
-        return None, "⚠️ 请先上传图像"
-    
-    if COUNT_MODEL is None:
-        return None, "❌ 计数模型未加载"
-    
-    try:
-        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)}"
+import zipfile
+import tempfile
+import shutil
 
-# ===== 跟踪功能 =====
-def track_video_handler(video_dir_input):
-    """跟踪处理函数"""
-    if video_dir_input is None or video_dir_input.strip() == "":
-        return None, "⚠️ 请输入视频帧目录路径"
+def track_video_handler(zip_file_obj):
+    """支持 ZIP 压缩包上传的 Tracking 处理函数"""
+    if zip_file_obj is None:
+        return None, "⚠️ 请上传包含视频帧的压缩包 (.zip)"
     
     if TRACK_MODEL is None:
         return None, "❌ 跟踪模型未加载"
     
-    # 检查目录是否存在
-    if not os.path.exists(video_dir_input):
-        return None, f"❌ 目录不存在: {video_dir_input}"
-    
     try:
-        print(f"🎬 Tracking - Video dir: {video_dir_input}")
+        # 创建临时目录
+        temp_dir = tempfile.mkdtemp()
+        print(f"📦 解压到临时目录: {temp_dir}")
+
+        # 解压 zip 文件
+        with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:
+            zip_ref.extractall(temp_dir)
+
+        print(f"🎬 Tracking - Video frame folder: {temp_dir}")
         
         result = run_track(
             TRACK_MODEL,
-            video_dir=video_dir_input,
+            video_dir=temp_dir,   # 传入解压目录
             box=None,
             device=TRACK_DEVICE,
             output_dir="tracked_results"
@@ -549,13 +521,15 @@ def track_video_handler(video_dir_input):
 - res_track.txt (CTC格式轨迹)
 - 其他跟踪数据文件
 """
-        
+
         print(f"✅ Tracking done - {num_tracks} tracks")
-        
+
         return None, result_text
-        
+
+    except zipfile.BadZipFile:
+        return None, "❌ 上传的文件不是有效的 ZIP 压缩包"
+
     except Exception as e:
-        print(f"❌ Tracking error: {e}")
         import traceback
         traceback.print_exc()
         return None, f"❌ 跟踪失败: {str(e)}"
@@ -673,55 +647,50 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
         
         # ===== Tab 3: Tracking =====
         with gr.Tab("🎬 跟踪 (Tracking)"):
-            gr.Markdown("## 视频细胞跟踪 - 时间序列分析")
-            
-            with gr.Row():
-                with gr.Column(scale=1):
-                    track_input = gr.Textbox(
-                        label="📁 视频帧目录路径",
-                        placeholder="例如: example_imgs/2D+Time/Fluo-N2DL-HeLa/train/02/",
-                        lines=2
-                    )
-                    track_btn = gr.Button("▶️ 运行跟踪", variant="primary", size="lg")
-                    
-                    gr.Markdown(
-                        """
-                        **使用说明:**
-                        1. 输入包含视频帧序列的目录路径
-                        2. 目录应包含: t000.tif, t001.tif, ...
-                        3. 点击 "运行跟踪"
-                        4. 结果将保存到 `tracked_results/` 目录
-                        
-                        **输入格式:**
-                        ```
-                        video_dir/
-                        ├── t000.tif
-                        ├── t001.tif
-                        ├── t002.tif
-                        └── ...
-                        ```
-                        
-                        **跟踪模式:** Greedy (快速)
-                        """
-                    )
+        gr.Markdown("## 视频细胞跟踪 - 支持 ZIP 压缩包上传")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                track_zip_upload = gr.File(
+                    label="📦 上传视频帧 ZIP 文件",
+                    file_types=[".zip"]
+                )
+                track_btn = gr.Button("▶️ 运行跟踪", variant="primary", size="lg")
                 
-                with gr.Column(scale=2):
-                    track_output = gr.Video(
-                        label="📹 跟踪结果视频 (暂不支持)",
-                        visible=False
-                    )
-                    track_status = gr.Textbox(
-                        label="📊 跟踪信息",
-                        lines=10
-                    )
+                gr.Markdown(
+                    """
+                    **使用说明:**
+                    1. 上传包含视频帧序列的压缩包 `.zip`
+                    2. 压缩包应直接包含 `.tif` 格式图像，如 t000.tif, t001.tif, ...
+                    3. 点击 "运行跟踪"
+                    4. 结果将保存到 `tracked_results/` 目录
+                    
+                    **压缩包示例结构:**
+                    ```
+                    frames.zip
+                    ├── t000.tif
+                    ├── t001.tif
+                    ├── t002.tif
+                    └── ...
+                    ```
+                    
+                    **跟踪模式:** Greedy (快速)
+                    """
+                )
             
-            # 绑定事件
-            track_btn.click(
-                fn=track_video_handler,
-                inputs=track_input,
-                outputs=[track_output, track_status]
-            )
-
+            with gr.Column(scale=2):
+                track_output = gr.Textbox(
+                    label="📊 跟踪信息",
+                    lines=12,
+                    interactive=False
+                )
+        
+        # 绑定事件：上传zip → 解压 → Tracking
+        track_btn.click(
+            fn=track_video_handler,               # 你刚才改好的函数
+            inputs=track_zip_upload,              # 文件上传
+            outputs=[None, track_output]          # 第二个是 Textbox 输出
+        )
     gr.Markdown(
         """
         ---