VisionLanguageGroup/MicroscopyMatching

1# inference_count.py2 3import torch4import numpy as np5from PIL import Image6import matplotlib.pyplot as plt7import tempfile8import os9from huggingface_hub import hf_hub_download10from counting import CountingModule11 12MODEL = None13DEVICE = torch.device("cpu")14 15def load_model(use_box=False):16    """17    load counting model from Hugging Face Hub18    19    Args:20        use_box: use bounding box as input (default: False)21    22    Returns:23        model: loaded counting model24        device: device25    """26    global MODEL, DEVICE27    28    try:29        print("🔄 Loading counting model...")30 31        MODEL = CountingModule(use_box=use_box)32        33        ckpt_path = hf_hub_download(34            repo_id="phoebe777777/111",35            filename="microscopy_matching_cnt.pth",36            token=None,37            force_download=False38        )39        40        print(f"✅ Checkpoint downloaded: {ckpt_path}")41        42        MODEL.load_state_dict(43            torch.load(ckpt_path, map_location="cpu"), 44            strict=True45        )46        MODEL.eval()47        48        if torch.cuda.is_available():49            DEVICE = torch.device("cuda")50            MODEL.move_to_device(DEVICE)51            print("✅ Model moved to CUDA")52        else:53            DEVICE = torch.device("cpu")54            MODEL.move_to_device(DEVICE)55            print("✅ Model on CPU")56        57        print("✅ Counting model loaded successfully")58        return MODEL, DEVICE59        60    except Exception as e:61        print(f"❌ Error loading counting model: {e}")62        import traceback63        traceback.print_exc()64        return None, torch.device("cpu")65 66 67@torch.no_grad()68def run(model, img_path, box=None, device="cpu", visualize=True):69    """70    Run counting inference on a single image71    72    Args:73        model: loaded counting model74        img_path: image path75        box: bounding box [[x1, y1, x2, y2], ...] or None76        device: device77        visualize: whether to generate visualization78    79    Returns:80        result_dict: {81            'density_map': numpy array,82            'count': float,83            'visualized_path': str (if visualize=True)84        }85    """86    print("DEVICE:", device)87    model.move_to_device(device)88    model.eval()89    if box is not None:90        use_box = True91    else:92        use_box = False93    model.use_box = use_box94 95    if model is None:96        return {97            'density_map': None,98            'count': 0,99            'visualized_path': None,100            'error': 'Model not loaded'101        }102    103    try:104        print(f"🔄 Running counting inference on {img_path}")105        106        with torch.no_grad():107            density_map, count = model(img_path, box)108        109        print(f"✅ Counting result: {count:.1f} objects")110        111        result = {112            'density_map': density_map,113            'count': count,114            'visualized_path': None115        }116 117        118        return result119        120    except Exception as e:121        print(f"❌ Counting inference error: {e}")122        import traceback123        traceback.print_exc()124        return {125            'density_map': None,126            'count': 0,127            'visualized_path': None,128            'error': str(e)129        }130 131 132def visualize_result(image_path, density_map, count):133    """134    Visualize counting results (consistent with your original visualization code)135    136    Args:137        image_path: original image path138        density_map: numpy array of predicted density map139        count140    141    Returns:142        output_path: temporary file path of the visualization result143    """144    try:145        import skimage.io as io146        147        img = io.imread(image_path)148        149        if len(img.shape) == 3 and img.shape[2] > 3:150            img = img[:, :, :3]151        if len(img.shape) == 2:152            img = np.stack([img]*3, axis=-1)153        154        img_show = img.squeeze()155        density_map_show = density_map.squeeze()156        157        img_show = (img_show - np.min(img_show)) / (np.max(img_show) - np.min(img_show) + 1e-8)158        159        fig, ax = plt.subplots(figsize=(8, 6))160        161        ax.imshow(img_show)162        ax.imshow(density_map_show, cmap='jet', alpha=0.5)163        ax.axis('off')164        165        plt.tight_layout()166        167        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')168        plt.savefig(temp_file.name, dpi=300)169        plt.close()170        171        print(f"✅ Visualization saved to {temp_file.name}")172        return temp_file.name173        174    except Exception as e:175        print(f"❌ Visualization error: {e}")176        import traceback177        traceback.print_exc()178        return image_path179 180 181if __name__ == "__main__":182    print("="*60)183    print("Testing Counting Model")184    print("="*60)185 186    model, device = load_model(use_box=False)187    188    if model is not None:189        print("\n" + "="*60)190        print("Model loaded successfully, testing inference...")191        print("="*60)192        193        test_image = "example_imgs/1977_Well_F-5_Field_1.png"194        195        if os.path.exists(test_image):196            result = run(197                model,198                test_image,199                box=None,200                device=device,201                visualize=True202            )203            204            if 'error' not in result:205                print("\n" + "="*60)206                print("Inference Results:")207                print("="*60)208                print(f"Count: {result['count']:.1f}")209                print(f"Density map shape: {result['density_map'].shape}")210                if result['visualized_path']:211                    print(f"Visualization saved to: {result['visualized_path']}")212            else:213                print(f"\n❌ Inference failed: {result['error']}")214        else:215            print(f"\n⚠️ Test image not found: {test_image}")216    else:217        print("\n❌ Model loading failed")218