diff --git a/paddleocr.py b/paddleocr.py index c5da7248d2cc7d778758a87309cfeaedcbd8ceb5..1f358061eaaf8825f61dd01b2204c1036c2482a2 100644 --- a/paddleocr.py +++ b/paddleocr.py @@ -179,7 +179,7 @@ def parse_args(mMain=True, add_help=True): parser.add_argument("--use_gpu", type=str2bool, default=True) parser.add_argument("--ir_optim", type=str2bool, default=True) parser.add_argument("--use_tensorrt", type=str2bool, default=False) - parser.add_argument("--gpu_mem", type=int, default=8000) + parser.add_argument("--gpu_mem", type=int, default=500) # params for text detector parser.add_argument("--image_dir", type=str) @@ -232,7 +232,7 @@ def parse_args(mMain=True, add_help=True): use_gpu=True, ir_optim=True, use_tensorrt=False, - gpu_mem=8000, + gpu_mem=500, image_dir='', det_algorithm='DB', det_model_dir=None, diff --git a/ppocr/utils/utility.py b/ppocr/utils/utility.py index 7bb4c906d298af54ed56e2805f487a2c22d1894b..3f6255b0a9f79cc820b42569f18f82cf40ca3b7a 100755 --- a/ppocr/utils/utility.py +++ b/ppocr/utils/utility.py @@ -52,12 +52,14 @@ def get_image_file_list(img_file): raise Exception("not found any img file in {}".format(img_file)) img_end = {'jpg', 'bmp', 'png', 'jpeg', 'rgb', 'tif', 'tiff', 'gif', 'GIF'} - if os.path.isfile(img_file) and imghdr.what(img_file) in img_end: + if os.path.isfile(img_file) and os.path.splitext(img_file)[-1][1:].lower( + ) in img_end: imgs_lists.append(img_file) elif os.path.isdir(img_file): for single_file in os.listdir(img_file): file_path = os.path.join(img_file, single_file) - if os.path.isfile(file_path) and imghdr.what(file_path) in img_end: + if os.path.isfile(file_path) and os.path.splitext(file_path)[-1][ + 1:].lower() in img_end: imgs_lists.append(file_path) if len(imgs_lists) == 0: raise Exception("not found any img file in {}".format(img_file)) @@ -77,4 +79,4 @@ def check_and_read_gif(img_path): frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB) imgvalue = frame[:, :, ::-1] return imgvalue, True - return None, False \ No newline at end of file + return None, False