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" # } # }