From 23658296ec85c36ed27195c9223c6dc94bce8714 Mon Sep 17 00:00:00 2001 From: littletomatodonkey Date: Mon, 27 Jul 2020 11:37:55 +0000 Subject: [PATCH] fix image list postfix --- ppocr/utils/utility.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/ppocr/utils/utility.py b/ppocr/utils/utility.py index c824d440..687e2acb 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 -- GitLab