From 664a441a2de9c388ed1083795902a06e4d5ca2f6 Mon Sep 17 00:00:00 2001 From: Konstantin Baierer Date: Fri, 23 Oct 2020 11:30:31 +0200 Subject: [PATCH] glob *.h5 files only in model dir, fix #7 --- sbb_binarize/sbb_binarize.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/sbb_binarize/sbb_binarize.py b/sbb_binarize/sbb_binarize.py index a664d6d..8775836 100644 --- a/sbb_binarize/sbb_binarize.py +++ b/sbb_binarize/sbb_binarize.py @@ -3,7 +3,8 @@ Tool to load model and binarize a given image. """ import sys -from os import listdir, environ, devnull +from glob import glob +from os import environ, devnull from os.path import join from warnings import catch_warnings, simplefilter @@ -191,7 +192,7 @@ class SbbBinarizer: if image_path is not None: image = cv2.imread(image) self.start_new_session() - list_of_model_files = listdir(self.model_dir) + list_of_model_files = glob('%s/*.h5' % self.model_dir) img_last = 0 for model_in in list_of_model_files: