File size: 5,674 Bytes
01857e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | #!/usr/bin/env python3
"""
Batch process TotalSegmentator_reduced and TotalSegmentatorMRI_reduced to RL_dataset format.
Datasets:
- TotalSegmentator_reduced
- TotalSegmentatorMRI_reduced
Input: /projects/p32685/datasets/biomedparse/{TotalSegmentator_reduced,TotalSegmentatorMRI_reduced}
Output: /projects/p32685/RL_dataset/{totalsegmentator,totalsegmentatormri}_10percent_processed/
"""
import os
import sys
import subprocess
from pathlib import Path
SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))
TWO_DATASETS = [
{
"name": "TotalSegmentator_reduced",
"raw_dir": "/projects/p32685/datasets/biomedparse/TotalSegmentator_reduced",
"description": "TotalSegmentator reduced - 117 anatomical structures (every 10th slice)"
},
{
"name": "TotalSegmentatorMRI_reduced",
"raw_dir": "/projects/p32685/datasets/biomedparse/TotalSegmentatorMRI_reduced",
"description": "TotalSegmentator MRI reduced - 50 anatomical structures (every 10th slice)"
}
]
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."""
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}%")
if not os.path.exists(raw_dir):
print(f"✗ Error: Raw data directory does not exist: {raw_dir}")
return False
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():
import argparse
parser = argparse.ArgumentParser(description="Batch process TotalSegmentator reduced datasets to RL_dataset")
parser.add_argument("--start_from", type=int, default=1, help="Start from dataset number (1-2)")
parser.add_argument("--end_at", type=int, default=2, help="End at dataset number (1-2)")
parser.add_argument("--skip_convert", action="store_true", help="Skip step 1 (conversion)")
parser.add_argument("--skip_packing", action="store_true", help="Skip step 2 (packing)")
parser.add_argument("--skip_precompute", action="store_true", help="Skip step 3 (precompute)")
parser.add_argument("--device", type=str, default="cuda:0", help="Device for template precomputation")
parser.add_argument("--template_percent", type=int, default=10, help="Template percentage (default: 10)")
args = parser.parse_args()
global DEVICE, TEMPLATE_PERCENT
DEVICE = args.device
TEMPLATE_PERCENT = args.template_percent
print("â•”" + "="*78 + "â•—")
print("â•‘" + " "*10 + "BATCH PROCESS TOTAL SEGMENTATOR REDUCED DATASETS" + " "*18 + "â•‘")
print("â•”" + "="*78 + "â•—")
print(f"\nDatasets: TotalSegmentator_reduced, TotalSegmentatorMRI_reduced")
print(f"Processing: {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)")
success_count = 0
failed = []
for i, cfg in enumerate(TWO_DATASETS, start=1):
if i < args.start_from or i > args.end_at:
continue
print(f"\n{'='*80}")
print(f"Dataset {i}/2: {cfg['name']}")
print(f"{'='*80}")
if process_dataset(cfg, args.skip_convert, args.skip_packing, args.skip_precompute):
success_count += 1
else:
failed.append(cfg['name'])
print("\n" + "="*80)
print("BATCH PROCESSING SUMMARY")
print("="*80)
n_total = args.end_at - args.start_from + 1
print(f"Successfully processed: {success_count}/{n_total} datasets")
if failed:
print(f"Failed: {', '.join(failed)}")
else:
print("\n✓ All datasets processed successfully!")
print(f"\nOutput: {OUTPUT_BASE_DIR}/")
# Output folder names omit _reduced (totalsegmentator, totalsegmentatormri)
out_names = ["totalsegmentator", "totalsegmentatormri"]
for i in range(args.start_from, args.end_at + 1):
out = f"{out_names[i-1]}_{TEMPLATE_PERCENT}percent_processed"
print(f" {out}/")
if __name__ == "__main__":
main()
|