paddleOCR代码学习-详细注释

可视化-代码

def draw_ocr_box_txt(image,
                     boxes,
                     txts,
                     scores=None,
                     drop_score=0.5,
                     font_path="./doc/simfang.ttf"):
    """
    把结果画在图片上
    image,  图片
    boxes,  预测的点坐标 [[(171, 1308), (1440, 1462), (1438, 1692), (162, 1539)], ...]
    txts,   文本        ["888", ...]
    scores = None,  分数
    drop_score = 0.5, 丢弃阈值
    font_path = "./doc/simfang.ttf"  字体
    """
    h, w = image.height, image.width
    img_left = image.copy()
    # PIL.Image.new() 使用给定的模式和大小创建一个新图像
    # 新建一张白色的图
    img_right = Image.new('RGB', (w, h), (255, 255, 255))
    # ImageDraw模块提供了图像对象的简单2D绘制。用户可以使用这个模块创建新的图像，注释或润饰已存在图像，为web应用实时产生各种图形。 Draw()创建一个可以在给定图像上绘图的对象。
    draw_left = ImageDraw.Draw(img_left)
    draw_right = ImageDraw.Draw(img_right)
    for idx, (box, txt) in enumerate(zip(boxes, txts)):
        if scores is not None and scores[idx] < drop_score:
            continue
        # rgb 颜色设置 这里随机生成RGB颜色
        color = (random.randint(0, 255), random.randint(0, 255),
                 random.randint(0, 255))
        # 根据点坐标画多边形 fill为填充颜色
        draw_left.polygon(box, fill=color)
        # 根据点坐标画多边形，outline为边框颜色
        draw_right.polygon(
            [
                box[0][0], box[0][1], box[1][0], box[1][1], box[2][0],
                box[2][1], box[3][0], box[3][1]
            ],
            outline=color)
        # 高度计算 sqrt 开平方根 四个点的顺序是 左上点 右上点，右下点，左下点，宽度计算为使用前两个点坐标，利用勾股定理计算
        box_height = math.sqrt((box[0][0] - box[3][0])**2 + (box[0][1] - box[3][
            1])**2)
        # 宽度计算，开平方根
        box_width = math.sqrt((box[0][0] - box[1][0])**2 + (box[0][1] - box[1][
            1])**2)
        # 判断是竖着的还是横着的框 height大于2倍的width，认为框是竖着的
        if box_height > 2 * box_width:
            # 字体大小
            font_size = max(int(box_width * 0.9), 10)
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
            # 初始纵坐标
            cur_y = box[0][1]
            # 把字符分开 竖起来写书写
            for c in txt:
                char_size = font.getsize(c) # char_size,字符的尺寸，2个维度（width, height）
                draw_right.text(
                    (box[0][0] + 3, cur_y), c, fill=(0, 0, 0), font=font)
                cur_y += char_size[1]
        # 框横着
        else:
            font_size = max(int(box_height * 0.8), 10)
            font = ImageFont.truetype(font_path, font_size, encoding="utf-8")
            draw_right.text(
                [box[0][0], box[0][1]], txt, fill=(0, 0, 0), font=font)
    # Image.blend 图像合成，两张图合成一张，0.5为不透明度，值越大，越不透明
    img = Image.blend(image, img_left, 0.5)
    # 新建一张图，设置RGB色彩(255, 255, 255)，设置宽度和高度，这里设置2倍的宽是为了粘贴2张图像
    img_show = Image.new('RGB', (w * 2, h), (255, 255, 255))
    # 把图像粘贴在img_show上
    img_show.paste(img, (0, 0, w, h))
    img_show.paste(img_right, (w, 0, w * 2, h))
    return np.array(img_show) # 转化为array格式后，opencv可以读取，也可以进入神经网络计算

predict_system.py

class TextSystem(object):
    def __init__(self, args):
        self.text_detector = predict_det.TextDetector(args)
        self.text_recognizer = predict_rec.TextRecognizer(args)
        self.use_angle_cls = args.use_angle_cls
        self.drop_score = args.drop_score
        if self.use_angle_cls:
            self.text_classifier = predict_cls.TextClassifier(args)

    def get_rotate_crop_image(self, img, points):
        '''
        剪切--做透视变换 和 旋转

        img_height, img_width = img.shape[0:2]
        left = int(np.min(points[:, 0]))
        right = int(np.max(points[:, 0]))
        top = int(np.min(points[:, 1]))
        bottom = int(np.max(points[:, 1]))
        img_crop = img[top:bottom, left:right, :].copy()
        points[:, 0] = points[:, 0] - left
        points[:, 1] = points[:, 1] - top
        '''
        img_crop_width = int(
            max(
                np.linalg.norm(points[0] - points[1]),
                np.linalg.norm(points[2] - points[3])))
        img_crop_height = int(
            max(
                np.linalg.norm(points[0] - points[3]),
                np.linalg.norm(points[1] - points[2])))
        pts_std = np.float32([[0, 0], [img_crop_width, 0],
                              [img_crop_width, img_crop_height],
                              [0, img_crop_height]])
        M = cv2.getPerspectiveTransform(points, pts_std)
        # cv2.warpPerspective()透视变换函数，可保持直线不变形，但是平行线可能不再平行
        dst_img = cv2.warpPerspective(
            img,
            M, (img_crop_width, img_crop_height),
            borderMode=cv2.BORDER_REPLICATE,
            flags=cv2.INTER_CUBIC)
        dst_img_height, dst_img_width = dst_img.shape[0:2]
        if dst_img_height * 1.0 / dst_img_width >= 1.5:
            dst_img = np.rot90(dst_img)
        return dst_img

    def print_draw_crop_rec_res(self, img_crop_list, rec_res):
        bbox_num = len(img_crop_list)
        for bno in range(bbox_num):
            cv2.imwrite("./output/img_crop_%d.jpg" % bno, img_crop_list[bno])
            logger.info(bno, rec_res[bno])

    def __call__(self, img):
        ori_im = img.copy()
        # 区域预测，返回预测区域四点和使用时间
        dt_boxes, elapse = self.text_detector(img)
        logger.info("dt_boxes num : {}, elapse : {}".format(
            len(dt_boxes), elapse))
        if dt_boxes is None:
            return None, None
        img_crop_list = []

        # 按从上到下、从左到右的顺序对文本框排序
        dt_boxes = sorted_boxes(dt_boxes)

        # 把预测的位置框裁剪出来一一做仿射变换，数据加入到列表中
        for bno in range(len(dt_boxes)):
            tmp_box = copy.deepcopy(dt_boxes[bno])
            # 裁剪--做透视变换和旋转
            img_crop = self.get_rotate_crop_image(ori_im, tmp_box)
            img_crop_list.append(img_crop)
        # 是否使用方向旋转
        if self.use_angle_cls:
            img_crop_list, angle_list, elapse = self.text_classifier(
                img_crop_list)
            logger.info("cls num  : {}, elapse : {}".format(
                len(img_crop_list), elapse))

        # 文本识别 返回结果和耗时
        rec_res, elapse = self.text_recognizer(img_crop_list)
        logger.info("rec_res num  : {}, elapse : {}".format(
            len(rec_res), elapse))
        # self.print_draw_crop_rec_res(img_crop_list, rec_res)
        filter_boxes, filter_rec_res = [], []
        # 这里根据 置信度 和 丢弃值 判断识别结果是否要丢弃，丢弃置信度低的结果，返回置信度高于丢弃值的结果
        for box, rec_reuslt in zip(dt_boxes, rec_res):
            text, score = rec_reuslt
            if score >= self.drop