improved argparse stuff

.gitignore

@@ -0,0 +1,4 @@
+venv/
+build/
+dist/
+livermask.egg-info/ 

livermask/livermask.py

@@ -10,120 +10,138 @@ import gdown
 from skimage.morphology import remove_small_holes, binary_dilation, binary_erosion, ball
 from skimage.measure import label, regionprops
 import warnings
+import argparse
+import pkg_resources
+
+
+# mute some warnings
 warnings.filterwarnings('ignore', '.*output shape of zoom.*')
 
 
 def intensity_normalization(volume, intensity_clipping_range):
-	result = np.copy(volume)
+    result = np.copy(volume)
 
-	result[volume < intensity_clipping_range[0]] = intensity_clipping_range[0]
-	result[volume > intensity_clipping_range[1]] = intensity_clipping_range[1]
+    result[volume < intensity_clipping_range[0]] = intensity_clipping_range[0]
+    result[volume > intensity_clipping_range[1]] = intensity_clipping_range[1]
 
-	min_val = np.amin(result)
-	max_val = np.amax(result)
-	if (max_val - min_val) != 0:
-		result = (result - min_val) / (max_val - min_val)
+    min_val = np.amin(result)
+    max_val = np.amax(result)
+    if (max_val - min_val) != 0:
+        result = (result - min_val) / (max_val - min_val)
 
-	return result
+    return result
 
 def post_process(pred):
-	return pred
+    return pred
 
 def get_model():
-	url = "https://drive.google.com/uc?id=12or5Q79at2BtLgQ7IaglNGPFGRlEgEHc"
-	output = "model.h5"
-	md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
-	gdown.cached_download(url, output, md5=md5) #, postprocess=gdown.extractall)
+    url = "https://drive.google.com/uc?id=12or5Q79at2BtLgQ7IaglNGPFGRlEgEHc"
+    output = "model.h5"
+    md5 = "ef5a6dfb794b39bea03f5496a9a49d4d"
+    gdown.cached_download(url, output, md5=md5) #, postprocess=gdown.extractall)
 
 def func(path, output):
 
-	cwd = "/".join(os.path.realpath(__file__).replace("\\", "/").split("/")[:-1]) + "/"
+    cwd = "/".join(os.path.realpath(__file__).replace("\\", "/").split("/")[:-1]) + "/"
 
-	name = cwd + "model.h5"
+    name = cwd + "model.h5"
 
-	# get model
-	get_model()
+    # get model
+    get_model()
 
-	# load model
-	model = load_model(name, compile=False)
+    # load model
+    model = load_model(name, compile=False)
 
-	print("preprocessing...")
-	nib_volume = nib.load(path)
-	new_spacing = [1., 1., 1.]
-	resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
-	data = resampled_volume.get_data().astype('float32')
+    print("preprocessing...")
+    nib_volume = nib.load(path)
+    new_spacing = [1., 1., 1.]
+    resampled_volume = resample_to_output(nib_volume, new_spacing, order=1)
+    data = resampled_volume.get_data().astype('float32')
 
-	curr_shape = data.shape
+    curr_shape = data.shape
 
-	# resize to get (512, 512) output images
-	img_size = 512
-	data = zoom(data, [img_size / data.shape[0], img_size / data.shape[1], 1.0], order=1)
+    # resize to get (512, 512) output images
+    img_size = 512
+    data = zoom(data, [img_size / data.shape[0], img_size / data.shape[1], 1.0], order=1)
 
-	# intensity normalization
-	intensity_clipping_range = [-150, 250] # HU clipping limits (Pravdaray's configs)
-	data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
+    # intensity normalization
+    intensity_clipping_range = [-150, 250] # HU clipping limits (Pravdaray's configs)
+    data = intensity_normalization(volume=data, intensity_clipping_range=intensity_clipping_range)
 
-	# fix orientation
-	data = np.rot90(data, k=1, axes=(0, 1))
-	data = np.flip(data, axis=0)
+    # fix orientation
+    data = np.rot90(data, k=1, axes=(0, 1))
+    data = np.flip(data, axis=0)
 
-	print("predicting...")
-	# predict on data
-	pred = np.zeros_like(data).astype(np.float32)
-	for i in tqdm(range(data.shape[-1]), "pred: "):
-		pred[..., i] = model.predict(np.expand_dims(np.expand_dims(np.expand_dims(data[..., i], axis=0), axis=-1), axis=0))[0, ..., 1]
-	del data 
+    print("predicting...")
+    # predict on data
+    pred = np.zeros_like(data).astype(np.float32)
+    for i in tqdm(range(data.shape[-1]), "pred: "):
+        pred[..., i] = model.predict(np.expand_dims(np.expand_dims(np.expand_dims(data[..., i], axis=0), axis=-1), axis=0))[0, ..., 1]
+    del data 
 
-	# threshold
-	pred = (pred >= 0.4).astype(int)
+    # threshold
+    pred = (pred >= 0.4).astype(int)
 
-	# fix orientation back
-	pred = np.flip(pred, axis=0)
-	pred = np.rot90(pred, k=-1, axes=(0, 1))
+    # fix orientation back
+    pred = np.flip(pred, axis=0)
+    pred = np.rot90(pred, k=-1, axes=(0, 1))
 
-	print("resize back...")
-	# resize back from 512x512
-	pred = zoom(pred, [curr_shape[0] / img_size, curr_shape[1] / img_size, 1.0], order=1)
-	pred = (pred >= 0.5).astype(np.float32)
+    print("resize back...")
+    # resize back from 512x512
+    pred = zoom(pred, [curr_shape[0] / img_size, curr_shape[1] / img_size, 1.0], order=1)
+    pred = (pred >= 0.5).astype(np.float32)
 
-	print("morphological post-processing...")
-	# morpological post-processing
-	# 1) first erode
-	pred = binary_erosion(pred.astype(bool), ball(3)).astype(np.float32)
+    print("morphological post-processing...")
+    # morpological post-processing
+    # 1) first erode
+    pred = binary_erosion(pred.astype(bool), ball(3)).astype(np.float32)
 
-	# 2) keep only largest connected component
-	labels = label(pred)
-	regions = regionprops(labels)
-	area_sizes = []
-	for region in regions:
-		area_sizes.append([region.label, region.area])
-	area_sizes = np.array(area_sizes)
-	tmp = np.zeros_like(pred)
-	tmp[labels == area_sizes[np.argmax(area_sizes[:, 1]), 0]] = 1
-	pred = tmp.copy()
-	del tmp, labels, regions, area_sizes
+    # 2) keep only largest connected component
+    labels = label(pred)
+    regions = regionprops(labels)
+    area_sizes = []
+    for region in regions:
+        area_sizes.append([region.label, region.area])
+    area_sizes = np.array(area_sizes)
+    tmp = np.zeros_like(pred)
+    tmp[labels == area_sizes[np.argmax(area_sizes[:, 1]), 0]] = 1
+    pred = tmp.copy()
+    del tmp, labels, regions, area_sizes
 
-	# 3) dilate
-	pred = binary_dilation(pred.astype(bool), ball(3))
+    # 3) dilate
+    pred = binary_dilation(pred.astype(bool), ball(3))
 
-	# 4) remove small holes
-	pred = remove_small_holes(pred.astype(bool), area_threshold=0.001*np.prod(pred.shape)).astype(np.float32)
+    # 4) remove small holes
+    pred = remove_small_holes(pred.astype(bool), area_threshold=0.001*np.prod(pred.shape)).astype(np.float32)
 
-	print("saving...")
-	pred = pred.astype(np.uint8)
-	img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
-	resampled_lab = resample_from_to(img, nib_volume, order=0)
-	nib.save(resampled_lab, output)
+    print("saving...")
+    pred = pred.astype(np.uint8)
+    img = nib.Nifti1Image(pred, affine=resampled_volume.affine)
+    resampled_lab = resample_from_to(img, nib_volume, order=0)
+    nib.save(resampled_lab, output)
 
 
 def main():
-	os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # disable GPU
+    # os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # disable GPU
+    version = pkg_resources.require("livermask")[0].version
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', metavar='--i', type=str, nargs='?',
+                    help="set path of which image to use.")
+    parser.add_argument('--output', metavar='--o', type=str, nargs='?',
+                    help="set path to store the output.")
+    parser.add_argument('--cpu', metavar='--o', action='store_true',
+                    help="force using the CPU even if a GPU is available.")
+    parser.add_argument('--version', metavar='--v',
+                    help='shows the current version of livermask.', version=version)
+    ret = parser.parse_args(sys.argv[1:]); print(ret)
 
-	path = sys.argv[1]
-	output = sys.argv[2]
+    # fix paths
+    ret.input = ret.input.replace("\\", "/")
+    ret.output = ret.output.replace("\\", "/")
 
-	func(path, output)
+    func(*vars(ret).values())
 
 
 if __name__ == "__main__":
-	main()
+    main()

setup.py

@@ -20,8 +20,8 @@ setuptools.setup(
         ]
      },
      install_requires=[
-        'tensorflow==1.13.1',
-        'numpy',
+        'numpy'
+        'tensorflow==2.6',
         'scipy',
         'tqdm',
         'nibabel',