Create inference_cout.py

inference_cout.py

@@ -0,0 +1,225 @@
+# inference_count.py
+# 计数模型推理模块 - 独立版本
+
+import torch
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import tempfile
+import os
+from huggingface_hub import hf_hub_download
+from counting_model import CountingModule
+
+MODEL = None
+DEVICE = torch.device("cpu")
+
+def load_model(use_box=False):
+    """
+    加载计数模型
+    
+    Args:
+        use_box: 是否使用边界框
+    
+    Returns:
+        model: 加载的模型
+        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="Shengxiao0709/cellsegmodel",
+            filename="microscopy_matching_cnt.pth",
+            token=None,
+            force_download=False
+        )
+        
+        print(f"✅ Checkpoint downloaded: {ckpt_path}")
+        
+        # 加载权重
+        MODEL.load_state_dict(
+            torch.load(ckpt_path, map_location="cpu"), 
+            strict=True
+        )
+        MODEL.eval()
+        
+        DEVICE = torch.device("cpu")
+        
+        print("✅ Counting model loaded successfully")
+        return MODEL, DEVICE
+        
+    except Exception as e:
+        print(f"❌ Error loading counting model: {e}")
+        import traceback
+        traceback.print_exc()
+        return None, torch.device("cpu")
+
+
+@torch.no_grad()
+def run(model, img_path, box=None, device="cpu", visualize=True):
+    """
+    运行计数推理
+    
+    Args:
+        model: 计数模型
+        img_path: 图像路径
+        box: 边界框 [[x1, y1, x2, y2], ...] 或 None
+        device: 设备
+        visualize: 是否生成可视化
+    
+    Returns:
+        result_dict: {
+            'density_map': numpy array,
+            'count': float,
+            'visualized_path': str (如果 visualize=True)
+        }
+    """
+    if model is None:
+        return {
+            'density_map': None,
+            'count': 0,
+            'visualized_path': None,
+            'error': 'Model not loaded'
+        }
+    
+    try:
+        print(f"🔄 Running counting inference on {img_path}")
+        
+        # 运行推理 (调用你的模型的 forward 方法)
+        density_map, count = model(img_path, box)
+        
+        print(f"✅ Counting result: {count:.1f} objects")
+        
+        result = {
+            'density_map': density_map,
+            'count': count,
+            'visualized_path': None
+        }
+        
+        # 可视化
+        if visualize:
+            viz_path = visualize_result(img_path, density_map, count)
+            result['visualized_path'] = viz_path
+        
+        return result
+        
+    except Exception as e:
+        print(f"❌ Counting inference error: {e}")
+        import traceback
+        traceback.print_exc()
+        return {
+            'density_map': None,
+            'count': 0,
+            'visualized_path': None,
+            'error': str(e)
+        }
+
+
+def visualize_result(image_path, density_map, count):
+    """
+    可视化计数结果 (与你原来的可视化代码一致)
+    
+    Args:
+        image_path: 原始图像路径
+        density_map: 密度图 (numpy array)
+        count: 计数值
+    
+    Returns:
+        output_path: 可视化结果的临时文件路径
+    """
+    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(1, 2, figsize=(12, 6))
+        
+        # 左图: 原始图像
+        ax[0].imshow(img_show)
+        ax[0].axis('off')
+        ax[0].set_title(f"Input image")
+        
+        # 右图: 密度图叠加
+        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}")
+        
+        plt.tight_layout()
+        
+        # 保存到临时文件
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
+        plt.savefig(temp_file.name, dpi=300)
+        plt.close()
+        
+        print(f"✅ Visualization saved to {temp_file.name}")
+        return temp_file.name
+        
+    except Exception as e:
+        print(f"❌ Visualization error: {e}")
+        import traceback
+        traceback.print_exc()
+        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:
+        print("\n" + "="*60)
+        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):
+            result = run(
+                model,
+                test_image,
+                box=None,
+                device=device,
+                visualize=True
+            )
+            
+            if 'error' not in result:
+                print("\n" + "="*60)
+                print("Inference Results:")
+                print("="*60)
+                print(f"Count: {result['count']:.1f}")
+                print(f"Density map shape: {result['density_map'].shape}")
+                if result['visualized_path']:
+                    print(f"Visualization saved to: {result['visualized_path']}")
+            else:
+                print(f"\n❌ Inference failed: {result['error']}")
+        else:
+            print(f"\n⚠️ Test image not found: {test_image}")
+    else:
+        print("\n❌ Model loading failed")