andreped commited on
Commit
05cdb10
1 Parent(s): 9f416d3

Added CLI arg to disable GPU + bump v1.3.0

Browse files
Files changed (2) hide show
  1. lungtumormask/__main__.py +6 -2
  2. setup.py +1 -1
lungtumormask/__main__.py CHANGED
@@ -19,9 +19,13 @@ def main():
19
  help='which radius to use for morphological post-processing segmentation smoothing.')
20
  parser.add_argument('--batch-size', metavar='batch-size', type=int, default=5,
21
  help='which batch size to use for lungmask inference.')
 
22
 
23
- argsin = sys.argv[1:]
24
- args = parser.parse_args(argsin)
 
 
 
25
 
26
  # import method here to enable faster testing
27
  from lungtumormask import mask
 
19
  help='which radius to use for morphological post-processing segmentation smoothing.')
20
  parser.add_argument('--batch-size', metavar='batch-size', type=int, default=5,
21
  help='which batch size to use for lungmask inference.')
22
+ parser.add_argument('--cpu', action='store_true', help='whether to force computation to happen on CPU only.')
23
 
24
+ args = parser.parse_args(sys.argv[1:])
25
+
26
+ # whether to force CPU computation
27
+ if args.cpu:
28
+ os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
29
 
30
  # import method here to enable faster testing
31
  from lungtumormask import mask
setup.py CHANGED
@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
3
  setup(
4
  name="lungtumormask",
5
  packages=find_packages(),
6
- version='1.2.1',
7
  author="Svein Ole M Sevle, Vemund Fredriksen, and André Pedersen",
8
  url="https://github.com/VemundFredriksen/LungTumorMask",
9
  license="MIT",
 
3
  setup(
4
  name="lungtumormask",
5
  packages=find_packages(),
6
+ version='1.3.0',
7
  author="Svein Ole M Sevle, Vemund Fredriksen, and André Pedersen",
8
  url="https://github.com/VemundFredriksen/LungTumorMask",
9
  license="MIT",