update counting and segmentation, tracking-basic

app.py

@@ -21,7 +21,7 @@ from inference_track import load_model as load_track_model, run as run_track
 
 # ===== 清理缓存目录 =====
 print("===== 清理缓存 =====")
-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)
@@ -138,8 +138,8 @@ def colorize_mask(mask: np.ndarray, num_colors: int = 512) -> np.ndarray:
 
 # ===== 分割功能 =====
 def segment_with_choice(use_box_choice, annot_value):
-    print(use_box_choice)
-    print(annot_value)
+    print("边界框选择:", use_box_choice)
+    print("注释值:", annot_value)
     """分割主函数 - 每个实例不同颜色+轮廓"""
     if annot_value is None or len(annot_value) < 1:
         print("❌ No annotation input")
@@ -256,34 +256,86 @@ def count_cells_handler(use_box_choice, annot_value):
             return None, f"❌ 计数失败: {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}")
         
         # 只提取密度图部分(假设visualized_path是拼接图,我们只要右半部分)
-        viz_path = result.get('visualized_path')
+        # viz_path = result.get('visualized_path')
         
         # 如果有density_map_path,直接使用
-        if 'density_map_path' in result:
-            density_path = result['density_map_path']
-        elif viz_path and os.path.exists(viz_path):
-            # 如果是拼接图,提取右半部分(密度图)
-            try:
-                viz_img = Image.open(viz_path)
-                w, h = viz_img.size
-                # 取右半部分
-                density_img = viz_img.crop((w//2, 0, w, h))
-                # 保存为新文件
-                temp_density = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-                density_img.save(temp_density.name)
-                density_path = temp_density.name
-            except:
-                density_path = viz_path
-        else:
-            density_path = viz_path
+        # if 'density_map_path' in result:
+        #     density_path = result['density_map_path']
+        # elif viz_path and os.path.exists(viz_path):
+        #     # 如果是拼接图,提取右半部分(密度图)
+        #     try:
+        #         viz_img = Image.open(viz_path)
+        #         w, h = viz_img.size
+        #         # 取右半部分
+        #         density_img = viz_img
+        #         # 保存为新文件
+        #         temp_density = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
+        #         density_img.save(temp_density.name)
+        #         density_path = temp_density.name
+        #     except:
+        #         density_path = viz_path
+        # else:
+        #     density_path = viz_path
+
+        # 读取原图
+        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"✅ 检测到 {count:.1f} 个细胞"
         
         print(f"✅ Counting done - Count: {count:.1f}")
+
+        return Image.fromarray(overlay), temp_density_file.name, result_text
         
-        return density_path, result_text
+        # return density_path, result_text
         
     except Exception as e:
         print(f"❌ Counting error: {e}")
@@ -295,6 +347,8 @@ def count_cells_handler(use_box_choice, annot_value):
 def find_tif_dir(root_dir):
     """递归查找第一个包含 .tif 文件的目录"""
     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
@@ -328,20 +382,18 @@ def track_video_handler(zip_file_obj):
         if 'error' in result:
             return None, f"❌ 跟踪失败: {result['error']}"
         
-        num_tracks = result['num_tracks']
         output_dir = result['output_dir']
         
         result_text = f"""✅ 跟踪完成!
 
-🎯 跟踪轨迹数量: {num_tracks}
-📁 结果保存在: {output_dir}
+        📁 结果保存在: {output_dir}
 
-包含的文件:
-- res_track.txt (CTC格式轨迹)
-- 其他跟踪数据文件
-"""
+        包含的文件:
+        - res_track.txt (CTC格式轨迹)
+        - 其他跟踪数据文件
+        """
 
-        print(f"✅ Tracking done - {num_tracks} tracks")
+        print(f"✅ Tracking done")
         return None, result_text
 
     except zipfile.BadZipFile:
@@ -575,12 +627,18 @@ with gr.Blocks(title="Microscopy Analysis Suite", theme=gr.themes.Soft()) as dem
                         label="📊 统计信息",
                         lines=2
                     )
+
+                    # 下载原始预测结果
+                    download_density_btn = gr.File(
+                        label="📥 下载原始预测 (.npy 格式)",
+                        visible=True
+                    )
             
             # 绑定事件
             count_btn.click(
                 fn=count_cells_handler,
                 inputs=[count_use_box_radio, count_annotator],
-                outputs=[count_output, count_status]
+                outputs=[count_output, download_density_btn, count_status]
             )
             
             # 初始化Gallery显示

counting.py

@@ -1,7 +1,5 @@
 # stable diffusion x loca
 import os
-# os.system("source /etc/network_turbo")
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
 import pprint
 from typing import Any, List, Optional
 import argparse

inference_count.py

@@ -49,7 +49,14 @@ def load_model(use_box=False):
         )
         MODEL.eval()
         
-        DEVICE = torch.device("cpu")
+        if torch.cuda.is_available():
+            DEVICE = torch.device("cuda")
+            MODEL.move_to_device(DEVICE)
+            print("✅ Model moved to CUDA")
+        else:
+            DEVICE = torch.device("cpu")
+            MODEL.move_to_device(DEVICE)
+            print("✅ Model on CPU")
         
         print("✅ Counting model loaded successfully")
         return MODEL, DEVICE
@@ -80,6 +87,15 @@ def run(model, img_path, box=None, device="cpu", visualize=True):
             'visualized_path': str (如果 visualize=True)
         }
     """
+    print("DEVICE:", device)
+    model.move_to_device(device)
+    model.eval()
+    if box is not None:
+        use_box = True
+    else:
+        use_box = False
+    model.use_box = use_box
+
     if model is None:
         return {
             'density_map': None,
@@ -92,7 +108,8 @@ def run(model, img_path, box=None, device="cpu", visualize=True):
         print(f"🔄 Running counting inference on {img_path}")
         
         # 运行推理 (调用你的模型的 forward 方法)
-        density_map, count = model(img_path, box)
+        with torch.no_grad():
+            density_map, count = model(img_path, box)
         
         print(f"✅ Counting result: {count:.1f} objects")
         
@@ -103,9 +120,9 @@ def run(model, img_path, box=None, device="cpu", visualize=True):
         }
         
         # 可视化
-        if visualize:
-            viz_path = visualize_result(img_path, density_map, count)
-            result['visualized_path'] = viz_path
+        # if visualize:
+        #     viz_path = visualize_result(img_path, density_map, count)
+        #     result['visualized_path'] = viz_path
         
         return result
         
@@ -153,18 +170,13 @@ def visualize_result(image_path, density_map, count):
         img_show = (img_show - np.min(img_show)) / (np.max(img_show) - np.min(img_show) + 1e-8)
         
         # 创建可视化 (与你原来的代码一致)
-        fig, ax = plt.subplots(1, 2, figsize=(12, 6))
-        
-        # 左图: 原始图像
-        ax[0].imshow(img_show)
-        ax[0].axis('off')
-        ax[0].set_title(f"Input image")
+        fig, ax = plt.subplots(figsize=(8, 6))
         
         # 右图: 密度图叠加
-        ax[1].imshow(img_show)
-        ax[1].imshow(density_map_show, cmap='jet', alpha=0.5)
-        ax[1].axis('off')
-        ax[1].set_title(f"Predicted density map, count: {count:.1f}")
+        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()
         

inference_seg.py

@@ -18,15 +18,30 @@ def load_model(use_box=False):
     )
     MODEL.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
     MODEL.eval()
-    DEVICE = torch.device("cpu")
+    if torch.cuda.is_available():
+        DEVICE = torch.device("cuda")
+        MODEL.move_to_device(DEVICE)
+        print("✅ Model moved to CUDA")
+    else:
+        DEVICE = torch.device("cpu")
+        MODEL.move_to_device(DEVICE)
+        print("✅ Model on CPU")
     return MODEL, DEVICE
 
 
 @torch.no_grad()
 def run(model, img_path, box=None, device="cpu"):
-    output = model(img_path, box=box)
+    print("DEVICE:", device)
+    model.move_to_device(device)
+    model.eval()
+    with torch.no_grad():
+        if box is not None:
+            use_box = True
+        else:
+            use_box = False
+        model.use_box = use_box
+        output = model(img_path, box=box)
     mask = output
-    mask = (mask > 0).astype(np.uint8)
     return mask
 # import os
 # import torch

inference_track.py

@@ -120,17 +120,18 @@ def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):
             masks,
             outdir=output_dir,
         )
+        print(f"✅ CTC results saved to {output_dir}")
         
-        num_tracks = len(track_graph.tracks())
+        # num_tracks = len(track_graph.tracks())
         
-        print(f"✅ Tracking completed: {num_tracks} tracks found")
+        print(f"✅ Tracking completed")
         
         result = {
             'track_graph': track_graph,
             'masks': masks,
             'masks_tracked': masks_tracked,
             'output_dir': output_dir,
-            'num_tracks': num_tracks
+            # 'num_tracks': num_tracks
         }
         
         return result

segmentation.py

@@ -1,5 +1,4 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "2"
 import pprint
 from typing import Any, List, Optional
 import argparse

tracking_one.py

@@ -1,5 +1,4 @@
 import os
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 import pprint
 from typing import Any, List, Optional
 import argparse