Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Adding thread control #39

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Adding thread control #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a5c35d8
Adding thread control
jrdalenberg Dec 5, 2022
a365d0e
Fix threads check
jrdalenberg Mar 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 15 additions & 1 deletion HD_BET/hd-bet
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,11 +1,13 @@
#!/usr/bin/env python

import os
import multiprocessing
from HD_BET.run import run_hd_bet
from HD_BET.utils import maybe_mkdir_p, subfiles
import HD_BET



if __name__ == "__main__":
print("\n########################")
print("If you are using hd-bet, please cite the following paper:")
Expand Down Expand Up @@ -33,6 +35,10 @@ if __name__ == "__main__":
'\'cpu\' to run on CPU. When using CPU you should '
'consider disabling tta. Default for -device is: 0',
required=False)
parser.add_argument('-threads', default=0, type=int, help='used to set the number of cpu threads. '
'Must be either int or str. Use 0 for max available cpus.'
'Default for -threads is: 0',
required=False)
parser.add_argument('-tta', default=1, required=False, type=int, help='whether to use test time data augmentation '
'(mirroring). 1= True, 0=False. Disable this '
'if you are using CPU to speed things up! '
Expand All @@ -58,6 +64,7 @@ if __name__ == "__main__":

mode = args.mode
device = args.device
threads = args.threads
tta = args.tta
pp = args.pp
save_mask = args.save_mask
Expand Down Expand Up @@ -90,6 +97,13 @@ if __name__ == "__main__":
output_files = [output_file_or_dir]
input_files = [input_file_or_dir]

max_cpu_count = multiprocessing.cpu_count()
if threads == 0:
threads = max_cpu_count
elif threads not in range(1, max_cpu_count + 1):
raise ValueError(f"Unknown value for threads: {threads}. Expected: value between 1 and maximum number of available threads ({max_cpu_count}) \
Tip: A value of 0 will pick the maximum available number.")

if tta == 0:
tta = False
elif tta == 1:
Expand Down Expand Up @@ -118,4 +132,4 @@ if __name__ == "__main__":
else:
raise ValueError("Unknown value for pp: %s. Expected: 0 or 1" % str(pp))

run_hd_bet(input_files, output_files, mode, config_file, device, pp, tta, save_mask, overwrite_existing)
run_hd_bet(input_files, output_files, mode, config_file, device, threads, pp, tta, save_mask, overwrite_existing)
4 changes: 3 additions & 1 deletion HD_BET/run.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ def apply_bet(img, bet, out_fname):


def run_hd_bet(mri_fnames, output_fnames, mode="accurate", config_file=os.path.join(HD_BET.__path__[0], "config.py"), device=0,
postprocess=False, do_tta=True, keep_mask=True, overwrite=True):
threads=0, postprocess=False, do_tta=True, keep_mask=True, overwrite=True):
"""

:param mri_fnames: str or list/tuple of str
Expand All @@ -35,6 +35,8 @@ def run_hd_bet(mri_fnames, output_fnames, mode="accurate", config_file=os.path.j
:return:
"""

torch.set_num_threads(threads)

list_of_param_files = []

if mode == 'fast':
Expand Down