未验证 提交 b131c31a 编写于 作者: D dyning 提交者: GitHub

Merge pull request #1035 from LDOUBLEV/fixocr

add comment
...@@ -91,12 +91,22 @@ def parse_args(): ...@@ -91,12 +91,22 @@ def parse_args():
def create_predictor(args, mode): def create_predictor(args, mode):
"""
create predictor for inference
:param args: params for prediction engine
:param mode: mode
:return: predictor
"""
if mode == "det": if mode == "det":
model_dir = args.det_model_dir model_dir = args.det_model_dir
elif mode == 'cls': elif mode == 'cls':
model_dir = args.cls_model_dir model_dir = args.cls_model_dir
else: elif mode == 'rec':
model_dir = args.rec_model_dir model_dir = args.rec_model_dir
else:
raise ValueError(
"'mode' of create_predictor() can only be one of ['det', 'cls', 'rec']"
)
if model_dir is None: if model_dir is None:
logger.info("not find {} model file path {}".format(mode, model_dir)) logger.info("not find {} model file path {}".format(mode, model_dir))
...@@ -144,6 +154,12 @@ def create_predictor(args, mode): ...@@ -144,6 +154,12 @@ def create_predictor(args, mode):
def draw_text_det_res(dt_boxes, img_path): def draw_text_det_res(dt_boxes, img_path):
"""
Visualize the results of detection
:param dt_boxes: The boxes predicted by detection model
:param img_path: Image path
:return: Visualized image
"""
src_im = cv2.imread(img_path) src_im = cv2.imread(img_path)
for box in dt_boxes: for box in dt_boxes:
box = np.array(box).astype(np.int32).reshape(-1, 2) box = np.array(box).astype(np.int32).reshape(-1, 2)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册