code/VESSA/data_preprocessing/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


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
        """
        features = 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
            "template": Value("string"),  # JSON string containing template list: [{"template_id": str, "similarity": float, "overlap": [int, ...]}, ...]
        })

        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")
        
        # Create intermediate data directories under output_base_dir
        split_output_dir = os.path.join(output_base_dir, split_name)
        os.makedirs(split_output_dir, exist_ok=True)
        
        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
                'template': None,
            }
            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():
                if k == 'template':
                    result[k].append(json.dumps(v))
                else:
                    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.")
    parser.add_argument("--dataset_root_dir", type=str, default="/projects/p32591/RL_dataset/lungx", help="Root directory containing dataset_info.json, train/, and val/ folders")
    parser.add_argument("--output_dir", type=str, default="/projects/p32591/RL_dataset/lungx_10percent_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_false", help="Recover dataset from backup entries instead of processing from scratch")
    args = parser.parse_args()

    # Create GenHFDataset instance
    gen_dataset = GenHFDataset(
        device=args.device,
    )
    
    if args.recover:
        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
    #     ) 



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