andreped commited on
Commit
df20506
1 Parent(s): 8315c39

Added feature to set batch-size through args [no ci]

Browse files
Files changed (1) hide show
  1. lungtumormask/__main__.py +7 -3
lungtumormask/__main__.py CHANGED
@@ -1,7 +1,6 @@
1
  import sys
2
  import argparse
3
  import os
4
- from lungtumormask import mask
5
 
6
  def path(string):
7
  if os.path.exists(string):
@@ -18,8 +17,13 @@ def main():
18
  help='which threshold to use for assigning voxel-wise classes.')
19
  parser.add_argument('--radius', metavar='radius', type=int, default=1,
20
  help='which radius to use for morphological post-processing segmentation smoothing.')
 
 
21
 
22
  argsin = sys.argv[1:]
23
  args = parser.parse_args(argsin)
24
-
25
- mask.mask(args.input, args.output, args.lung_filter, args.threshold, args.radius)
 
 
 
 
1
  import sys
2
  import argparse
3
  import os
 
4
 
5
  def path(string):
6
  if os.path.exists(string):
 
17
  help='which threshold to use for assigning voxel-wise classes.')
18
  parser.add_argument('--radius', metavar='radius', type=int, default=1,
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
28
+
29
+ mask.mask(args.input, args.output, args.lung_filter, args.threshold, args.radius, args.batch_size)