diff --git a/ppocr/utils/utility.py b/ppocr/utils/utility.py index c824d4404ad5f1fbee0d672fbc06d365f7d4eade..687e2acbc8021ddaad074c5445478cbba2ad2589 100755 --- a/ppocr/utils/utility.py +++ b/ppocr/utils/utility.py @@ -14,6 +14,7 @@ import logging import os +import imghdr def initial_logger(): @@ -61,13 +62,14 @@ def get_image_file_list(img_file): if img_file is None or not os.path.exists(img_file): raise Exception("not found any img file in {}".format(img_file)) - img_end = ['jpg', 'png', 'jpeg', 'JPEG', 'JPG', 'bmp'] - if os.path.isfile(img_file) and img_file.split('.')[-1] in img_end: + img_end = {'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff'} + if os.path.isfile(img_file) and imghdr.what(img_file) in img_end: imgs_lists.append(img_file) elif os.path.isdir(img_file): for single_file in os.listdir(img_file): - if single_file.split('.')[-1] in img_end: - imgs_lists.append(os.path.join(img_file, single_file)) + file_path = os.path.join(img_file, single_file) + if imghdr.what(file_path) in img_end: + imgs_lists.append(file_path) if len(imgs_lists) == 0: raise Exception("not found any img file in {}".format(img_file)) return imgs_lists