Don't break overall processing on a bad image (#10216)
parent
fac03876f3
commit
1dad0a980f
18
paddleocr.py
18
paddleocr.py
|
@ -512,12 +512,12 @@ class PaddleOCR(predict_system.TextSystem):
|
|||
|
||||
def ocr(self, img, det=True, rec=True, cls=True):
|
||||
"""
|
||||
ocr with paddleocr
|
||||
OCR with PaddleOCR
|
||||
args:
|
||||
img: img for ocr, support ndarray, img_path and list or ndarray
|
||||
det: use text detection or not. If false, only rec will be exec. Default is True
|
||||
rec: use text recognition or not. If false, only det will be exec. Default is True
|
||||
cls: use angle classifier or not. Default is True. If true, the text with rotation of 180 degrees can be recognized. If no text is rotated by 180 degrees, use cls=False to get better performance. Text with rotation of 90 or 270 degrees can be recognized even if cls=False.
|
||||
img: img for OCR, support ndarray, img_path and list or ndarray
|
||||
det: use text detection or not. If False, only rec will be exec. Default is True
|
||||
rec: use text recognition or not. If False, only det will be exec. Default is True
|
||||
cls: use angle classifier or not. Default is True. If True, the text with rotation of 180 degrees can be recognized. If no text is rotated by 180 degrees, use cls=False to get better performance. Text with rotation of 90 or 270 degrees can be recognized even if cls=False.
|
||||
"""
|
||||
assert isinstance(img, (np.ndarray, list, str, bytes))
|
||||
if isinstance(img, list) and det == True:
|
||||
|
@ -525,7 +525,7 @@ class PaddleOCR(predict_system.TextSystem):
|
|||
exit(0)
|
||||
if cls == True and self.use_angle_cls == False:
|
||||
logger.warning(
|
||||
'Since the angle classifier is not initialized, the angle classifier will not be uesd during the forward process'
|
||||
'Since the angle classifier is not initialized, it will not be used during the forward process'
|
||||
)
|
||||
|
||||
img = check_img(img)
|
||||
|
@ -541,6 +541,9 @@ class PaddleOCR(predict_system.TextSystem):
|
|||
ocr_res = []
|
||||
for idx, img in enumerate(imgs):
|
||||
dt_boxes, rec_res, _ = self.__call__(img, cls)
|
||||
if not dt_boxes and not rec_res:
|
||||
ocr_res.append(None)
|
||||
continue
|
||||
tmp_res = [[box.tolist(), res]
|
||||
for box, res in zip(dt_boxes, rec_res)]
|
||||
ocr_res.append(tmp_res)
|
||||
|
@ -549,6 +552,9 @@ class PaddleOCR(predict_system.TextSystem):
|
|||
ocr_res = []
|
||||
for idx, img in enumerate(imgs):
|
||||
dt_boxes, elapse = self.text_detector(img)
|
||||
if not dt_boxes:
|
||||
ocr_res.append(None)
|
||||
continue
|
||||
tmp_res = [box.tolist() for box in dt_boxes]
|
||||
ocr_res.append(tmp_res)
|
||||
return ocr_res
|
||||
|
|
|
@ -65,15 +65,25 @@ class TextSystem(object):
|
|||
self.crop_image_res_index += bbox_num
|
||||
|
||||
def __call__(self, img, cls=True):
|
||||
time_dict = {'det': 0, 'rec': 0, 'csl': 0, 'all': 0}
|
||||
time_dict = {'det': 0, 'rec': 0, 'cls': 0, 'all': 0}
|
||||
|
||||
if img is None:
|
||||
logger.debug("no valid image provided")
|
||||
return None, None, time_dict
|
||||
|
||||
start = time.time()
|
||||
ori_im = img.copy()
|
||||
dt_boxes, elapse = self.text_detector(img)
|
||||
time_dict['det'] = elapse
|
||||
logger.debug("dt_boxes num : {}, elapse : {}".format(
|
||||
len(dt_boxes), elapse))
|
||||
|
||||
if dt_boxes is None:
|
||||
return None, None
|
||||
logger.debug("no dt_boxes found, elapsed : {}".format(elapse))
|
||||
end = time.time()
|
||||
time_dict['all'] = end - start
|
||||
return None, None, time_dict
|
||||
else:
|
||||
logger.debug("dt_boxes num : {}, elapsed : {}".format(
|
||||
len(dt_boxes), elapse))
|
||||
img_crop_list = []
|
||||
|
||||
dt_boxes = sorted_boxes(dt_boxes)
|
||||
|
@ -89,12 +99,12 @@ class TextSystem(object):
|
|||
img_crop_list, angle_list, elapse = self.text_classifier(
|
||||
img_crop_list)
|
||||
time_dict['cls'] = elapse
|
||||
logger.debug("cls num : {}, elapse : {}".format(
|
||||
logger.debug("cls num : {}, elapsed : {}".format(
|
||||
len(img_crop_list), elapse))
|
||||
|
||||
rec_res, elapse = self.text_recognizer(img_crop_list)
|
||||
time_dict['rec'] = elapse
|
||||
logger.debug("rec_res num : {}, elapse : {}".format(
|
||||
logger.debug("rec_res num : {}, elapsed : {}".format(
|
||||
len(rec_res), elapse))
|
||||
if self.args.save_crop_res:
|
||||
self.draw_crop_rec_res(self.args.crop_res_save_dir, img_crop_list,
|
||||
|
|
Loading…
Reference in New Issue