VisionLanguageGroup/MicroscopyMatching

1# inference_track.py2 3import torch4import numpy as np5import os6from pathlib import Path7from tqdm import tqdm8from huggingface_hub import hf_hub_download9from tracking_one import TrackingModule10from models.tra_post_model.tracking import graph_to_ctc11 12MODEL = None13DEVICE = torch.device("cpu")14 15def load_model(use_box=False):16    """17    load tracking model from Hugging Face Hub18    19    Args:20        use_box: use bounding box as input (default: False)21    22    Returns:23        model: loaded tracking model24        device25    """26    global MODEL, DEVICE27    28    try:29        print("🔄 Loading tracking model...")30        31        # 初始化模型32        MODEL = TrackingModule(use_box=use_box)33        34        # Load checkpoint from Hugging Face Hub35        ckpt_path = hf_hub_download(36            repo_id="phoebe777777/111",37            filename="microscopy_matching_tra.pth",38            token=None,39            force_download=False40        )41        42        print(f"✅ Checkpoint downloaded: {ckpt_path}")43        44        # Load weights45        MODEL.load_state_dict(46            torch.load(ckpt_path, map_location="cpu"), 47            strict=True48        )49        MODEL.eval()50        51        # Move model to device52        if torch.cuda.is_available():53            DEVICE = torch.device("cuda")54            MODEL.move_to_device(DEVICE)55            print("✅ Model moved to CUDA")56        else:57            DEVICE = torch.device("cpu")58            MODEL.move_to_device(DEVICE)59            print("✅ Model on CPU")60        61        print("✅ Tracking model loaded successfully")62        return MODEL, DEVICE63        64    except Exception as e:65        print(f"❌ Error loading tracking model: {e}")66        import traceback67        traceback.print_exc()68        return None, torch.device("cpu")69 70 71@torch.no_grad()72def run(model, video_dir, box=None, device="cpu", output_dir="tracked_results"):73    """74    run tracking inference on video frames75    76    Args:77        model: loaded tracking model78        video_dir: directory of video frame sequence (contains consecutive image files)79        box: bounding box (optional)80        device: device81        output_dir: output directory82    83    Returns:84        result_dict: {85            'track_graph': TrackGraph object containing tracking results,86            'masks': tracked masks (T, H, W),87            'output_dir': output directory path,88            'num_tracks': number of tracked trajectories89        }90    """91    if model is None:92        return {93            'track_graph': None,94            'masks': None,95            'output_dir': None,96            'num_tracks': 0,97            'error': 'Model not loaded'98        }99    100    try:101        print(f"🔄 Running tracking inference on {video_dir}")102        103        # Run tracking104        track_graph, masks = model.track(105            file_dir=video_dir,106            boxes=box,107            mode="greedy",  # Optional: "greedy", "greedy_nodiv", "ilp"108            dataname="tracking_result"109        )110        111        # 创建输出目录112        if not os.path.exists(output_dir):113            os.makedirs(output_dir)114        115        # Convert tracking results to CTC format and save116        print("🔄 Converting to CTC format...")117        ctc_tracks, masks_tracked = graph_to_ctc(118            track_graph,119            masks,120            outdir=output_dir,121        )122        print(f"✅ CTC results saved to {output_dir}")123        124        125        print(f"✅ Tracking completed")126        127        result = {128            'track_graph': track_graph,129            'masks': masks,130            'masks_tracked': masks_tracked,131            'output_dir': output_dir,132        }133        134        return result135        136    except Exception as e:137        print(f"❌ Tracking inference error: {e}")138        import traceback139        traceback.print_exc()140        return {141            'track_graph': None,142            'masks': None,143            'output_dir': None,144            'num_tracks': 0,145            'error': str(e)146        }147 148 149def visualize_tracking_result(masks_tracked, output_path):150    """151    visualize tracking results152    153    Args:154        masks_tracked: masks with tracking results (T, H, W)155        output_path: output video file path156    157    Returns:158        output_path: output video file path159    """160    try:161        import cv2162        import matplotlib.pyplot as plt163        from matplotlib import cm164        165        T, H, W = masks_tracked.shape166        167        # create a color map for unique track IDs168        unique_ids = np.unique(masks_tracked)169        num_colors = len(unique_ids)170        cmap = cm.get_cmap('tab20', num_colors)171        172        # create video writer173        fourcc = cv2.VideoWriter_fourcc(*'mp4v')174        out = cv2.VideoWriter(output_path, fourcc, 5.0, (W, H))175        176        for t in range(T):177            frame = masks_tracked[t]178            179            # create colored image180            colored_frame = np.zeros((H, W, 3), dtype=np.uint8)181            for i, obj_id in enumerate(unique_ids):182                if obj_id == 0:183                    continue184                mask = (frame == obj_id)185                color = np.array(cmap(i % num_colors)[:3]) * 255186                colored_frame[mask] = color187            188            # convert to BGR (OpenCV format)189            colored_frame_bgr = cv2.cvtColor(colored_frame, cv2.COLOR_RGB2BGR)190            out.write(colored_frame_bgr)191        192        out.release()193        print(f"✅ Visualization saved to {output_path}")194        return output_path195        196    except Exception as e:197        print(f"❌ Visualization error: {e}")198        return None199