add tracking display

app.py

@@ -13,6 +13,8 @@ import tempfile
 import zipfile
 from skimage import measure
 from matplotlib import cm
+from glob import glob
+from natsort import natsorted
 
 # ===== 导入三个推理模块 =====
 from inference_seg import load_model as load_seg_model, run as run_seg
@@ -353,55 +355,447 @@ def find_tif_dir(root_dir):
             return dirpath
     return None
 
-def track_video_handler(zip_file_obj):
-    """支持 ZIP 压缩包上传的 Tracking 处理函数"""
+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):
+    """递归查找第一个包含有效 .tif 文件的目录"""
+    for dirpath, dirnames, filenames in os.walk(root_dir):
+        if '__MACOSX' in dirpath:
+            continue
+        
+        potential_tifs = [
+            os.path.join(dirpath, f) 
+            for f in filenames 
+            if f.lower().endswith(('.tif', '.tiff')) and not f.startswith('._')
+        ]
+        
+        if not potential_tifs:
+            continue
+        
+        valid_tifs = [f for f in potential_tifs if is_valid_tiff(f)]
+        
+        if valid_tifs:
+            print(f"✅ Found {len(valid_tifs)} valid TIFF files in: {dirpath}")
+            return dirpath
+    
+    return None
+
+def create_ctc_results_zip(output_dir):
+    """
+    Create a ZIP file with CTC format results
+    
+    Parameters:
+    -----------
+    output_dir : str
+        Directory containing tracking results (res_track.txt, etc.)
+    
+    Returns:
+    --------
+    zip_path : str
+        Path to created ZIP file
+    """
+    # Create temp directory for ZIP
+    temp_zip_dir = tempfile.mkdtemp()
+    zip_filename = f"tracking_results_{time.strftime('%Y%m%d_%H%M%S')}.zip"
+    zip_path = os.path.join(temp_zip_dir, zip_filename)
+    
+    print(f"📦 Creating results ZIP: {zip_path}")
+    
+    # Create ZIP with all tracking results
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        # Add all files from output directory
+        for root, dirs, files in os.walk(output_dir):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, output_dir)
+                zipf.write(file_path, arcname)
+                print(f"  📄 Added: {arcname}")
+        
+        # Add a README with summary
+        readme_content = f"""Tracking Results Summary
+            ========================
+
+            Generated: {time.strftime('%Y-%m-%d %H:%M:%S')}
+
+            Files:
+            ------
+            - res_track.txt: CTC format tracking data
+            Format: track_id start_frame end_frame parent_id
+            
+            - Segmentation masks
+
+            For more information on CTC format:
+            http://celltrackingchallenge.net/
+            """
+        zipf.writestr("README.txt", readme_content)
+    
+    print(f"✅ ZIP created: {zip_path} ({os.path.getsize(zip_path) / 1024:.1f} KB)")
+    return zip_path
+
+def 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("nipy_spectral")
+    except:
+        from matplotlib import cm
+        cmap = cm.get_cmap("nipy_spectral")
+    
+    frames = []
+    alpha = 0.5  # 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])
+                
+                # 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]
+
+
+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)"
+        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"📦 使用边界框: {box_array}")
+        
+        # Extract input ZIP
         temp_dir = tempfile.mkdtemp()
-        print(f"📦 解压到临时目录: {temp_dir}")
+        print(f"\n📦 解压到临时目录: {temp_dir}")
 
         with zipfile.ZipFile(zip_file_obj.name, 'r') as zip_ref:
-            zip_ref.extractall(temp_dir)
-
-        tif_dir = find_tif_dir(temp_dir)
+            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_count} 文件, 跳过: {skipped_count} 文件")
+        
+        # Find valid TIFF directory
+        tif_dir = find_valid_tif_dir(temp_dir)
+        
         if tif_dir is None:
-            return None, "❌ 解压后未找到任何 .tif 图像"
+            return None, "❌ 未找到有效的TIFF文件", None, None
+        
+        # Validate TIFF files
+        tif_files = sorted(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, "❌ 没有有效的TIFF文件", None, None
+        
+        print(f"📈 使用 {len(valid_tif_files)} 个TIFF文件")
         
-        print(f"🎬 Tracking - Found .tif in: {tif_dir}")
+        # 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结果将保存到: {output_temp_dir}")
+
+        # Run tracking with optional bounding box
         result = run_track(
             TRACK_MODEL,
             video_dir=tif_dir,
-            box=None,
+            box=box_array,  # Pass bounding box if specified
             device=TRACK_DEVICE,
-            output_dir="tracked_results"
+            output_dir=output_temp_dir
         )
         
         if 'error' in result:
-            return None, f"❌ 跟踪失败: {result['error']}"
+            return None, f"❌ 跟踪失败: {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
         
-        output_dir = result['output_dir']
+        bbox_info = ""
+        if box_array:
+            bbox_info = f"\n🔲 使用边界框: [{box_array[0][0]}, {box_array[0][1]}, {box_array[0][2]}, {box_array[0][3]}]"
         
         result_text = f"""✅ 跟踪完成!
 
-        📁 结果保存在: {output_dir}
+🖼️  处理帧数: {len(valid_tif_files)}{bbox_info}
 
-        包含的文件:
-        - res_track.txt (CTC格式轨迹)
-        - 其他跟踪数据文件
-        """
-
-        print(f"✅ Tracking done")
-        return None, result_text
+📥 点击下方按钮下载CTC格式结果
+结果包含:
+- res_track.txt (CTC格式轨迹数据)
+- 其他跟踪相关文件
+- README.txt (结果说明)
+"""
+        
+        print(f"\n✅ Tracking完成")
+        
+        # Clean up input temp directory (keep output temp for download)
+        if temp_dir:
+            try:
+                shutil.rmtree(temp_dir)
+                print(f"🗑️  清理输入临时目录")
+            except:
+                pass
+        
+        return results_zip, result_text, gr.update(visible=True), tracking_video
 
     except zipfile.BadZipFile:
-        return None, "❌ 上传的文件不是有效的 ZIP 压缩包"
+        return None, "❌ 不是有效的ZIP文件", None, None
     except Exception as e:
         import traceback
         traceback.print_exc()
-        return None, f"❌ 跟踪失败: {str(e)}"
+        
+        # Clean up on error
+        for d in [temp_dir, output_temp_dir]:
+            if d:
+                try:
+                    shutil.rmtree(d)
+                except:
+                    pass
+        
+        return None, f"❌ 跟踪失败: {str(e)}", None, None
+
+
 
 # ===== 示例图像 =====
 example_images = ["003_img.png", "1977_Well_F-5_Field_1.png"]
@@ -692,29 +1086,138 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         label="📦 上传视频帧 ZIP 文件",
                         file_types=[".zip"]
                     )
+                    
+                    # First frame annotation for bounding box
+                    track_first_frame_annotator = BBoxAnnotator(
+                        label="🖼️ 首帧边界框标注 (可选)",
+                        categories=["cell"]
+                    )
+                    
+                    with gr.Row():
+                        track_use_box_radio = gr.Radio(
+                            choices=["Yes", "No"],
+                            value="No",
+                            label="🔲 使用边界框?"
+                        )
+                    
                     track_btn = gr.Button("▶️ 运行跟踪", variant="primary", size="lg")
                     
                     gr.Markdown(
-                        """
+                        '''
                         **使用说明:**
-                        1. 上传包含 `.tif` 图像的 ZIP 压缩包
-                        2. 点击 "运行跟踪"
-                        3. 结果保存到 `tracked_results/` 目录
-                        """
+                        1. 上传包含 `.tif` 图像序列的 ZIP 文件
+                        2. (可选) 在首帧上标注边界框并选择 "Yes"
+                        3. 点击 "运行跟踪"
+                        4. 下载CTC格式结果
+                        
+                        **注意:**
+                        - 确保TIFF文件按时间顺序命名 (如: t000.tif, t001.tif...)
+                        - 边界框将应用于整个视频序列
+                        - 避免使用macOS的压缩工具
+                        '''
                     )
                 
                 with gr.Column(scale=2):
+                    # Preview of first frame
+                    track_first_frame_preview = gr.Image(
+                        label="📸 首帧预览",
+                        type="pil",
+                        height=300
+                    )
+                    
                     track_output = gr.Textbox(
                         label="📊 跟踪信息",
-                        lines=12,
+                        lines=8,
                         interactive=False
                     )
+                    
+                    # Download button for CTC results
+                    track_download = gr.File(
+                        label="📥 下载跟踪结果 (CTC格式)",
+                        visible=False
+                    )
+            
+            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
+                
+                import tifffile
+                
+                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:
+                            # ===== NORMALIZATION STARTS HERE ===== 
+                            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
+                                # ===== NORMALIZATION ENDS HERE =====
+                            except Exception as e:
+                                print(f"⚠️  Error normalizing first frame: {e}")
+                                # Fallback to original file
+                                return first_frame
+                except Exception as e:
+                    print(f"⚠️  Error loading first frame: {e}")
+                    pass
+                return None
             
-            dummy = gr.Textbox(visible=False)
+            # Load first frame into annotator when ZIP uploaded
+            track_zip_upload.change(
+                fn=load_first_frame_for_annotation,
+                inputs=track_zip_upload,
+                outputs=track_first_frame_annotator
+            )
+            
+            # Run tracking
             track_btn.click(
                 fn=track_video_handler,
-                inputs=track_zip_upload,
-                outputs=[dummy, track_output]
+                inputs=[track_use_box_radio, track_first_frame_annotator, track_zip_upload],
+                outputs=[track_download, track_output, track_download, track_first_frame_preview]
             )
     
     gr.Markdown(
@@ -722,9 +1225,7 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
         ---
         ### 💡 技术说明
         
-        **分割 (Segmentation)** - 基于 Stable Diffusion 特征的实例分割  
-        **计数 (Counting)** - 密度图估计  
-        **跟踪 (Tracking)** - Trackastra 跟踪算法
+        **MicroscopyMatching** - 基于 Stable Diffusion 的显微图像分析工具套件
         """
     )