code/VESSA_data_preprocess/template_manager.py

"""
TemplateManager - Template Matching with FAISS

Efficient template matching with on-demand template loading:
1. Pass template_dir when calling match()
2. Templates are cached to avoid reloading
3. FAISS search for similarity ranking
4. Return first match with class overlap

Template .npz file format:
    Required fields:
    - 'cls_feature': DINOv2 CLS token feature (384-dim for vits14)
    - 'classes': List of class numbers in this template, e.g., [1, 2, 3]
    - 'original_size': (width, height) of original image
    - 'modality': Modality string (e.g., 'MRI' for ACDC, 'CT', etc.)
    - 'template_image': PIL RGB image as numpy array (H, W, 3)
    - 'template_mask': PIL Gray Multi-label segmentation mask with class numbers (H, W)

"""

import os
import json
import numpy as np
from PIL import Image
from typing import Union, Tuple, List, Dict, Optional
import torch
import torchvision.transforms as T
import faiss
import matplotlib.pyplot as plt


class TemplateManager:
    """
    Template matching manager with on-demand template loading
    
    Pipeline:
    1. Load templates from template_dir (cached)
    2. FAISS search to rank all templates by similarity
    3. Check class overlap from most similar
    4. Return first match with class overlap
    
    Features:
    - On-demand template loading (pass template_dir to match())
    - Template caching (avoid reloading same folder)
    - FAISS search for efficient similarity ranking
    - Early stopping when match found
    - Optional modality filtering
    """
    
    def __init__(
        self,
        device: str = "cuda",
        model_name: str = "dinov2_vits14",
        image_size: int = 224,
        distance_metric: str = "cosine"
    ):
        """
        Initialize TemplateManager without templates
        
        Args:
            device: "cuda" or "cpu"
            model_name: DINOv2 model name
            image_size: Image preprocessing size
            distance_metric: 'cosine' or 'l2'
        """
        self.device = torch.device(device if torch.cuda.is_available() else "cpu")
        self.distance_metric = distance_metric
        
        # Load DINOv2 model
        print(f"Loading DINOv2 model on {self.device}...")
        self.model = torch.hub.load('facebookresearch/dinov2', model_name)
        self.model = self.model.to(self.device).eval()
        
        # Image preprocessing
        self.transform = T.Compose([
            T.Resize(image_size),
            T.CenterCrop(image_size),
            T.ToTensor(),
            T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        ])
        
        # Template cache: {template_dir: (templates, faiss_index, template_ids_list)}
        self.template_cache = {}
        
        print(f"✓ TemplateManager initialized (ready to load templates)")
    
    def _load_templates(self, template_dir: str):
        """Load templates from directory and build FAISS index"""
        templates = {}
        
        if not os.path.exists(template_dir):
            raise ValueError(f"Template directory not found: {template_dir}")
        
        for filename in os.listdir(template_dir):
            if filename.endswith('.npz'):
                template_id = os.path.splitext(filename)[0]
                npz_path = os.path.join(template_dir, filename)
                
                try:
                    data = np.load(npz_path, allow_pickle=True)
                    
                    # Load feature and normalize
                    cls_feature = data['cls_feature']
                    if self.distance_metric == 'cosine':
                        cls_feature = cls_feature / np.linalg.norm(cls_feature)
                    
                    # Load classes
                    classes = data.get('classes', [])
                    if isinstance(classes, np.ndarray):
                        classes = classes.tolist()
                    
                    # Load template image
                    template_image = data.get('template_image', None)
                    if template_image is not None:
                        template_image = Image.fromarray(template_image.astype('uint8'))
                    
                    # Load template mask
                    template_mask = data.get('template_mask', None)
                    
                    # Load sizes
                    original_size = data.get('original_size', None)
                    if original_size is not None and isinstance(original_size, np.ndarray):
                        original_size = tuple(original_size)
                    
                    processed_size = data.get('processed_size', None)
                    if processed_size is not None and isinstance(processed_size, np.ndarray):
                        processed_size = tuple(processed_size)
                    
                    # Load modality from NPZ
                    modality = data.get('modality', None)
                    if modality is not None and isinstance(modality, np.ndarray):
                        modality = str(modality)  # Convert numpy string to Python string
                    
                    templates[template_id] = {
                        'cls_feature': cls_feature,
                        'template_image': template_image,
                        'template_mask': template_mask,
                        'classes': classes,
                        'original_size': original_size,
                        'processed_size': processed_size,
                        'modality': modality
                    }
                except Exception as e:
                    print(f"Warning: Failed to load {filename}: {e}")
        
        if not templates:
            raise ValueError(f"No valid templates found in {template_dir}")
        
        # Build FAISS index
        template_ids_list = list(templates.keys())
        features = np.stack([templates[tid]['cls_feature'] for tid in template_ids_list]).astype('float32')
        
        dimension = features.shape[1]
        if self.distance_metric == 'cosine':
            index = faiss.IndexFlatIP(dimension)
        else:
            index = faiss.IndexFlatL2(dimension)
        
        index.add(features)
        
        print(f"✓ Loaded {len(templates)} templates from {template_dir}")
        
        return templates, index, template_ids_list
    
    def _extract_feature(self, image: Union[str, Image.Image]):
        """Extract normalized DINOv2 feature from image"""
        if isinstance(image, str):
            img = Image.open(image).convert('RGB')
        elif isinstance(image, Image.Image):
            img = image.convert('RGB')
        else:
            raise ValueError("Image must be a file path or PIL Image")
        
        tensor = self.transform(img).unsqueeze(0).to(self.device)
        
        with torch.no_grad():
            feature = self.model(tensor).squeeze(0).cpu().numpy()
        
        if self.distance_metric == 'cosine':
            feature = feature / np.linalg.norm(feature)
        
        return feature
    
    def match(
        self,
        target_image: Union[str, Image.Image],
        template_dir: str,
        target_labels: Optional[List[int]] = None,
        target_modality: Optional[str] = None,
        return_mask: bool = True,
        match_strategy: str = "partial"
    ) -> Union[Tuple[Optional[Image.Image], Optional[float]], Tuple[Optional[Image.Image], Optional[np.ndarray], Optional[float]]]:
        """
        Find the most matching template using FAISS similarity search
        
        Args:
            target_image: Target image path or PIL Image
            template_dir: Directory containing template .npz files
            target_labels: List of class numbers, e.g., [1, 2, 3]
            target_modality: Image modality (e.g., 'T1', 'CT', 'MRI')
            return_mask: Whether to return the template mask
            match_strategy: Class matching strategy:
                - "partial": Match if template has ANY of the target classes (at least 1 overlap)
                - "complete": Match only if template has ALL target classes
                - "exact": Match only if template has exactly the same classes
            
        Returns:
            If return_mask=True:
                (matched_template_image, matched_template_mask, similarity_score)
                Returns (None, None, None) if no matching template found
            If return_mask=False:
                (matched_template_image, similarity_score)
                Returns (None, None) if no matching template found
        """
        # Load templates from cache or disk
        if template_dir not in self.template_cache:
            print(f"Loading templates from {template_dir}...")
            templates, index, template_ids_list = self._load_templates(template_dir)
            self.template_cache[template_dir] = (templates, index, template_ids_list)
        else:
            templates, index, template_ids_list = self.template_cache[template_dir]
        
        # Extract target feature
        target_feature = self._extract_feature(target_image)
        print(f"Target feature shape: {target_feature.shape}")
        query = target_feature.reshape(1, -1).astype('float32')
        
        # FAISS search all templates
        num_templates = len(template_ids_list)
        D, I = index.search(query, num_templates)
        
        # Check class overlap from most similar
        if target_labels:
            target_class_values = set(target_labels)  # Direct set from numbers
            
            for similarity, template_idx in zip(D[0], I[0]):
                template_id = template_ids_list[template_idx]
                template_data = templates[template_id]
                
                # Check modality
                template_modality = template_data['modality']
                if target_modality and template_modality and target_modality != template_modality:
                    continue
                
                # Check class overlap based on strategy
                template_classes = set(template_data['classes'])
                
                match_found = False
                if match_strategy == "partial":
                    # At least one class overlap
                    match_found = len(target_class_values.intersection(template_classes)) > 0
                elif match_strategy == "complete":
                    # Template must contain all target classes (can have more)
                    match_found = target_class_values.issubset(template_classes)
                elif match_strategy == "exact":
                    # Template must have exactly the same classes
                    match_found = target_class_values == template_classes
                else:
                    raise ValueError(f"Unknown match_strategy: {match_strategy}. Use 'partial', 'complete', or 'exact'.")
                
                if match_found:
                    # Found match
                    matched_image = template_data['template_image']
                    matched_mask = template_data['template_mask']
                    
                    if matched_image is None:
                        continue
                    
                    print(f"✓ Match found: template classes {sorted(template_classes)} vs target {sorted(target_labels)} (strategy: {match_strategy})")
                    
                    if return_mask:
                        return matched_image, matched_mask, float(similarity)
                    else:
                        return matched_image, float(similarity)
            
            # No matching template found with class overlap
            if return_mask:
                return None, None, None
            else:
                return None, None
        
        # No class constraint - return most similar
        template_id = template_ids_list[I[0][0]]
        template_data = templates[template_id]
        matched_image = template_data['template_image']
        matched_mask = template_data['template_mask']
        similarity = float(D[0][0])
        
        if matched_image is None:
            if return_mask:
                return None, None, None
            else:
                return None, None
        
        if return_mask:
            return matched_image, matched_mask, similarity
        else:
            return matched_image, similarity
    
    def clear_cache(self, template_dir: Optional[str] = None):
        """
        Clear template cache
        
        Args:
            template_dir: Specific directory to clear, or None to clear all
        """
        if template_dir is None:
            self.template_cache.clear()
            print("✓ Cleared all template cache")
        elif template_dir in self.template_cache:
            del self.template_cache[template_dir]
            print(f"✓ Cleared cache for {template_dir}")
    
    def get_cache_info(self):
        """Get information about cached templates"""
        info = {
            'num_cached_dirs': len(self.template_cache),
            'cached_dirs': {}
        }
        
        for template_dir, (templates, _, _) in self.template_cache.items():
            info['cached_dirs'][template_dir] = {
                'num_templates': len(templates),
                'template_ids': list(templates.keys())[:5]  # First 5 IDs
            }
        
        return info
    
    def visualize_match(
        self,
        target_image: Union[str, Image.Image],
        matched_image: Optional[Image.Image],
        matched_mask: Optional[np.ndarray],
        similarity: Optional[float],
        save_path: Optional[str] = None,
        show_mask: bool = True
    ):
        """
        Visualize input image and matched template image side by side
        
        Args:
            target_image: Input image (file path or PIL Image)
            matched_image: Matched template image
            matched_mask: Matched template mask
            similarity: Similarity score
            save_path: Path to save visualization (optional)
            show_mask: Whether to overlay mask on matched image
        """
        # Load input image
        if isinstance(target_image, str):
            target_img = Image.open(target_image).convert('RGB')
        elif isinstance(target_image, Image.Image):
            target_img = target_image.convert('RGB')
        else:
            raise ValueError("target_image must be a file path or PIL Image")
        
        # If no match found
        if matched_image is None:
            fig, ax = plt.subplots(1, 1, figsize=(6, 6))
            ax.imshow(target_img)
            ax.set_title("Input Image\n(No matching template found)", fontsize=12, fontweight='bold')
            ax.axis('off')
        else:
            # Create subplots
            if show_mask and matched_mask is not None:
                fig, axes = plt.subplots(1, 3, figsize=(18, 6))
            else:
                fig, axes = plt.subplots(1, 2, figsize=(12, 6))
            
            # Show input image
            axes[0].imshow(target_img)
            axes[0].set_title("Input Image", fontsize=12, fontweight='bold')
            axes[0].axis('off')
            
            # Show matched template image
            axes[1].imshow(matched_image)
            title = f"Matched Template"
            if similarity is not None:
                title += f"\nSimilarity: {similarity:.4f}"
            axes[1].set_title(title, fontsize=12, fontweight='bold', color='green')
            axes[1].axis('off')
            
            # If mask available, show overlay
            if show_mask and matched_mask is not None:
                # Create overlay image
                matched_array = np.array(matched_image)
                
                # Create colored mask (different colors for different classes)
                mask_colored = np.zeros_like(matched_array)
                unique_classes = np.unique(matched_mask)
                colors = plt.cm.rainbow(np.linspace(0, 1, len(unique_classes)))
                
                for idx, class_val in enumerate(unique_classes):
                    if class_val == 0:  # Skip background
                        continue
                    mask_colored[matched_mask == class_val] = (colors[idx][:3] * 255).astype(np.uint8)
                
                # Blend mask with image
                alpha = 0.4
                overlay = matched_array.copy()
                mask_indices = matched_mask > 0
                overlay[mask_indices] = (
                    matched_array[mask_indices] * (1 - alpha) + 
                    mask_colored[mask_indices] * alpha
                ).astype(np.uint8)
                
                axes[2].imshow(overlay)
                axes[2].set_title("Template + Mask Overlay", fontsize=12, fontweight='bold')
                axes[2].axis('off')
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            print(f"✓ Visualization saved to: {save_path}")
        
        plt.show()
        
        return fig


############################################################################################
if __name__ == "__main__":
    # Configuration
    TEMPLATE_DIR_1 = "/projects/p32685/RL_dataset/acdc_processed/template"
    TARGET_IMAGE = "/projects/p32685/RL_dataset/acdc/train/images/patient001_slice0_1.jpg"
    
    # Initialize (no templates loaded yet)
    manager = TemplateManager(device="cuda:0")
    
    # Example 1: Match with template folder
    print("\n--- Example 1: Match with Template Folder ---")
    assert os.path.exists(TARGET_IMAGE), "Target image not found"

    target_labels = [1, 2, 3]  # Just class numbers
    
    # Templates loaded on first call
    # Try different matching strategies:
    # - "partial": template has at least one of [1,2,3] (default, most flexible)
    # - "complete": template must have all [1,2,3] (can have more like [1,2,3,4])
    # - "exact": template must have exactly [1,2,3]
    matched_image, matched_mask, similarity = manager.match(
        target_image=TARGET_IMAGE,
        template_dir=TEMPLATE_DIR_1,  # Pass folder here!
        target_labels=target_labels,
        target_modality="MRI",
        return_mask=True,
        match_strategy="complete"  # Change to "partial", "complete", or "exact"
    )
    
    # Check if match found
    if matched_image is not None:
        print(f"✓ Best match found!")
        print(f"  Similarity: {similarity:.4f}")
        print(f"  Template image size: {matched_image.size}")
        print(f"  Mask shape: {matched_mask.shape if matched_mask is not None else None}")
        
        # Visualize matching results
        print("\nVisualizing matching results...")
        manager.visualize_match(
            target_image=TARGET_IMAGE,
            matched_image=matched_image,
            matched_mask=matched_mask,
            similarity=similarity,
            save_path="match_result.png",  # Save visualization
            show_mask=True  # Show mask overlay
        )
    else:
        print(f"✗ No matching template found with specified criteria")
        # Show input image even if no match found
        manager.visualize_match(
            target_image=TARGET_IMAGE,
            matched_image=None,
            matched_mask=None,
            similarity=None,
            save_path="no_match_result.png"
        )