|
7 | 7 | import re |
8 | 8 | import subprocess |
9 | 9 | from pathlib import Path |
| 10 | +from multiprocessing.pool import Pool |
10 | 11 |
|
11 | 12 |
|
12 | 13 | def get_args(): |
13 | 14 | parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, description=__doc__) |
14 | 15 |
|
15 | 16 | parser.add_argument('-o', '--output', type=Path, action='store', dest='outdir', metavar='OUTPUT_DIR', |
16 | 17 | default=Path('.'), help="Path to defacing outputs directory.") |
| 18 | + parser.add_argument('-n', '--n-cpus', type=int, default=1, |
| 19 | + help='Number of parallel processes to run. ' |
| 20 | + 'Defaults to 1, meaning "serial processing" when not provided.') |
17 | 21 | return parser.parse_args() |
18 | 22 |
|
19 | 23 |
|
@@ -50,8 +54,21 @@ def generate_3d_renders(defaced_img, render_outdir): |
50 | 54 | def main(): |
51 | 55 | args = get_args() |
52 | 56 | defaced_imgs = list(args.outdir.rglob('defaced.nii.gz')) |
53 | | - for img in defaced_imgs: |
54 | | - generate_3d_renders(img, img.parent) |
| 57 | + |
| 58 | + if args.n_cpus == 1: |
| 59 | + print(f'Running in Serial, one render at a time') |
| 60 | + for img in defaced_imgs: |
| 61 | + generate_3d_renders(img, img.parent) |
| 62 | + |
| 63 | + elif args.n_cpus > 1: |
| 64 | + print(f'Running in Parallel with {args.n_cpus} cores') |
| 65 | + # initialize pool |
| 66 | + with Pool(processes=args.n_cpus) as p: |
| 67 | + p.starmap( generate_3d_renders, zip( defaced_imgs, |
| 68 | + [img.parent for img in defaced_imgs] ) ) |
| 69 | + |
| 70 | + else: |
| 71 | + raise ValueError(f"Provided --n-cpus of {args.n_cpus} is not a valid number of CPUs.") |
55 | 72 |
|
56 | 73 | # prep for visual inspection using visualqc deface |
57 | 74 | print(f"Preparing for QC by visual inspection...\n") |
|
0 commit comments