OCRの読み取りミスを修正する手作業をAIに代わりにやってもらってみた

import codecs
import csv
import json
import os
import time
from typing import Dict, List

import cv2
import numpy as np
import openai
import requests

WIDTH_CONSOLIDATE = 3
SPLITTER = ","
# OPENAI_MODEL = "gpt-4-0613"
OPENAI_MODEL = "gpt-3.5-turbo-0613"

prompt_template = """
OCRで複数のセルの文字列を読み込んだのですが、誤って1つのセルとして読み込んでしまいました。前後の行のように"{splitter}"で分割してください。分割した結果のみ出力してください。

# 前後の行
{near_rows}

# 分割したい文字列
{target_row}

# 分割結果
"""

# 矩形描画
def disp_rects(rects, img, thickness):
    image = img.copy()
    for i, rect in enumerate(rects):
        color = np.random.randint(0, 255, 3).tolist()
        cv2.drawContours(image, rects, i, color, thickness)
    return image

# 近似座標の集約
def consolidation(list: List[int]):
    result: List[int] = []
    min = 0
    counter = 0
    for val in list:
        if min == 0:
            min = val
            keep = val
        else:
            if val > keep + WIDTH_CONSOLIDATE:  # 3ドット以内をまとめる
                if counter > 2:  # 得意な検出は排除する
                    result.append(int(min + (keep - min) / 2))
                min = val
                counter = 0
            counter += 1
            keep = val

    if counter > 2:  # 特異な検出は排除する
        result.append(int(min + (keep - min) / 2))

    return result

# 座標検出
def detect_point(rects: List[np.ndarray]):
    # 全X,Y検出
    x_list = []
    y_list = []
    for rect in rects:
        for i in range(4):
            x, y = rect[i]
            if not x in x_list:
                x_list.append(x)
            if not y in y_list:
                y_list.append(y)

    x_list.sort()
    y_list.sort()

    # 近似値の集約
    x_list = consolidation(x_list)
    y_list = consolidation(y_list)

    return x_list, y_list

# LINE描画
def disp_line(x_list, y_list, img):
    image = img.copy()

    x_min = min(x_list)
    x_max = max(x_list)
    y_min = min(y_list)
    y_max = max(y_list)

    for x in x_list:
        cv2.line(
            image,
            pt1=(x, y_min),
            pt2=(x, y_max),
            color=(0, 0, 255),
            thickness=1,
            lineType=cv2.LINE_4,
            shift=0,
        )

    for y in y_list:
        cv2.line(
            image,
            pt1=(x_min, y),
            pt2=(x_max, y),
            color=(0, 0, 255),
            thickness=1,
            lineType=cv2.LINE_4,
            shift=0,
        )

    return image

def estimate_table_coords(org_img: np.ndarray):
    # グレースケール変換
    gray_image = cv2.cvtColor(org_img, cv2.COLOR_BGR2GRAY)
    # エッジ抽出
    edges_image = cv2.Canny(gray_image, 1, 100, apertureSize=3)
    # 膨張処理
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    dilate_image = cv2.dilate(edges_image, kernel)

    # 輪郭抽出
    contours, hierarchy = cv2.findContours(dilate_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # 面積でフィルタリング
    rects: List[np.ndarray] = []
    for cnt, hrchy in zip(contours, hierarchy[0]):
        if cv2.contourArea(cnt) < 3000:
            continue  # 面積が一定の大きさを満たさないものを除く
        if cv2.contourArea(cnt) > 20000:
            continue  # 面積が一定の大きさを超えるものを除く
        if hrchy[3] == -1:
            continue  # ルートノードは除く
        # 輪郭を囲む長方形を計算する。
        rect = cv2.minAreaRect(cnt)
        rect_points = cv2.boxPoints(rect).astype(int)
        rects.append(rect_points)

    # 座標検出
    x_list, y_list = detect_point(rects)

    return x_list, y_list

def detect_empty_cell_sets(line: List[str]):
    index_sets_empty_cell: List[List[int]] = []
    index_set = []
    for i, cell in enumerate(line):
        if cell == "":
            index_set.append(i)
            if i == (len(line) - 1) or line[i + 1] != "":
                index_sets_empty_cell.append(index_set)
                index_set = []

    return index_sets_empty_cell

def search_bbox(
    x_list: List[int],
    y_list: List[int],
    i_line: int,
    read_results: Dict,
    index_set: List[int],
):
    # coordinate from image
    left = x_list[index_set[0]]
    right = x_list[index_set[-1] + 1]
    top = y_list[i_line]
    bottom = y_list[i_line + 1]

    for line in read_results["lines"]:
        # coordinate from OCR
        boundingBox = line["boundingBox"]
        x0 = boundingBox[0]
        x1 = boundingBox[4]
        y0 = boundingBox[1]
        y1 = boundingBox[5]

        # if bbox in empty cell
        if left <= x0 and x1 <= right and top <= y0 and y1 <= bottom:
            return line

    return None

def extract_near_cell_values(
    content: List[List[str]],
    i_line: int,
    index_set: List[int],
    n_get_near_rows: int = 3,
):
    near_cells_list: List[List[str]] = []
    n_get_line = 0
    for line in content[1:]:  # 先頭行はヘッダーなので除く
        near_cells = line[index_set[0] : index_set[-1] + 1]
        if all([cell != "" for cell in near_cells]):
            near_cells_list.append(near_cells)
            n_get_line += 1
            if n_get_line >= n_get_near_rows:
                break
    return near_cells_list

def estimate_split_with_AI(
    ocr_text: str,
    near_cells_list: List[List[str]],
):
    prompt = prompt_template.format(
        splitter=SPLITTER,
        near_rows="\n".join([SPLITTER.join(line) for line in near_cells_list]),
        target_row=ocr_text,
    )
    messages = [{"role": "user", "content": prompt}]

    response = openai.ChatCompletion.create(
        model=OPENAI_MODEL,
        temperature=0.0,
        messages=messages,
    )
    content: str = response["choices"][0]["message"]["content"]  # type: ignore
    estimated_result = content.split(SPLITTER)
    print(content)
    print(estimated_result)

    return estimated_result

def revise_table(
    x_list: List[int],
    y_list: List[int],
    response: Dict,
    content: List[List[str]],
):
    content_ = content.copy()
    read_results = response["analyzeResult"]["readResults"][0]

    # detect empty cell
    for i_line, content_line in enumerate(content):
        index_sets_empty_cell = detect_empty_cell_sets(content_line)
        index_sets_empty_cell_filtered = [
            index_set for index_set in index_sets_empty_cell if all([index < len(x_list) for index in index_set])
        ]
        for index_set_empty in index_sets_empty_cell_filtered:
            if len(index_set_empty) == 1:  # 1つのセルが空の場合、今回の推定の対象外なので何もしない
                print(f"WARNING: {i_line} {index_set_empty} cannot be estimated")
                pass
            else:  # 2つ以上のセルが連続で空の場合、その座標に該当するOCRの認識結果をあてはめる
                # 該当するOCRの認識結果を探す
                line_ocr = search_bbox(x_list, y_list, i_line, read_results, index_set_empty)
                if line_ocr is None:  # その箇所をOCRが認識していなければ処理しない（できない）
                    continue

                near_cells_list = extract_near_cell_values(content, i_line, index_set_empty)

                ocr_text = line_ocr["text"]
                estimated_cells = estimate_split_with_AI(ocr_text, near_cells_list)
                if len(index_set_empty) != len(estimated_cells):
                    print(f"WARNING: {i_line} {index_set_empty} ignored")

                for index, cell in zip(index_set_empty, estimated_cells):
                    content_[i_line][index] = cell

    return content_

def main(
    filepath_img: str,
    filepath_ocr_json: str,
    filepath_csv_input: str,
    filepath_csv_output: str,
):
    # 入力
    img_fax_warped = cv2.imread(filepath_img)

    with codecs.open(filepath_ocr_json, "r", encoding="utf8") as f:
        response = json.load(f)

    with open(filepath_csv_input, "r", encoding="utf8") as f:
        reader = csv.reader(f)
        content = [row for row in reader]

    # 画像に写っている表の格子の座標を推定する
    x_list, y_list = estimate_table_coords(img_fax_warped)

    # 空のセルを推定する
    content_revised = revise_table(x_list, y_list, response, content)

    # 出力

    with open(filepath_csv_output, mode="w", encoding="utf_8_sig") as f:
        for line in content_revised:
            f.write(",".join(line))
            f.write("\n")

if __name__ == "__main__":
    filepath_img = os.path.join(os.path.dirname(__file__), "fax.png")
    filepath_ocr_json = os.path.join(os.path.dirname(__file__), "..", "output_read3.2.json")
    filepath_csv_input = os.path.join(os.path.dirname(__file__), "..", "output.csv")
    filepath_csv_output = os.path.join(os.path.dirname(__file__), "..", "output_rev.csv")
    main(
        filepath_img=filepath_img,
        filepath_ocr_json=filepath_ocr_json,
        filepath_csv_input=filepath_csv_input,
        filepath_csv_output=filepath_csv_output,
    )