code/VESSA_data_preprocess/process_six_tasks.py

#!/usr/bin/env python3
"""
Batch process six biomedical datasets.

This script processes six datasets sequentially:
- BraTS2020
- SPIDER
- STS-2D-Tooth
- STS-3D-Tooth
- TotalSegmentator
- TotalSegmentatorMRI

Each dataset goes through:
- Converts BiomedParse format to HF format
- Creates processed datasets with templates
- Output structure:
  RL_dataset/
    brats2020/
      brats2020_10percent_processed/
    spider/
      spider_10percent_processed/
    ...
"""

import os
import sys
import subprocess
from pathlib import Path

# Add script directory to path
SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))

# Six dataset configurations
SIX_DATASETS = [
    {
        "name": "BraTS2020",
        "raw_dir": "/projects/p32685/datasets/biomedparse/BraTS2020",
        "description": "BraTS 2020 - Brain Tumour Segmentation"
    },
    {
        "name": "SPIDER",
        "raw_dir": "/projects/p32685/datasets/biomedparse/SPIDER",
        "description": "SPIDER - Spine Segmentation Dataset"
    },
    {
        "name": "STS-2D-Tooth",
        "raw_dir": "/projects/p32685/datasets/biomedparse/STS-2D-Tooth",
        "description": "STS-2D-Tooth - 2D Tooth Segmentation Dataset"
    },
    {
        "name": "STS-3D-Tooth",
        "raw_dir": "/projects/p32685/datasets/biomedparse/STS-3D-Tooth",
        "description": "STS-3D-Tooth - 3D Tooth Segmentation Dataset"
    },
    {
        "name": "TotalSegmentator",
        "raw_dir": "/projects/p32685/datasets/biomedparse/TotalSegmentator",
        "description": "TotalSegmentator v2.0.1 - 117 anatomical structures segmentation"
    },
    {
        "name": "TotalSegmentatorMRI",
        "raw_dir": "/projects/p32685/datasets/biomedparse/TotalSegmentatorMRI",
        "description": "TotalSegmentator MRI v2.0.0 - 50 anatomical structures segmentation"
    }
]

# Configuration
OUTPUT_BASE_DIR = "/projects/p32685/RL_dataset"
TEMPLATE_PERCENT = 10
DEVICE = "cuda:0"


def process_dataset(dataset_config, skip_convert=False, skip_packing=False, skip_precompute=False):
    """
    Process a single dataset.
    
    Args:
        dataset_config: Dictionary with dataset configuration
        skip_convert: Skip step 1 (conversion)
        skip_packing: Skip step 2 (packing)
        skip_precompute: Skip step 3 (precompute)
    """
    dataset_name = dataset_config["name"]
    raw_dir = dataset_config["raw_dir"]
    
    print("\n" + "="*80)
    print(f"Processing {dataset_name}: {dataset_config['description']}")
    print("="*80)
    print(f"Raw data: {raw_dir}")
    print(f"Output base: {OUTPUT_BASE_DIR}")
    print(f"Template percent: {TEMPLATE_PERCENT}%")
    
    # Check if raw directory exists
    if not os.path.exists(raw_dir):
        print(f"✗ Error: Raw data directory does not exist: {raw_dir}")
        return False
    
    # Build command
    cmd = [
        "python3",
        os.path.join(SCRIPT_DIR, "end_to_end_preprocessing.py"),
        "--dataset_name", dataset_name,
        "--template_percent", str(TEMPLATE_PERCENT),
        "--raw_data_dir", raw_dir,
        "--output_base_dir", OUTPUT_BASE_DIR,
        "--device", DEVICE
    ]
    
    if skip_convert:
        cmd.append("--skip_convert")
    if skip_packing:
        cmd.append("--skip_packing")
    if skip_precompute:
        cmd.append("--skip_precompute")
    
    print(f"\nRunning: {' '.join(cmd)}")
    
    try:
        result = subprocess.run(cmd, check=True, cwd=SCRIPT_DIR)
        if result.returncode == 0:
            print(f"\n✓ {dataset_name} processed successfully!")
            return True
        else:
            print(f"\n✗ {dataset_name} failed with return code {result.returncode}")
            return False
    except subprocess.CalledProcessError as e:
        print(f"\n✗ {dataset_name} failed: {e}")
        return False
    except Exception as e:
        print(f"\n✗ {dataset_name} failed with exception: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Main function to process all six datasets."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Batch process six biomedical datasets")
    parser.add_argument("--start_from", type=int, default=1,
                       help="Start from dataset number (1-6)")
    parser.add_argument("--end_at", type=int, default=6,
                       help="End at dataset number (1-6)")
    parser.add_argument("--skip_convert", action="store_true",
                       help="Skip step 1 (conversion) for all datasets")
    parser.add_argument("--skip_packing", action="store_true",
                       help="Skip step 2 (packing) for all datasets")
    parser.add_argument("--skip_precompute", action="store_true",
                       help="Skip step 3 (precompute) for all datasets")
    parser.add_argument("--device", type=str, default="cuda:0",
                       help="Device for template precomputation")
    parser.add_argument("--template_percent", type=int, default=10,
                       help="Percentage of training data to use as templates (default: 10)")
    
    args = parser.parse_args()
    
    global DEVICE, TEMPLATE_PERCENT
    DEVICE = args.device
    TEMPLATE_PERCENT = args.template_percent
    
    print("╔" + "="*78 + "╗")
    print("║" + " "*20 + "BATCH PROCESS SIX DATASETS" + " "*33 + "║")
    print("╔" + "="*78 + "╗")
    print(f"\nTotal datasets: {len(SIX_DATASETS)}")
    print(f"Processing datasets: {args.start_from} to {args.end_at}")
    print(f"Output base: {OUTPUT_BASE_DIR}")
    print(f"Template percent: {TEMPLATE_PERCENT}%")
    print(f"Device: {DEVICE}")
    
    if args.skip_convert:
        print("⚠ Skipping Step 1 (conversion)")
    if args.skip_packing:
        print("⚠ Skipping Step 2 (packing)")
    if args.skip_precompute:
        print("⚠ Skipping Step 3 (precompute)")
    
    # Process datasets
    success_count = 0
    failed_datasets = []
    
    for i, dataset_config in enumerate(SIX_DATASETS, start=1):
        if i < args.start_from or i > args.end_at:
            continue
        
        print(f"\n{'='*80}")
        print(f"Dataset {i}/{len(SIX_DATASETS)}: {dataset_config['name']}")
        print(f"{'='*80}")
        
        success = process_dataset(
            dataset_config,
            skip_convert=args.skip_convert,
            skip_packing=args.skip_packing,
            skip_precompute=args.skip_precompute
        )
        
        if success:
            success_count += 1
        else:
            failed_datasets.append(dataset_config['name'])
    
    # Summary
    print("\n" + "="*80)
    print("BATCH PROCESSING SUMMARY")
    print("="*80)
    print(f"Successfully processed: {success_count}/{args.end_at - args.start_from + 1} datasets")
    
    if failed_datasets:
        print(f"\nFailed datasets:")
        for dataset in failed_datasets:
            print(f"  - {dataset}")
    else:
        print("\n✓ All datasets processed successfully!")
    
    print(f"\nOutput structure:")
    print(f"  {OUTPUT_BASE_DIR}/")
    for i in range(args.start_from, args.end_at + 1):
        dataset_name = SIX_DATASETS[i-1]['name']
        dataset_name_lower = dataset_name.lower()
        print(f"    {dataset_name_lower}/")
        print(f"      {dataset_name_lower}_{TEMPLATE_PERCENT}percent_processed/")


if __name__ == "__main__":
    main()