code/VESSA_data_preprocess/process_all_msd_tasks.py

#!/usr/bin/env python3
"""
Batch process all 10 MSD tasks.

This script processes all MSD tasks sequentially:
- Converts each task to HF format
- Creates 10% processed datasets
- Output structure:
  RL_dataset/
    MSD_task01/
      Task01_BrainTumour/
      Task01_BrainTumour_10percent_processed/
    MSD_task02/
      Task02_Heart/
      Task02_Heart_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))

# MSD task configurations
MSD_TASKS = [
    {
        "name": "Task01_BrainTumour",
        "raw_dir": "/projects/p32685/MSD/Task01_BrainTumour",
        "description": "Brain Tumour Segmentation"
    },
    {
        "name": "Task02_Heart",
        "raw_dir": "/projects/p32685/MSD/Task02_Heart",
        "description": "Heart Segmentation"
    },
    {
        "name": "Task03_Liver",
        "raw_dir": "/projects/p32685/MSD/Task03_Liver",
        "description": "Liver Segmentation"
    },
    {
        "name": "Task04_Hippocampus",
        "raw_dir": "/projects/p32685/MSD/Task04_Hippocampus",
        "description": "Hippocampus Segmentation"
    },
    {
        "name": "Task05_Prostate",
        "raw_dir": "/projects/p32685/MSD/Task05_Prostate",
        "description": "Prostate Segmentation"
    },
    {
        "name": "Task06_Lung",
        "raw_dir": "/projects/p32685/MSD/Task06_Lung",
        "description": "Lung Tumor Segmentation"
    },
    {
        "name": "Task07_Pancreas",
        "raw_dir": "/projects/p32685/MSD/Task07_Pancreas",
        "description": "Pancreas Segmentation"
    },
    {
        "name": "Task08_HepaticVessel",
        "raw_dir": "/projects/p32685/MSD/Task08_HepaticVessel",
        "description": "Hepatic Vessel Segmentation"
    },
    {
        "name": "Task09_Spleen",
        "raw_dir": "/projects/p32685/MSD/Task09_Spleen",
        "description": "Spleen Segmentation"
    },
    {
        "name": "Task10_Colon",
        "raw_dir": "/projects/p32685/MSD/Task10_Colon",
        "description": "Colon Tumor Segmentation"
    }
]

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


def process_task(task_config, skip_convert=False, skip_packing=False, skip_precompute=False):
    """
    Process a single MSD task.
    
    Args:
        task_config: Dictionary with task configuration
        skip_convert: Skip step 1 (conversion)
        skip_packing: Skip step 2 (packing)
        skip_precompute: Skip step 3 (precompute)
    """
    task_name = task_config["name"]
    raw_dir = task_config["raw_dir"]
    
    print("\n" + "="*80)
    print(f"Processing {task_name}: {task_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", task_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✓ {task_name} processed successfully!")
            return True
        else:
            print(f"\n✗ {task_name} failed with return code {result.returncode}")
            return False
    except subprocess.CalledProcessError as e:
        print(f"\n✗ {task_name} failed: {e}")
        return False
    except Exception as e:
        print(f"\n✗ {task_name} failed with exception: {e}")
        import traceback
        traceback.print_exc()
        return False


def main():
    """Main function to process all MSD tasks."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Batch process all MSD tasks")
    parser.add_argument("--start_from", type=int, default=1,
                       help="Start from task number (1-10)")
    parser.add_argument("--end_at", type=int, default=10,
                       help="End at task number (1-10)")
    parser.add_argument("--skip_convert", action="store_true",
                       help="Skip step 1 (conversion) for all tasks")
    parser.add_argument("--skip_packing", action="store_true",
                       help="Skip step 2 (packing) for all tasks")
    parser.add_argument("--skip_precompute", action="store_true",
                       help="Skip step 3 (precompute) for all tasks")
    parser.add_argument("--device", type=str, default="cuda:0",
                       help="Device for template precomputation")
    
    args = parser.parse_args()
    
    global DEVICE
    DEVICE = args.device
    
    print("╔" + "="*78 + "╗")
    print("║" + " "*15 + "BATCH PROCESS ALL MSD TASKS" + " "*38 + "║")
    print("╔" + "="*78 + "╗")
    print(f"\nTotal tasks: {len(MSD_TASKS)}")
    print(f"Processing tasks: {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 tasks
    success_count = 0
    failed_tasks = []
    
    for i, task_config in enumerate(MSD_TASKS, start=1):
        if i < args.start_from or i > args.end_at:
            continue
        
        print(f"\n{'='*80}")
        print(f"Task {i}/{len(MSD_TASKS)}: {task_config['name']}")
        print(f"{'='*80}")
        
        success = process_task(
            task_config,
            skip_convert=args.skip_convert,
            skip_packing=args.skip_packing,
            skip_precompute=args.skip_precompute
        )
        
        if success:
            success_count += 1
        else:
            failed_tasks.append(task_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} tasks")
    
    if failed_tasks:
        print(f"\nFailed tasks:")
        for task in failed_tasks:
            print(f"  - {task}")
    else:
        print("\n✓ All tasks processed successfully!")
    
    print(f"\nOutput structure:")
    print(f"  {OUTPUT_BASE_DIR}/")
    for i in range(args.start_from, args.end_at + 1):
        task_name = MSD_TASKS[i-1]['name']
        task_num = f"task{i:02d}"
        print(f"    MSD_{task_num}/")
        print(f"      {task_name}/")
        print(f"      {task_name}_{TEMPLATE_PERCENT}percent_processed/")


if __name__ == "__main__":
    main()