OCR_web.py

from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
import math
import numpy as np
import cv2
import re
import tempfile
from pathlib import Path

def image_split(imageBytes):
    # Read image from imageBytes
    image = cv2.imdecode(np.frombuffer(imageBytes, np.uint8), cv2.IMREAD_COLOR)
    (h, w) = image.shape[:2]

    sub_images = []
    # If height is greater than or equal to 2000
    if h >= 2000:
        # Split image into multiple images
        num = int(math.ceil(h / 2000))
        height = int(h / num)

        for i in range(num):
            # Calculate the dimensions of the sub-image
            startY = i * height
            endY = (i + 1) * height
            width = w
            dim = (width, endY - startY)
            # Double check that the dimensions are within bounds
            if dim[0] <= w and dim[1] <= h:
                # Resize the sub-image
                sub_image = cv2.resize(image[startY:endY, 0:w], dim, interpolation=cv2.INTER_AREA)
                # Save the sub-image
                # cv2.imwrite('test_{}.png'.format(i),sub_image)
                # Add the sub-image to the list
                sub_images.append(sub_image)
    else:
        sub_images.append(image)
    return sub_images

def image_resize(image):
    #image = cv2.imdecode(np.frombuffer(imageBytes, np.uint8), cv2.IMREAD_COLOR)
    (h,w) = image.shape[:2]
    WIDTH = 541
    # Calculate the ratio of height and width
    r = WIDTH / float(w)
    # Create a dimension with the calculated width and height
    dim = (WIDTH, int(h * r))
    # Resize the image
    image = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)

    return image

def msg_type(line_list, image, ori_width):
    resize_ratio = 541 / ori_width
    if len(line_list) == 2:
        # inner_lines = line_list[0]
        words = line_list[1][0]
        x = int(line_list[0][0][0] * resize_ratio)
        y_min = int(line_list[0][0][1] * resize_ratio)
        y_max = int(line_list[0][3][1] * resize_ratio)
        y_mid = y_min + int((y_max-y_min)/2) # 计算首个文字的正中心起始位置，确保后续计算对话框背景颜色更准确
        color = image[y_mid][x]
        # print("1111")

        return words, (x, y_mid), (y_min, y_max), color
    else:
        return None, (None, None), None

def date_delete(words, line_text, line, dialog, res):

    data_pattern1 = re.compile(r'^\d{1,2}:\d{1,2}$')  # 剔除24：24的时间戳
    if bool(data_pattern1.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern2 = re.compile(r'(凌晨|上午|下午|晚上)\d{1,2}:\d{1,2}$')  # 剔除凌晨|上午|下午|晚上24：24的时间戳
    if bool(data_pattern2.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern3 = re.compile(r'^(\d{4})年(\d{1,2})月(\d{1,2})日\d{1,2}:\d{1,2}$')  # 剔除2024年4月12日24：24的时间戳
    if bool(data_pattern3.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern4 = re.compile(r'^(\d{4})年(\d{1,2})月(\d{1,2})日(凌晨|上午|下午|晚上)\d{1,2}:\d{1,2}$')  # 剔除2024年4月12日凌晨|上午|下午|晚上24：24的时间戳
    if bool(data_pattern4.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern5 = re.compile(r'^昨天\d{1,2}:\d{1,2}$')  # 剔除昨天24：24的时间戳
    if bool(data_pattern5.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern6 = re.compile(r'昨天(凌晨|上午|下午|晚上)\d{1,2}:\d{1,2}$')  # 剔除昨天凌晨|上午|下午|晚上24：24的时间戳
    if bool(data_pattern6.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern7 = re.compile(r'^(\d{1,2})月(\d{1,2})日\d{1,2}:\d{1,2}$')  # 剔除4月12日24：24的时间戳
    if bool(data_pattern7.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern8 = re.compile(r'^(\d{1,2})月(\d{1,2})日(凌晨|上午|下午|晚上)\d{1,2}:\d{1,2}$')  # 剔除4月12日(凌晨|上午|下午|晚上)24：24的时间戳
    if bool(data_pattern8.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern9 = re.compile(r'(星期一|星期二|星期三|星期四|星期五|星期六|星期日)\d{1,2}:\d{1,2}$')  # 剔除星期一24：24的时间戳
    if bool(data_pattern9.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

    data_pattern10 = re.compile(r'(星期一|星期二|星期三|星期四|星期五|星期六|星期日)(凌晨|上午|下午|晚上)\d{1,2}:\d{1,2}$')  # 剔除星期一(凌晨|上午|下午|晚上)24：24的时间戳
    if bool(data_pattern10.match(words)):
        if line == res[-1]:
            dialog.append(line_text)
        return True

def OCR_text(img_path: str, output_path: str) -> Path:
    ocr = PaddleOCR(use_angle_cls=True, lang="ch", ocr_version="PP-OCRv4")

    with open(img_path, "rb") as f:
        imageBytes = f.read()
    image_list = image_split(imageBytes)

    result = ocr.ocr(img_path, cls=True)
    dialog = []
    line_text = ''
    # y_max_last = 0
    filename = ''

    # result0 = result[0]
    # image_RGB = Image.open(img_path).convert('RGB')
    # boxes = [line[0] for line in result0]
    # txts = [line[1][0] for line in result0]
    # scores = [line[1][1] for line in result0]
    # im_show = draw_ocr(image_RGB, boxes, txts, scores, font_path='./ppocr_img/fonts/simfang.ttf')
    # imshow = Image.fromarray(im_show)
    # imshow.save('result.jpg')

    output_txt_path = Path(output_path) / "result.txt"

    for image in image_list:
        ori_width = image.shape[1]

        image = image_resize(image)
        for idx in range(len(result)):
            res = result[idx]
            for line in res:
                # print(line)
                words, (x, y_mid), (y_min, y_max), color = msg_type(line, image, ori_width)

                if line[1][1] < 0.7:
                    if line == res[-1]:
                        dialog.append(line_text)
                    continue

                if words in ['...', '我通过了你的朋友验证请求，现在', '我们可以开始聊天了', '以上是打招呼的内容', '三',
                             '+', '你撤回了一条消息', 'HD', ':', '按住说话', '以下是新消息', '川']:
                    if line == res[-1]:
                        dialog.append(line_text)
                    continue

                if date_delete(words, line_text, line, dialog, res):
                    continue

                data_pattern = re.compile(r'^\d{1,3}"$')  # 剔除语音
                if bool(data_pattern.match(words)):
                    if line == res[-1]:
                        dialog.append(line_text)
                    continue

                ### 使用颜色判断（待进一步测试）
                if color[0] > 160 and color[1] > 160 and color[2] > 160:  # 背景颜色为白色，则判断为speaker1
                    if line_text == '':
                        line_text = 'Speaker1：' + words
                        y_max_last = y_max
                    elif (y_min - y_max_last) < (y_max - y_min):
                        line_text = line_text + words
                        y_max_last = y_max
                    else:
                        dialog.append(line_text)
                        line_text = 'Speaker1：' + words
                        y_max_last = y_max
                else:  # 否则判定为speaker2
                    if line_text == '':
                        line_text = 'Speaker2:' + words
                        y_max_last = y_max
                    elif (y_min - y_max_last) < (y_max - y_min):
                        line_text = line_text + words
                        y_max_last = y_max
                    else:
                        dialog.append(line_text)
                        line_text = 'Speaker2:' + words
                        y_max_last = y_max

                line_point = line
                test_point = res[-1]
                if line == res[-1]:
                    dialog.append(line_text)

            ### 保存识别到的结果为txt
            if filename == '':
                with open(output_txt_path, 'a+', encoding='utf-8') as f:
                    for dia in dialog:
                        f.write(dia + '\n')
            else:
                with open(filename.replace('/', '.') + '.txt', 'a+', encoding='utf-8') as f:
                    for dia in dialog:
                        f.write(dia + '\n')

    print(f'output dir is: {output_txt_path}')
    return output_txt_path





if __name__ == '__main__':
    img_path = './image/3.png'
    output_path = tempfile.mkdtemp()
    OCR_text(img_path, output_path)