File size: 3,948 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 | #!/usr/bin/env python3
"""
测试 Step 3 (precompute template matches) 的脚本
用法: python test_step3.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_step3_precompute
import argparse
def main():
parser = argparse.ArgumentParser(description="Test Step 3 (precompute template matches)")
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)")
parser.add_argument("--device", type=str, default="cuda:0",
help="Device for template precomputation (default: cuda:0)")
parser.add_argument("--topk", type=int, default=1,
help="Top-k matches per sample (default: 1)")
args = parser.parse_args()
# 确定 processed 目录（Step 2 的输出）
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
processed_dir = os.path.join(args.output_base_dir, f"MSD_{task_num.lower()}",
f"{args.dataset_name}_{args.template_percent}percent_processed")
else:
processed_dir = os.path.join(args.output_base_dir,
f"{args.dataset_name.lower()}_{args.template_percent}percent_processed")
# 检查 processed 目录是否存在
if not os.path.exists(processed_dir):
print(f"✗ Error: Processed directory not found: {processed_dir}")
print(f" Make sure Step 2 has completed successfully.")
sys.exit(1)
# 检查必要的子目录
hf_dataset_path = os.path.join(processed_dir, "hf_dataset")
template_dir = os.path.join(processed_dir, "template")
backup_dir = os.path.join(processed_dir, "backup_entries")
missing_dirs = []
if not os.path.isdir(hf_dataset_path):
missing_dirs.append(f"hf_dataset ({hf_dataset_path})")
if not os.path.isdir(template_dir):
missing_dirs.append(f"template ({template_dir})")
if not os.path.isdir(backup_dir):
missing_dirs.append(f"backup_entries ({backup_dir})")
if missing_dirs:
print(f"✗ Error: Missing required directories:")
for d in missing_dirs:
print(f" - {d}")
print(f" Make sure Step 2 has completed successfully.")
sys.exit(1)
print("="*80)
print("TESTING STEP 3: PRECOMPUTE TEMPLATE MATCHES")
print("="*80)
print(f"Dataset: {args.dataset_name}")
print(f"Processed directory: {processed_dir}")
print(f"Device: {args.device}")
print(f"Top-k matches: {args.topk}")
print("="*80)
try:
run_step3_precompute(
processed_dir=processed_dir,
device=args.device,
topk=args.topk
)
print("\n" + "="*80)
print("✓ STEP 3 TEST PASSED")
print("="*80)
print(f"Processed directory: {processed_dir}")
return 0
except Exception as e:
print("\n" + "="*80)
print("✗ STEP 3 TEST FAILED")
print("="*80)
print(f"Error: {e}")
import traceback
traceback.print_exc()
return 1
if __name__ == "__main__":
sys.exit(main())
|