#!/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()