File size: 3,078 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 | #!/usr/bin/env python3
"""
测试 Step 2 (packing) 的脚本
用法: python test_step2.py --dataset_name Task01_BrainTumour
"""
import sys
import os
from pathlib import Path
# 添加当前目录到路径
SCRIPT_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(SCRIPT_DIR))
from end_to_end_preprocessing import run_step2_packing
import argparse
def main():
parser = argparse.ArgumentParser(description="Test Step 2 (packing)")
parser.add_argument("--dataset_name", type=str, required=True,
help="Dataset name (e.g., Task01_BrainTumour)")
parser.add_argument("--output_base_dir", type=str, default="/projects/p32685/RL_dataset",
help="Base output directory")
parser.add_argument("--template_percent", type=int, default=10,
help="Template percentage (default: 10)")
args = parser.parse_args()
# 确定输入目录（HF 格式）
if args.dataset_name.startswith("Task") and args.dataset_name.endswith(("BrainTumour", "Heart", "Liver", "Hippocampus",
"Prostate", "Lung", "Pancreas", "HepaticVessel",
"Spleen", "Colon")):
task_num = args.dataset_name.split("_")[0] # Task01
dataset_dir = os.path.join(args.output_base_dir, f"MSD_{task_num.lower()}", args.dataset_name)
else:
dataset_dir = os.path.join(args.output_base_dir, args.dataset_name.lower())
# 检查输入目录是否存在
if not os.path.exists(dataset_dir):
print(f"✗ Error: Dataset directory not found: {dataset_dir}")
print(f" Make sure Step 1 has completed successfully.")
sys.exit(1)
# 检查 dataset_info.json 是否存在
dataset_info_path = os.path.join(dataset_dir, "dataset_info.json")
if not os.path.exists(dataset_info_path):
print(f"✗ Error: dataset_info.json not found in {dataset_dir}")
print(f" Step 1 may not have completed successfully.")
sys.exit(1)
print("="*80)
print("TESTING STEP 2: PACKING")
print("="*80)
print(f"Dataset: {args.dataset_name}")
print(f"Input directory: {dataset_dir}")
print(f"Output base: {args.output_base_dir}")
print(f"Template percent: {args.template_percent}%")
print("="*80)
try:
processed_dir = run_step2_packing(
dataset_dir=dataset_dir,
output_base_dir=args.output_base_dir,
dataset_name=args.dataset_name,
template_percent=args.template_percent
)
print("\n" + "="*80)
print("✓ STEP 2 TEST PASSED")
print("="*80)
print(f"Output directory: {processed_dir}")
return 0
except Exception as e:
print("\n" + "="*80)
print("✗ STEP 2 TEST FAILED")
print("="*80)
print(f"Error: {e}")
import traceback
traceback.print_exc()
return 1
if __name__ == "__main__":
sys.exit(main())
|