code/VESSA_data_preprocess/packing_HF_dataset.py

import os
import json
import argparse
import sys
import random
import cv2
import numpy as np
from datasets import Dataset, DatasetDict, Features, Value, Image, Sequence
from PIL import Image as PILImage
from tqdm import tqdm
from pathlib import Path
import torch
import torchvision.transforms as T
import glob


# ========== Template Feature Extraction Functions ==========

def load_dinov2_model(device="cuda"):
    """Load DINOv2 model for feature extraction"""
    print(f"Loading DINOv2 model on {device}...")
    model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vits14')
    model = model.to(device).eval()
    return model

# Image preprocessing transform for DINOv2
transform = T.Compose([
    T.Resize(224),
    T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

def extract_template_features(template_data_dict, output_dir, device="cuda"):
    """
    Extract DINOv2 features from template images and save with masks according to TemplateManager format.
    
    Template .npz file format (required by TemplateManager):
        - '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)
    
    Args:
        template_data_dict: Dictionary mapping template_id -> {
            'image': PIL Image (RGB),
            'mask': PIL Image (grayscale) or numpy array with class values,
            'modality': str (e.g., 'MRI', 'CT')
        }
        output_dir: Directory to save the .npz feature files
        device: Device to run model on
    
    Returns:
        Dictionary mapping template_id -> feature array info
    """
    os.makedirs(output_dir, exist_ok=True)
    
    # Load model
    model = load_dinov2_model(device)
    
    print(f"\n=== Extracting Template Features (TemplateManager Format) ===")
    print(f"Number of templates: {len(template_data_dict)}")
    print(f"Output directory: {output_dir}")
    
    results = {}
    
    for i, (template_id, data) in enumerate(template_data_dict.items()):
        try:
            img = data['image']
            mask = data['mask']
            modality = data.get('modality', 'MRI')
            
            # Ensure image is RGB
            if img.mode != 'RGB':
                img = img.convert('RGB')
            
            # Get original size
            orig_width, orig_height = img.size
            
            # Process mask to extract classes
            if isinstance(mask, PILImage.Image):
                mask_array = np.array(mask)
            else:
                mask_array = mask
            
            # Extract unique class values (excluding background 0)
            unique_classes = np.unique(mask_array)
            classes = [int(c) for c in unique_classes if c > 0]
            
            # Apply transform for feature extraction
            tensor = transform(img).unsqueeze(0).to(device)
            
            # Extract DINOv2 features
            with torch.no_grad():
                cls_feature = model(tensor).squeeze(0).cpu().numpy()
            
            # Save features with TemplateManager format
            output_path = os.path.join(output_dir, f"{template_id}.npz")
            
            # Convert image to numpy array
            template_image_array = np.array(img)  # Shape: (H, W, 3)
            
            save_dict = {
                'cls_feature': cls_feature,  # Shape: (feature_dim,) e.g., (384,)
                'classes': np.array(classes),  # List of class numbers
                'original_size': np.array([orig_width, orig_height]),
                'modality': np.array(modality),  # Modality string
                'template_image': template_image_array,  # RGB image array (H, W, 3)
                'template_mask': mask_array,  # Grayscale mask with class numbers (H, W)
            }
            
            np.savez(output_path, **save_dict)
            
            results[template_id] = {
                'path': output_path,
                'original_size': (orig_width, orig_height),
                'cls_feature_shape': cls_feature.shape,
                'classes': classes,
                'modality': modality
            }
            
            if (i + 1) % 10 == 0 or (i + 1) == len(template_data_dict):
                print(f"[{i+1}/{len(template_data_dict)}] Processed: {template_id} (classes: {classes})")
            
        except Exception as e:
            print(f"Error processing template {template_id}: {e}")
            import traceback
            traceback.print_exc()
    
    # Save a summary JSON file
    summary_path = os.path.join(output_dir, "template_summary.json")
    with open(summary_path, 'w') as f:
        json_results = {}
        for template_id, info in results.items():
            json_results[template_id] = {
                'path': info['path'],
                'original_size': list(info['original_size']),
                'cls_feature_shape': list(info['cls_feature_shape']),
                'classes': info['classes'],
                'modality': info['modality']
            }
        json.dump(json_results, f, indent=2)
    
    print(f"\n✓ Extracted features for {len(results)} templates")
    print(f"✓ Summary saved to: {summary_path}")
    
    return results


def extract_templates_from_dataset(image_dir, mask_dir, num_templates=20, modality="MRI", 
                                   patient_filter=None, device="cuda"):
    """
    Extract template images and masks from a dataset directory.
    
    Args:
        image_dir: Directory containing images
        mask_dir: Directory containing masks (should have same filenames as images)
        num_templates: Number of template images to process (-1 for all)
        modality: Modality string (default: 'MRI' for ACDC dataset)
        patient_filter: List of patient IDs to process (e.g., ['patient001', 'patient002'])
        device: Device to run model on
    
    Returns:
        Dictionary of template data ready for feature extraction
    """
    print(f"\n=== Loading Template Images ===")
    print(f"Image directory: {image_dir}")
    print(f"Mask directory: {mask_dir}")
    print(f"Modality: {modality}")
    if patient_filter:
        print(f"Patient filter: {patient_filter}")
    
    # Find all images in the directory
    if patient_filter:
        # Filter for specific patients
        image_paths = []
        for patient_id in patient_filter:
            patient_images = sorted(glob.glob(os.path.join(image_dir, f"{patient_id}*.jpg")))
            patient_images.extend(sorted(glob.glob(os.path.join(image_dir, f"{patient_id}*.png"))))
            image_paths.extend(patient_images)
        image_paths = sorted(list(set(image_paths)))  # Remove duplicates
    else:
        image_paths = sorted(glob.glob(os.path.join(image_dir, "*.jpg")))
        image_paths.extend(sorted(glob.glob(os.path.join(image_dir, "*.png"))))
    
    if not image_paths:
        print(f"No images found in {image_dir}")
        return None
    
    print(f"Found {len(image_paths)} images")
    
    # Limit to num_templates if specified
    if num_templates > 0 and num_templates < len(image_paths):
        image_paths = image_paths[:num_templates]
        print(f"Processing first {num_templates} images as templates")
    
    # Load images and masks into dictionary
    template_data_dict = {}
    for img_path in tqdm(image_paths, desc="Loading templates"):
        # Extract template ID from filename (without extension)
        template_id = os.path.splitext(os.path.basename(img_path))[0]
        
        # Construct mask path (assuming same filename but in mask_dir)
        mask_path = None
        for ext in ['.png', '.jpg', '.npy']:
            potential_mask = os.path.join(mask_dir, f"{template_id}{ext}")
            if os.path.exists(potential_mask):
                mask_path = potential_mask
                break
        
        if mask_path is None:
            continue
        
        try:
            # Load image as PIL RGB
            img = PILImage.open(img_path).convert('RGB')
            
            # Load mask
            if mask_path.endswith('.npy'):
                mask_array = np.load(mask_path)
                mask = PILImage.fromarray(mask_array.astype(np.uint8))
            else:
                mask = PILImage.open(mask_path)
                # Convert to grayscale if needed
                if mask.mode != 'L':
                    mask = mask.convert('L')
            
            template_data_dict[template_id] = {
                'image': img,
                'mask': mask,
                'modality': modality
            }
            
        except Exception as e:
            print(f"Error loading {img_path} or {mask_path}: {e}")
    
    if not template_data_dict:
        print("No templates loaded successfully")
        return None
    
    print(f"✓ Successfully loaded {len(template_data_dict)} template image-mask pairs")
    
    return template_data_dict


class GenHFDataset:
    """
    A class for generating and processing datasets for HuggingFace format.
    """
    
    def __init__(self, device="cuda:0"):
        """
        Initialize the GenHFDataset class.
        
        Args:
            device (str): Device to run inference on
        """
        self.device = device
        
    def _load_dataset_info(self, dataset_root_dir):
        """
        Load dataset information from dataset_info.json file.
        
        Args:
            dataset_root_dir (str): Root directory containing dataset_info.json
            
        Returns:
            dict: Dataset information
        """
        dataset_json_path = os.path.join(dataset_root_dir, "dataset_info.json")
        if not os.path.exists(dataset_json_path):
            raise FileNotFoundError(f"dataset_info.json not found in {dataset_root_dir}")
        
        with open(dataset_json_path, 'r', encoding='utf-8') as f:
            dataset_info = json.load(f)
        
        return dataset_info
    
    def print_dataset_info(self, dataset_root_dir):
        """
        Print detailed information about the dataset from dataset_info.json.
        
        Args:
            dataset_root_dir (str): Root directory containing dataset_info.json
        """
        dataset_info = self._load_dataset_info(dataset_root_dir)
        
        print("=" * 60)
        print("DATASET INFORMATION")
        print("=" * 60)
        print(f"Dataset: {dataset_info.get('dataset_name', 'Unknown')}")
        
        # Optional fields with defaults
        if 'description' in dataset_info:
            print(f"Description: {dataset_info['description']}")
        if 'modality' in dataset_info:
            print(f"Modality: {dataset_info['modality']}")
        if 'total_references' in dataset_info:
            print(f"Total references: {dataset_info['total_references']}")
        if 'processed_references' in dataset_info:
            print(f"Processed references: {dataset_info['processed_references']}")
        
        splits = dataset_info.get('splits', {})
        if splits:
            print(f"\nSplits:")
            for split_name, count in splits.items():
                print(f"  {split_name}: {count}")
        
        categories = dataset_info.get('categories', {})
        if categories:
            print(f"\nCategories (top 10):")
            sorted_categories = sorted(categories.items(), key=lambda x: int(x[1]), reverse=True)[:10]
            for cat_id, count in sorted_categories:
                print(f"  Category {cat_id}: {count} samples")
        
        mask_types = dataset_info.get('mask_types', {})
        if mask_types:
            print(f"\nMask types:")
            for mask_type, count in mask_types.items():
                print(f"  {mask_type}: {count}")
        
        print("=" * 60)
    
    def _get_labels_from_dataset_info(self, dataset_info):
        """
        Extract labels from dataset_info.json.
        
        Args:
            dataset_info (dict): Dataset information dictionary
            
        Returns:
            dict: Labels dictionary with class mapping
        """
        # Get class mapping from dataset_info
        class_mapping = dataset_info.get('class_mapping', {})
        classes = dataset_info.get('classes', [])
        
        # If class_mapping exists, use it
        if class_mapping:
            labels = class_mapping
        # Otherwise, build from classes list (assuming 0 is background)
        elif classes:
            labels = {"0": "background"}
            for idx, class_name in enumerate(classes, start=1):
                labels[str(idx)] = class_name
        else:
            raise ValueError("No class information found in dataset_info.json")
        
        # Convert labels dict to list of label names (excluding background)
        label_names = []
        for class_id, class_name in labels.items():
            if class_id != '0':  # Skip background class
                label_names.append(class_name)
        
        return {
            'labels': labels,
            'label_names': label_names
        }
    
    def create_local_dataset(self, train_data, val_data, output_dir):
        """
        Create and save a HuggingFace DatasetDict from processed train and validation data.
            
        Args:
            train_data (dict): Training data dictionary
            val_data (dict): Validation data dictionary
            output_dir (str): Output directory for the dataset
            
        Returns:
            DatasetDict: The created dataset
        """
        base_features = {
            "id": Value("string"),
            "dataset_name": Value("string"),
            "image": Value("string"),  # File path to image
            "label_mask": Value("string"),  # File path to mask
            "labels": Value("string"),  # JSON string containing per-class information
        }
        has_template = (
            (train_data and "template" in train_data)
            or (val_data and "template" in val_data)
        )
        if has_template:
            base_features["template"] = Value("string")
        features = Features(base_features)

        def process_split(split_data, split_name="data"):
            # Handle empty data
            if not split_data or len(split_data.get('image', [])) == 0:
                return {}
                
            processed_data = split_data.copy()
            
            print(f"Processing {split_name} split: {len(split_data['image'])} images")
            print(f"Note: Images and masks are stored as file paths, not loaded into memory")
            
            # Keep original paths without loading images
            processed_data['image'] = split_data['image']  # Keep as file paths
            processed_data['label_mask'] = split_data['label_mask']  # Keep as file paths
            
            return processed_data

        dataset_dict = {}
        
        # Only add train dataset if train_data is not empty
        if train_data and len(train_data.get('image', [])) > 0:
            dataset_dict['train'] = Dataset.from_dict(process_split(train_data, "train"), features=features)
        
        # Only add validation set if it exists and is not empty
        if val_data and len(val_data.get('image', [])) > 0:
            dataset_dict['val'] = Dataset.from_dict(process_split(val_data, "validation"), features=features)
        
        # Check if we have any data
        if not dataset_dict:
            raise ValueError("No valid data found in either train or validation sets")
        
        dataset = DatasetDict(dataset_dict)
        dataset.save_to_disk(output_dir)
        print(f"DATASET SAVED TO: {output_dir}")
        return dataset

    def generate_dataset(self, data_dir, output_base_dir, split_name, subset_size=0):
        """
        Generate a list of entries for a given dataset directory without using annotation files.
        
        Args:
            data_dir (str): Path to the dataset split (e.g., /root/data/acdc/train)
            output_base_dir (str): Base output directory for intermediate files
            split_name (str): Name of the split (e.g., 'train', 'val')
            subset_size (int): Number of samples to process. If 0, use all samples.
            
        Returns:
            List[dict]: List of entry dicts
        """
        # Get dataset name and labels from dataset_info.json
        dataset_root_dir = os.path.dirname(data_dir)  # Go up one level from train/val to root
        dataset_info = self._load_dataset_info(dataset_root_dir)
        dataset_name = dataset_info.get('dataset_name', 'Unknown')
        
        # Get labels from dataset_info
        label_data = self._get_labels_from_dataset_info(dataset_info)
        labels = label_data['labels']
        label_names = label_data['label_names']
        
        image_dir = os.path.join(data_dir, "images")
        mask_dir = os.path.join(data_dir, "masks")
        
        # Check if directories exist
        if not os.path.exists(image_dir):
            raise ValueError(f"Images directory not found: {image_dir}")
        if not os.path.exists(mask_dir):
            raise ValueError(f"Masks directory not found: {mask_dir}")
        
        # Get all image files (support both .jpg and .png)
        image_files = [f for f in os.listdir(image_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
        
        if not image_files:
            raise ValueError(f"No image files found in {image_dir}")
        
        # Sample subset_size if specified, otherwise use all files
        if subset_size > 0:
            num_files = min(subset_size, len(image_files))
            selected_files = random.sample(image_files, num_files)
        else:
            selected_files = image_files
        
        print(f"Processing {len(selected_files)} image files from {len(image_files)} total files")
        
        entries = []
        for image_file in tqdm(selected_files, desc=f"Processing {split_name}"):
            # Get base name (remove extension)
            base_name = os.path.splitext(image_file)[0]
            
            # Get full image path
            image_path = os.path.join(image_dir, image_file)
            
            
            if not os.path.exists(image_path):
                print(f"Warning: Image file {image_path} not found")
                continue
            
            # Check if corresponding mask exists (try .png extension for mask)
            mask_file = base_name + ".png"
            mask_path = os.path.join(mask_dir, mask_file)
            
            if not os.path.exists(mask_path):
                print(f"Warning: Mask file {mask_path} not found for image {image_file}")
                continue
            
            # Load image to verify it's valid
            image = cv2.imread(image_path)
            if image is None:
                print(f"Warning: Could not load image {image_path}")
                continue
            
            # Load mask to verify it's valid
            gt_mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            if gt_mask is None:
                print(f"Warning: Could not load mask {mask_path}")
                continue
            
            # Create entry for this sample
            entry = {
                "id": base_name,
                "dataset_name": dataset_name,
                'image': image_path,
                'label_mask': mask_path,  # multi-class mask, num-class read from dataset_info.json
                'labels': json.dumps(labels),  # Store labels from dataset_info.json
            }
            entries.append(entry)
        
        return entries

    def entries_to_dict(self, entries):
        """
        Convert a list of entry dicts to a dict of lists for HuggingFace Dataset construction.
            
        Args:
            entries (list): List of entry dicts
            
        Returns:
            dict: Dict of lists keyed by entry fields
        """
        if not entries:
            return {}
        result = {k: [] for k in entries[0].keys()}
        for entry in entries:
            for k, v in entry.items():
                result[k].append(v)
        return result
    
    def save_entries_as_backup(self, entries, output_dir, split_name):
        """
        Save entries as backup JSON files before creating the dataset.
        
        Args:
            entries (list): List of entry dictionaries
            output_dir (str): Output directory
            split_name (str): Split name (train/val)
        """
        backup_dir = os.path.join(output_dir, "backup_entries")
        os.makedirs(backup_dir, exist_ok=True)
        
        backup_file = os.path.join(backup_dir, f"{split_name}_entries.json")
        
        # Convert entries to serializable format
        serializable_entries = []
        for entry in entries:
            serializable_entry = entry.copy()
            # Ensure all paths are strings
            serializable_entry['image'] = str(entry['image'])
            serializable_entry['label_mask'] = str(entry['label_mask'])
            serializable_entries.append(serializable_entry)
        
        with open(backup_file, 'w', encoding='utf-8') as f:
            json.dump(serializable_entries, f, indent=2, ensure_ascii=False)
        
        print(f"Backup saved to: {backup_file}")
        return backup_file

    def load_entries_from_backup(self, backup_file):
        """
        Load entries from backup JSON file.
        
        Args:
            backup_file (str): Path to backup JSON file
            
        Returns:
            list: List of entry dictionaries
        """
        with open(backup_file, 'r', encoding='utf-8') as f:
            entries = json.load(f)
        
        print(f"Loaded {len(entries)} entries from backup: {backup_file}")
        return entries

    def recover_from_backup(self, output_dir):
        """
        Recover dataset from backup entries and create HuggingFace dataset.
        
        Args:
            output_dir (str): Output directory containing backup_entries
            
        Returns:
            DatasetDict: The recovered dataset
        """
        backup_dir = os.path.join(output_dir, "backup_entries")
        if not os.path.exists(backup_dir):
            raise FileNotFoundError(f"Backup directory not found: {backup_dir}")
        
        train_entries = []
        val_entries = []
        
        # Load train entries
        train_backup = os.path.join(backup_dir, "train_entries.json")
        if os.path.exists(train_backup):
            train_entries = self.load_entries_from_backup(train_backup)
        
        # Load val entries
        val_backup = os.path.join(backup_dir, "val_entries.json")
        if os.path.exists(val_backup):
            val_entries = self.load_entries_from_backup(val_backup)
        
        if not train_entries and not val_entries:
            raise ValueError("No backup entries found")
        
        # Convert entries to dict format
        train_data = self.entries_to_dict(train_entries)
        val_data = self.entries_to_dict(val_entries) if val_entries else {}
        
        # Create the dataset
        dataset = self.create_local_dataset(
            train_data=train_data,
            val_data=val_data,
            output_dir=os.path.join(output_dir, "hf_dataset")
        )
        
        print("Dataset recovered successfully from backup!")
        return dataset

    def process_full_dataset(self, dataset_root_dir, output_dir, 
                           subset_size=0):
        """
        Process both train and validation datasets and create a complete HuggingFace dataset.
        
        Args:
            dataset_root_dir (str): Root directory containing train/ and val/ folders
            output_dir (str): Output directory for the final dataset
            subset_size (int): Number of samples to process per split (0 means all samples)
            
        Returns:
            DatasetDict: The complete dataset
        """
        # Define train and val directories
        train_dir = os.path.join(dataset_root_dir, "train")
        val_dir = os.path.join(dataset_root_dir, "val")
        
        # Check if at least one of train or val directories exists
        if not os.path.exists(train_dir) and not os.path.exists(val_dir):
            raise FileNotFoundError(f"Neither train nor val directory found in {dataset_root_dir}")
        
        # Process training set if it exists
        train_entries = []
        if os.path.exists(train_dir):
            print("Generating train entries...")
            train_entries = self.generate_dataset(
                data_dir=train_dir,
                output_base_dir=output_dir,
                split_name="train",
                subset_size=subset_size,
            )
            print(f"Train entries: {len(train_entries)}")
        else:
            print("Train directory not found, skipping training set")
        
        # Always save train entries as backup, even if empty
        self.save_entries_as_backup(train_entries, output_dir, "train")
###########################################################################################
        # Process validation set if it exists
        val_entries = []
        if os.path.exists(val_dir):
            print("Generating val entries...")
            try:
                val_entries = self.generate_dataset(
                    data_dir=val_dir,
                    output_base_dir=output_dir,
                    split_name="val",
                    subset_size=subset_size,
                )
                print(f"Val entries: {len(val_entries)}")
            except Exception as e:
                print(f"Warning: Failed to process validation set: {e}")
                val_entries = []
        else:
            print("Val directory not found, skipping validation set")
        
        # Always save val entries as backup, even if empty
        self.save_entries_as_backup(val_entries, output_dir, "val")
###########################################################################################
        # Convert entries to dict format
        train_data = self.entries_to_dict(train_entries)
        val_data = self.entries_to_dict(val_entries) if val_entries else {}

        # Try to create the dataset with error handling
        try:
            dataset = self.create_local_dataset(
                train_data=train_data,
                val_data=val_data,
                output_dir=os.path.join(output_dir, "hf_dataset")
            )
            print("Dataset created successfully!")
            return dataset
        except Exception as e:
            print(f"Error creating dataset: {e}")
            print("Entries have been saved as backup. You can recover them from the backup_entries directory.")
            print(f"Backup location: {os.path.join(output_dir, 'backup_entries')}")
            raise



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate and save HuggingFace dataset from image/mask pairs with template extraction.")
    parser.add_argument("--dataset_root_dir", type=str, default="/projects/p32685/RL_dataset/acdc", help="Root directory containing dataset_info.json, train/, and val/ folders")
    parser.add_argument("--output_dir", type=str, default="/projects/p32685/RL_dataset/acdc_processed", help="Output HuggingFace dataset directory")
    parser.add_argument("--device", type=str, default="cuda:0", help="Device to run inference on")
    parser.add_argument("--subset_size", type=int, default=0, help="Number of samples to process per split (0 means all samples)")
    parser.add_argument("--recover", action="store_true", help="Recover dataset from backup entries instead of processing from scratch")
    
    # Template extraction arguments
    parser.add_argument("--extract_templates", action="store_true", help="Extract template features from training images")
    parser.add_argument("--num_templates", type=int, default=20, help="Number of template images to extract features from (-1 for all matching patients)")
    parser.add_argument("--template_patients", type=str, default="patient001,patient002,patient003,patient004", 
                       help="Comma-separated list of patient IDs to use as templates")
    parser.add_argument("--modality", type=str, default="MRI", help="Dataset modality (e.g., MRI, CT)")
    
    args = parser.parse_args()

    # Create GenHFDataset instance
    gen_dataset = GenHFDataset(
        device=args.device,
    )
    
    if args.recover:
        # Recover from backup
        print("Recovering dataset from backup...")
        dataset = gen_dataset.recover_from_backup(
            output_dir=args.output_dir
        )
    else:
        dataset = gen_dataset.process_full_dataset(
            dataset_root_dir=args.dataset_root_dir,
            output_dir=args.output_dir,
            subset_size=args.subset_size
        )
    
    # Extract templates if requested
    if args.extract_templates:
        print("\n" + "="*60)
        print("TEMPLATE FEATURE EXTRACTION")
        print("="*60)
        
        # Parse patient filter
        patient_filter = None
        if args.template_patients:
            patient_filter = [p.strip() for p in args.template_patients.split(',')]
        
        # Define paths
        train_image_dir = os.path.join(args.dataset_root_dir, "train", "images")
        train_mask_dir = os.path.join(args.dataset_root_dir, "train", "masks")
        template_output_dir = os.path.join(args.output_dir, "template")
        
        # Extract templates from training images
        template_data = extract_templates_from_dataset(
            image_dir=train_image_dir,
            mask_dir=train_mask_dir,
            num_templates=args.num_templates,
            modality=args.modality,
            patient_filter=patient_filter,
            device=args.device
        )
        
        if template_data:
            # Extract DINOv2 features
            results = extract_template_features(
                template_data_dict=template_data,
                output_dir=template_output_dir,
                device=args.device
            )
            
            if results:
                print("\n" + "="*60)
                print("TEMPLATE EXTRACTION COMPLETE")
                print("="*60)
                print(f"Total templates processed: {len(results)}")
                print(f"Template features saved to: {template_output_dir}")
        else:
            print("No templates extracted!")
    
    print("\n" + "="*60)
    print("ALL PROCESSING COMPLETE")
    print("="*60) 



# Dataset Structure Requirements
# Please organize all datasets into the following structure:
# dataset/
# ├── train/
# ├── val/
# └── dataset_info.json

# Required Components:
# Two subdirectories and one JSON file - train, val, and dataset_info.json
# dataset_info.json must include:
# "dataset_name": The name of the dataset (REQUIRED)
# "classes": List of class names (excluding background) OR
# "class_mapping": Dict mapping class_id to class_name (including background)

# Both train and val subdirectories must contain two folders:
# images/
# masks/

# Example structure:
# dataset/
# ├── train/
# │   ├── images/
# │   │   ├── image001.jpg
# │   │   └── image002.jpg
# │   └── masks/
# │       ├── image001.png
# │       └── image002.png
# ├── val/
# │   ├── images/
# │   └── masks/
# └── dataset_info.json (example below)
        # {
        #     "dataset_name": "ACDC",
        #     "description": "Automated Cardiac Diagnosis Challenge - cardiac MRI segmentation dataset",
        #     "modality": "MRI",
        #     "classes": ["right ventricle", "myocardium", "left ventricle"],
        #     "num_classes": 4,
        #     "class_mapping": {
        #         "0": "background",
        #         "1": "right ventricle",
        #         "2": "myocardium",
        #         "3": "left ventricle"
        #     }
        # }