update

app.py

@@ -18,7 +18,6 @@ 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
@@ -27,7 +26,6 @@ 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")
@@ -39,7 +37,6 @@ if os.path.exists(cache_path):
     except:
         pass
 
-# ===== 全局模型变量 =====
 SEG_MODEL = None
 SEG_DEVICE = torch.device("cpu")
 
@@ -50,7 +47,6 @@ 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
@@ -76,14 +72,12 @@ def load_all_models():
 
 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"""
+    """Save feedback to 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)
@@ -102,13 +96,11 @@ def save_feedback_to_hf(query_id, feedback_type, feedback_text=None, img_path=No
     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}",
@@ -117,19 +109,17 @@ def save_feedback_to_hf(query_id, feedback_type, feedback_text=None, img_path=No
             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文件"""
+    """Save feedback to local JSON file"""
     feedback_data = {
         "query_id": query_id,
         "feedback_type": feedback_type,
@@ -154,9 +144,8 @@ def save_feedback(query_id, feedback_type, feedback_text=None, img_path=None, bb
     with feedback_file.open("w") as f:
         json.dump(feedback_data, f, indent=4, ensure_ascii=False)
 
-# ===== 辅助函数 =====
 def parse_first_bbox(bboxes):
-    """解析第一个边界框"""
+    """Parse the first bounding box from the annotation input, supports dict or list format"""
     if not bboxes:
         return None
     b = bboxes[0]
@@ -169,7 +158,7 @@ def parse_first_bbox(bboxes):
     return None
 
 def parse_bboxes(bboxes):
-    """解析所有边界框"""
+    """Parse all bounding boxes from the annotation input"""
     if not bboxes:
         return None
     
@@ -185,7 +174,7 @@ def parse_bboxes(bboxes):
     return result
 
 def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
-    """将实例掩码转换为彩色图像"""
+    """Convert a 2D mask of instance IDs to a color image for visualization."""
     def hsv_to_rgb(h, s, v):
         i = int(h * 6.0)
         f = h * 6.0 - i
@@ -210,10 +199,9 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
     color_idx = mask % num_colors
     return palette_arr[color_idx]
 
-# ===== 分割功能 =====
 # @spaces.GPU
 def segment_with_choice(use_box_choice, annot_value):
-    """分割主函数 - 每个实例不同颜色+轮廓"""
+    """Segmentation handler - supports bounding box, returns colorized overlay and original mask path"""
     if annot_value is None or len(annot_value) < 1:
         print("❌ No annotation input")
         return None, None
@@ -224,18 +212,12 @@ def segment_with_choice(use_box_choice, annot_value):
     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))
@@ -243,13 +225,11 @@ def segment_with_choice(use_box_choice, annot_value):
         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)
@@ -276,24 +256,19 @@ def segment_with_choice(use_box_choice, annot_value):
         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]  # 黄色轮廓
@@ -302,7 +277,7 @@ def segment_with_choice(use_box_choice, annot_value):
 
     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"""
@@ -315,11 +290,6 @@ def count_cells_handler(use_box_choice, annot_value):
     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
@@ -341,7 +311,6 @@ def count_cells_handler(use_box_choice, annot_value):
         
         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}")
@@ -365,20 +334,16 @@ def count_cells_handler(use_box_choice, annot_value):
             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
         
@@ -400,7 +365,6 @@ def count_cells_handler(use_box_choice, annot_value):
 
         return Image.fromarray(overlay), temp_density_file.name, result_text
         
-        # return density_path, result_text
         
     except Exception as e:
         print(f"❌ Counting error: {e}")
@@ -408,7 +372,7 @@ def count_cells_handler(use_box_choice, annot_value):
         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):
@@ -502,14 +466,13 @@ def create_ctc_results_zip(output_dir):
     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)
@@ -568,14 +531,7 @@ def create_tracking_visualization(tif_dir, output_dir, valid_tif_files):
         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
@@ -710,19 +666,19 @@ def create_tracking_visualization(tif_dir, output_dir, valid_tif_files):
 # @spaces.GPU
 def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj):
     """
-    支持 ZIP 压缩包上传的 Tracking 处理函数 - 支持首帧边界框
+    Tracking handler - processes a ZIP of TIF frames, supports bounding box, returns visualization and results ZIP
     
     Parameters:
     -----------
     use_box_choice : str
-        "Yes" or "No" - 是否使用边界框
+        "Yes" or "No" - whether to use bounding box annotation for tracking
     first_frame_annot : tuple or None
         (image_path, bboxes) from BBoxAnnotator, only used if user annotated first frame
     zip_file_obj : File
         Uploaded ZIP file containing TIF sequence
     """
     if zip_file_obj is None:
-        return None, "⚠️ 请上传包含视频帧的压缩包 (.zip)", None, None
+        return None, "⚠️  Please upload a ZIP file containing video frames (.zip)", None, None
     
     temp_dir = None
     output_temp_dir = None
@@ -734,11 +690,6 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj):
             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
@@ -880,9 +831,8 @@ def track_video_handler(use_box_choice, first_frame_annot, zip_file_obj):
 
 
 
-# ===== 示例图像 =====
+# ===== Example Images =====
 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")]
 
@@ -909,7 +859,6 @@ with gr.Blocks(
         object-fit: contain !important;
     }
 
-    /* 强制密度图容器和图片高度 */
     #density_map_output {
         height: 500px !important;
     }
@@ -939,7 +888,7 @@ with gr.Blocks(
     
     # 全局状态
     current_query_id = gr.State(str(uuid.uuid4()))
-    user_uploaded_examples = gr.State(example_images_seg.copy())  # 初始化时包含原始示例
+    user_uploaded_examples = gr.State(example_images_seg.copy())
     
     with gr.Tabs():
         # ===== Tab 1: Segmentation =====
@@ -1031,27 +980,27 @@ with gr.Blocks(
                         visible=False
                     )
             
-            # 绑定事件: 运行分割
+            # click event for segmentation
             run_seg_btn.click(
                 fn=segment_with_choice,
                 inputs=[use_box_radio, annotator],
                 outputs=[seg_output, download_mask_btn]
             )
 
-            # 清空按钮事件
+            # click event for clear button
             clear_btn.click(
                 fn=lambda: None,
                 inputs=None,
                 outputs=annotator
             )
             
-            # 初始化Gallery显示
+            # init Gallery with example images
             demo.load(
                 fn=lambda: example_images_seg.copy(),
                 outputs=example_gallery
             )
             
-            # 绑定事件: 上传示例图片
+            # click event for image uploader
             def add_to_gallery(img_path, current_imgs):
                 if not img_path:
                     return current_imgs
@@ -1072,7 +1021,7 @@ with gr.Blocks(
                 outputs=example_gallery
             )
             
-            # 绑定事件: 点击Gallery加载
+            # click event for Gallery selection
             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]
@@ -1084,7 +1033,7 @@ with gr.Blocks(
                 outputs=annotator
             )
             
-            # 绑定事���: 提交反馈
+            # click event for submitting feedback
             def submit_user_feedback(query_id, score, comment, annot_val):
                 try:
                     img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
@@ -1097,7 +1046,7 @@ with gr.Blocks(
                     #     img_path=img_path,
                     #     bboxes=bboxes
                     # )
-                    # 使用 HF 存储
+
                     save_feedback_to_hf(
                         query_id=query_id,
                         feedback_type=f"score_{int(score)}",
@@ -1259,14 +1208,14 @@ with gr.Blocks(
                 outputs=[count_output, download_density_btn, count_status]
             )
 
-            # 清空按钮事件
+            # Clear selection
             clear_btn.click(
                 fn=lambda: None,
                 inputs=None,
                 outputs=count_annotator
             )
 
-            # 绑定事件: 提交反馈
+            # Submit feedback
             def submit_user_feedback(query_id, score, comment, annot_val):
                 try:
                     img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
@@ -1279,7 +1228,7 @@ with gr.Blocks(
                     #     img_path=img_path,
                     #     bboxes=bboxes
                     # )
-                    # 使用 HF 存储
+
                     save_feedback_to_hf(
                         query_id=query_id,
                         feedback_type=f"score_{int(score)}",
@@ -1581,14 +1530,14 @@ with gr.Blocks(
                 outputs=[track_download, track_output, track_download, track_first_frame_preview]
             )
 
-            # 清空按钮事件
+            # Clear selection
             clear_btn.click(
                 fn=lambda: None,
                 inputs=None,
                 outputs=track_first_frame_annotator
             )
 
-            # 绑定事件: 提交反馈
+            # Submit feedback
             def submit_user_feedback(query_id, score, comment, annot_val):
                 try:
                     img_path = annot_val[0] if annot_val and len(annot_val) > 0 else None
@@ -1601,7 +1550,7 @@ with gr.Blocks(
                     #     img_path=img_path,
                     #     bboxes=bboxes
                     # )
-                    # 使用 HF 存储
+
                     save_feedback_to_hf(
                         query_id=query_id,
                         feedback_type=f"score_{int(score)}",

inference_count.py

@@ -1,5 +1,4 @@
 # inference_count.py
-# 计数模型推理模块 - 独立版本
 
 import torch
 import numpy as np
@@ -15,24 +14,22 @@ DEVICE = torch.device("cpu")
 
 def load_model(use_box=False):
     """
-    加载计数模型
+    load counting model from Hugging Face Hub
     
     Args:
-        use_box: 是否使用边界框
+        use_box: use bounding box as input (default: False)
     
     Returns:
-        model: 加载的模型
-        device: 设备
+        model: loaded counting model
+        device: device
     """
     global MODEL, DEVICE
     
     try:
         print("🔄 Loading counting model...")
-        
-        # 初始化模型
+
         MODEL = CountingModule(use_box=use_box)
         
-        # 从 Hugging Face Hub 下载权重
         ckpt_path = hf_hub_download(
             repo_id="phoebe777777/111",
             filename="microscopy_matching_cnt.pth",
@@ -42,7 +39,6 @@ def load_model(use_box=False):
         
         print(f"✅ Checkpoint downloaded: {ckpt_path}")
         
-        # 加载权重
         MODEL.load_state_dict(
             torch.load(ckpt_path, map_location="cpu"), 
             strict=True
@@ -71,20 +67,20 @@ def load_model(use_box=False):
 @torch.no_grad()
 def run(model, img_path, box=None, device="cpu", visualize=True):
     """
-    运行计数推理
+    Run counting inference on a single image
     
     Args:
-        model: 计数模型
-        img_path: 图像路径
-        box: 边界框 [[x1, y1, x2, y2], ...] 或 None
-        device: 设备
-        visualize: 是否生成可视化
+        model: loaded counting model
+        img_path: image path
+        box: bounding box [[x1, y1, x2, y2], ...] or None
+        device: device
+        visualize: whether to generate visualization
     
     Returns:
         result_dict: {
             'density_map': numpy array,
             'count': float,
-            'visualized_path': str (如果 visualize=True)
+            'visualized_path': str (if visualize=True)
         }
     """
     print("DEVICE:", device)
@@ -107,7 +103,6 @@ def run(model, img_path, box=None, device="cpu", visualize=True):
     try:
         print(f"🔄 Running counting inference on {img_path}")
         
-        # 运行推理 (调用你的模型的 forward 方法)
         with torch.no_grad():
             density_map, count = model(img_path, box)
         
@@ -118,11 +113,7 @@ def run(model, img_path, box=None, device="cpu", visualize=True):
             'count': count,
             'visualized_path': None
         }
-        
-        # 可视化
-        # if visualize:
-        #     viz_path = visualize_result(img_path, density_map, count)
-        #     result['visualized_path'] = viz_path
+
         
         return result
         
@@ -140,47 +131,39 @@ def run(model, img_path, box=None, device="cpu", visualize=True):
 
 def visualize_result(image_path, density_map, count):
     """
-    可视化计数结果 (与你原来的可视化代码一致)
+    Visualize counting results (consistent with your original visualization code)
     
     Args:
-        image_path: 原始图像路径
-        density_map: 密度图 (numpy array)
-        count: 计数值
+        image_path: original image path
+        density_map: numpy array of predicted density map
+        count
     
     Returns:
-        output_path: 可视化结果的临时文件路径
+        output_path: temporary file path of the visualization result
     """
     try:
         import skimage.io as io
         
-        # 读取原始图像
         img = io.imread(image_path)
         
-        # 处理不同格式的图像
         if len(img.shape) == 3 and img.shape[2] > 3:
             img = img[:, :, :3]
         if len(img.shape) == 2:
             img = np.stack([img]*3, axis=-1)
         
-        # 归一化显示
         img_show = img.squeeze()
         density_map_show = density_map.squeeze()
         
-        # 归一化图像
         img_show = (img_show - np.min(img_show)) / (np.max(img_show) - np.min(img_show) + 1e-8)
         
-        # 创建可视化 (与你原来的代码一致)
         fig, ax = plt.subplots(figsize=(8, 6))
         
-        # 右图: 密度图叠加
         ax.imshow(img_show)
         ax.imshow(density_map_show, cmap='jet', alpha=0.5)
         ax.axis('off')
-        # ax.set_title(f"Predicted density map, count: {count:.1f}")
         
         plt.tight_layout()
         
-        # 保存到临时文件
         temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
         plt.savefig(temp_file.name, dpi=300)
         plt.close()
@@ -195,13 +178,11 @@ def visualize_result(image_path, density_map, count):
         return image_path
 
 
-# ===== 测试代码 =====
 if __name__ == "__main__":
     print("="*60)
     print("Testing Counting Model")
     print("="*60)
-    
-    # 测试模型加载
+
     model, device = load_model(use_box=False)
     
     if model is not None:
@@ -209,7 +190,6 @@ if __name__ == "__main__":
         print("Model loaded successfully, testing inference...")
         print("="*60)
         
-        # 测试推理
         test_image = "example_imgs/1977_Well_F-5_Field_1.png"
         
         if os.path.exists(test_image):

inference_seg.py

@@ -43,45 +43,3 @@ def run(model, img_path, box=None, device="cpu"):
         output = model(img_path, box=box)
     mask = output
     return mask
-# import os
-# import torch
-# import numpy as np
-# from huggingface_hub import hf_hub_download
-# from segmentation import SegmentationModule
-
-# MODEL = None
-# DEVICE = torch.device("cpu")
-
-# def load_model(use_box=False):
-#     global MODEL, DEVICE
-
-#     # === 优化1: 使用 /data 缓存模型，避免写入 .cache ===
-#     cache_dir = "/data/cellseg_model_cache"
-#     os.makedirs(cache_dir, exist_ok=True)
-
-#     ckpt_path = hf_hub_download(
-#         repo_id="Shengxiao0709/cellsegmodel",
-#         filename="microscopy_matching_seg.pth",
-#         token=None,
-#         local_dir=cache_dir,              # ✅ 下载到 /data
-#         local_dir_use_symlinks=False,     # ✅ 避免软链接问题
-#         force_download=False              # ✅ 已存在时不重复下载
-#     )
-
-#     # === 优化2: 加载模型 ===
-#     MODEL = SegmentationModule(use_box=use_box)
-#     state_dict = torch.load(ckpt_path, map_location="cpu")
-#     MODEL.load_state_dict(state_dict, strict=False)
-#     MODEL.eval()
-
-#     DEVICE = torch.device("cpu")
-#     print(f"✅ Model loaded from {ckpt_path}")
-#     return MODEL, DEVICE
-
-
-# @torch.no_grad()
-# def run(model, img_path, box=None, device="cpu"):
-#     output = model(img_path, box=box)
-#     mask = output["pred"]
-#     mask = (mask > 0).astype(np.uint8)
-#     return mask

inference_track.py

@@ -1,5 +1,4 @@
 # inference_track.py
-# 视频跟踪模型推理模块
 
 import torch
 import numpy as np
@@ -15,14 +14,14 @@ DEVICE = torch.device("cpu")
 
 def load_model(use_box=False):
     """
-    加载跟踪模型
+    load tracking model from Hugging Face Hub
     
     Args:
-        use_box: 是否使用边界框
+        use_box: use bounding box as input (default: False)
     
     Returns:
-        model: 加载的模型
-        device: 设备
+        model: loaded tracking model
+        device
     """
     global MODEL, DEVICE
     
@@ -32,7 +31,7 @@ def load_model(use_box=False):
         # 初始化模型
         MODEL = TrackingModule(use_box=use_box)
         
-        # 从 Hugging Face Hub 下载权重
+        # Load checkpoint from Hugging Face Hub
         ckpt_path = hf_hub_download(
             repo_id="phoebe777777/111",
             filename="microscopy_matching_tra.pth",
@@ -42,14 +41,14 @@ def load_model(use_box=False):
         
         print(f"✅ Checkpoint downloaded: {ckpt_path}")
         
-        # 加载权重
+        # Load weights
         MODEL.load_state_dict(
             torch.load(ckpt_path, map_location="cpu"), 
             strict=True
         )
         MODEL.eval()
         
-        # 设置设备
+        # Move model to device
         if torch.cuda.is_available():
             DEVICE = torch.device("cuda")
             MODEL.move_to_device(DEVICE)
@@ -72,21 +71,21 @@ def load_model(use_box=False):
 @torch.no_grad()
 def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
     """
-    运行视频跟踪推理
+    run tracking inference on video frames
     
     Args:
-        model: 跟踪模型
-        video_dir: 视频帧序列目录 (包含连续的图像文件)
-        box: 边界框 (可选)
-        device: 设备
-        output_dir: 输出目录
+        model: loaded tracking model
+        video_dir: directory of video frame sequence (contains consecutive image files)
+        box: bounding box (optional)
+        device: device
+        output_dir: output directory
     
     Returns:
         result_dict: {
-            'track_graph': TrackGraph对象,
-            'masks': 分割掩码数组 (T, H, W),
-            'output_dir': 输出目录路径,
-            'num_tracks': 跟踪轨迹数量
+            'track_graph': TrackGraph object containing tracking results,
+            'masks': tracked masks (T, H, W),
+            'output_dir': output directory path,
+            'num_tracks': number of tracked trajectories
         }
     """
     if model is None:
@@ -101,11 +100,11 @@ def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
     try:
         print(f"🔄 Running tracking inference on {video_dir}")
         
-        # 运行跟踪
+        # Run tracking
         track_graph, masks = model.track(
             file_dir=video_dir,
             boxes=box,
-            mode="greedy",  # 可选: "greedy", "greedy_nodiv", "ilp"
+            mode="greedy",  # Optional: "greedy", "greedy_nodiv", "ilp"
             dataname="tracking_result"
         )
         
@@ -113,7 +112,7 @@ def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
         if not os.path.exists(output_dir):
             os.makedirs(output_dir)
         
-        # 转换为CTC格式并保存
+        # Convert tracking results to CTC format and save
         print("🔄 Converting to CTC format...")
         ctc_tracks, masks_tracked = graph_to_ctc(
             track_graph,
@@ -122,7 +121,6 @@ def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
         )
         print(f"✅ CTC results saved to {output_dir}")
         
-        # num_tracks = len(track_graph.tracks())
         
         print(f"✅ Tracking completed")
         
@@ -131,7 +129,6 @@ def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
             'masks': masks,
             'masks_tracked': masks_tracked,
             'output_dir': output_dir,
-            # 'num_tracks': num_tracks
         }
         
         return result
@@ -151,36 +148,35 @@ def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
 
 def visualize_tracking_result(masks_tracked, output_path):
     """
-    可视化跟踪结果 (可选)
+    visualize tracking results
     
     Args:
-        masks_tracked: 跟踪后的掩码 (T, H, W)
-        output_path: 输出视频路径
+        masks_tracked: masks with tracking results (T, H, W)
+        output_path: output video file path
     
     Returns:
-        output_path: 视频文件路径
+        output_path: output video file path
     """
     try:
         import cv2
         import matplotlib.pyplot as plt
         from matplotlib import cm
         
-        # 获取时间帧数
         T, H, W = masks_tracked.shape
         
-        # 创建颜色映射
+        # create a color map for unique track IDs
         unique_ids = np.unique(masks_tracked)
         num_colors = len(unique_ids)
         cmap = cm.get_cmap('tab20', num_colors)
         
-        # 创建视频写入器
+        # create video writer
         fourcc = cv2.VideoWriter_fourcc(*'mp4v')
         out = cv2.VideoWriter(output_path, fourcc, 5.0, (W, H))
         
         for t in range(T):
             frame = masks_tracked[t]
             
-            # 创建彩色图像
+            # create colored image
             colored_frame = np.zeros((H, W, 3), dtype=np.uint8)
             for i, obj_id in enumerate(unique_ids):
                 if obj_id == 0:
@@ -189,7 +185,7 @@ def visualize_tracking_result(masks_tracked, output_path):
                 color = np.array(cmap(i % num_colors)[:3]) * 255
                 colored_frame[mask] = color
             
-            # 转换为BGR (OpenCV格式)
+            # convert to BGR (OpenCV format)
             colored_frame_bgr = cv2.cvtColor(colored_frame, cv2.COLOR_RGB2BGR)
             out.write(colored_frame_bgr)