diff --git a/sbb_binarize/cli.py b/sbb_binarize/cli.py new file mode 100644 index 0000000..3de8820 --- /dev/null +++ b/sbb_binarize/cli.py @@ -0,0 +1,23 @@ +""" +sbb_binarize CLI +""" + +from argparse import ArgumentParser + +from .sbb_binarize import SbbBinarizer + +def main(): + parser = ArgumentParser() + + parser.add_argument('-i', '--image', default=None, help='image.') + parser.add_argument('-p', '--patches', default=False, help='by setting this parameter to true you let the model to see the image in patches.') + parser.add_argument('-s', '--save', default=False, help='save prediction with a given name here. The name and format should be given (outputname.tif).') + parser.add_argument('-m', '--model', default=None, help='models directory.') + + options = parser.parse_args() + + binarizer = SbbBinarizer(options.image, options.model, options.patches, options.save) + binarizer.run() + +if __name__ == "__main__": + main() diff --git a/sbb_binarize/sbb_binarize.py b/sbb_binarize/sbb_binarize.py index 7faf1be..6c8f8fd 100644 --- a/sbb_binarize/sbb_binarize.py +++ b/sbb_binarize/sbb_binarize.py @@ -2,7 +2,6 @@ Tool to load model and binarize a given image. """ -from argparse import ArgumentParser from os import listdir from os.path import join from warnings import catch_warnings, simplefilter @@ -16,7 +15,7 @@ import tensorflow as tf with catch_warnings(): simplefilter("ignore") -class sbb_binarize: +class SbbBinarizer: # TODO use True/False for patches def __init__(self, image, model, patches='false', save=None): @@ -218,21 +217,3 @@ class sbb_binarize: img_last = (img_last[:, :] == 0)*255 if self.save: cv2.imwrite(self.save, img_last) - -def main(): - parser = ArgumentParser() - - parser.add_argument('-i', '--image', dest='inp1', default=None, help='image.') - parser.add_argument('-p', '--patches', dest='inp3', default=False, help='by setting this parameter to true you let the model to see the image in patches.') - parser.add_argument('-s', '--save', dest='inp4', default=False, help='save prediction with a given name here. The name and format should be given (outputname.tif).') - parser.add_argument('-m', '--model', dest='inp2', default=None, help='models directory.') - - options = parser.parse_args() - - possibles = globals() - possibles.update(locals()) - x = sbb_binarize(options.inp1, options.inp2, options.inp3, options.inp4) - x.run() - -if __name__ == "__main__": - main() diff --git a/setup.py b/setup.py index 3796c51..980dd2f 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,7 @@ setup( install_requires=install_requires, entry_points={ 'console_scripts': [ - 'sbb_binarize=sbb_binarize.sbb_binarize:main', + 'sbb_binarize=sbb_binarize.cli:main', ] }, )