Spaces:
Runtime error
Runtime error
File size: 2,311 Bytes
906fcb9 caf6ee7 70d0e22 caf6ee7 1baebae caf6ee7 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae 906fcb9 70d0e22 906fcb9 caf6ee7 906fcb9 1baebae 906fcb9 70d0e22 906fcb9 1baebae 906fcb9 1baebae 906fcb9 1baebae | 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 | import argparse
import logging
import os
from pathlib import Path
import yaml
from src.preprocessing.generate_heatmap import get_heatmap
from src.preprocessing.histogram_match import histmatch
from src.preprocessing.prostate_mask import get_segmask
from src.preprocessing.register_and_crop import register_files
from src.utils import setup_logging, validate_steps
def parse_args():
parser = argparse.ArgumentParser(description="File preprocessing")
parser.add_argument("--config", type=str, help="Path to YAML config file")
parser.add_argument(
"--steps",
nargs="+", # ← list of strings
choices=[
"register_and_crop",
"histogram_match",
"get_segmentation_mask",
"get_heatmap",
], # ← restrict allowed values
required=True,
help="Steps to execute (one or more)",
)
parser.add_argument("--t2_dir", default=None, help="Path to T2W files")
parser.add_argument("--dwi_dir", default=None, help="Path to DWI files")
parser.add_argument("--adc_dir", default=None, help="Path to ADC files")
parser.add_argument("--seg_dir", default=None, help="Path to segmentation masks")
parser.add_argument("--output_dir", default=None, help="Path to output folder")
parser.add_argument(
"--margin", default=0.2, type=float, help="Margin to center crop the images"
)
parser.add_argument("--project_dir", default=None, help="Project directory")
args = parser.parse_args()
if args.config:
with open(args.config) as config_file:
config = yaml.safe_load(config_file)
args.__dict__.update(config)
return args
if __name__ == "__main__":
args = parse_args()
if args.project_dir is None:
args.project_dir = Path(__file__).resolve().parent # Set project directory
FUNCTIONS = {
"register_and_crop": register_files,
"histogram_match": histmatch,
"get_segmentation_mask": get_segmask,
"get_heatmap": get_heatmap,
}
args.logfile = os.path.join(args.output_dir, "preprocessing.log")
setup_logging(args.logfile)
logging.info("Starting preprocessing")
validate_steps(args.steps)
for step in args.steps:
func = FUNCTIONS[step]
args = func(args)
|