blog

$ cd [poetryで初期化したプロジェクトディレクトリ]
$ poetry add pymupdf

# add-apt-repository ppa:alex-p/tesseract-ocr5
# apt update
# apt install tesseract-ocr

# apt install libtesseract-dev

$ cd [作業ディレクトリ]
$ git clone https://github.com/sirfz/tesserocr.git
$ cd tesserocr
$ CPPFLAGS=-march=native python3 setup.py build_ext -I/usr/local/include
$ cd [poetryで初期化したプロジェクトディレクトリ]
$ poetry add [作業ディレクトリ]/tesserocr

$ cd [poetryで初期化したプロジェクトディレクトリ]
$ poetry add opencv-python

$ cd [poetryで初期化したプロジェクトディレクトリ]
$ poetry add pandas

import matplotlib.pyplot as plt


def imshow(image):
    fig, ax = plt.subplots(figsize=(9, 9))
    ax.imshow(image, cmap=plt.cm.gray)
    plt.tight_layout()
    plt.show()

import io

import fitz
import pymupdf
from PIL import Image, ImageOps

pdf_file_name = "docs/pism/53-2-231.pdf"
page_index = 2

with fitz.open(pdf_file_name) as pdf_file:
    page = pdf_file[page_index]
    image_metadata = page.get_images()[0]
    xref = image_metadata[0]
    pdf_image = pdf_file.extract_image(xref)
    image_bytes = pdf_image["image"]
    image_ext = pdf_image["ext"]
    image_size = (pdf_image["width"], pdf_image["height"])

image = Image.open(io.BytesIO(image_bytes))
imshow(image)

import cv2
import numpy as np

image_array = np.frombuffer(image_bytes, np.uint8)
print(image_array)

[137  80  78 ...  66  96 130]

grayscale_image = cv2.imdecode(image_array, cv2.IMREAD_UNCHANGED)
imshow(grayscale_image)

print(grayscale_image)
print(grayscale_image[(grayscale_image > 0) & (grayscale_image < 255)])

[[0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 ...
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]
 [0 0 0 ... 0 0 0]]
[]

binary_image = cv2.threshold(grayscale_image, 20, 255, cv2.THRESH_BINARY)
print(binary_image)

(20.0, array([[0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       ...,
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0],
       [0, 0, 0, ..., 0, 0, 0]], dtype=uint8))

inverted_image = cv2.bitwise_not(grayscale_image)
imshow(inverted_image)

dilated_image = cv2.dilate(grayscale_image, None, iterations=2)
imshow(dilated_image)

contours, hierarchy = cv2.findContours(
    dilated_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)

contour_image = np.zeros_like(grayscale_image)
cv2.drawContours(contour_image, contours, -1, 255, 2)
imshow(contour_image)

rectangle_contours = []
for contour in contours:
    perimeter = cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, 0.02 * perimeter, True)
    if len(approx) == 4:
        rectangle_contours.append(approx)

rectangle_contour_image = np.zeros_like(grayscale_image)
cv2.drawContours(rectangle_contour_image, rectangle_contours, -1, 255, 1)
imshow(rectangle_contour_image)

max_area = 0
sorted_rectangle_contours = sorted(
    rectangle_contours, key=lambda x: cv2.contourArea(x)
)
print([cv2.contourArea(x) for x in sorted_rectangle_contours[-3:]])

top3_rectangle_contour_image = np.zeros_like(grayscale_image)
cv2.drawContours(
    top3_rectangle_contour_image, sorted_rectangle_contours[-3:], -1, 255, 1
)

imshow(top3_rectangle_contour_image)

[396750.0, 396864.0, 439901.0]

tables = []
for contour in sorted_rectangle_contours[-3:]:
    rect = cv2.boundingRect(contour)
    print(rect)
    table = grayscale_image[
        rect[1] : rect[1] + rect[3], rect[0] : rect[0] + rect[2]
    ]
    tables.append(table)
    imshow(table)

(0, 734, 1381, 289)

(0, 443, 1381, 289)

(0, 122, 1381, 320)

table_a = tables[2]
table_b = tables[1]
table_c = tables[0]

target_table_name = "a"
target_table = eval("table_{}".format(target_table_name))

def get_hline_image(table):
    kernel = np.ones([1, 6])
    image = cv2.erode(table, kernel, iterations=2)
    image = cv2.dilate(image, kernel, iterations=100)
    return image


hline_image = get_hline_image(target_table)
imshow(hline_image)

def get_hline_image(table):
    kernel = np.zeros((11, 11), np.uint8)
    kernel[5] = 1
    image = cv2.erode(table, kernel, iterations=2)
    image = cv2.dilate(image, kernel, iterations=60)
    return image


hline_image = get_hline_image(target_table)
imshow(hline_image)

def get_vline_image(table):
    kernel = np.zeros((11, 11), np.uint8)
    kernel[:, 5] = 1
    image = cv2.erode(table, kernel, iterations=2)
    image = cv2.dilate(image, kernel, iterations=60)
    return image


vline_image = get_vline_image(target_table)
imshow(vline_image)

def get_line_image(table):
    hline = get_hline_image(table)
    vline = get_vline_image(table)
    image = cv2.add(hline, vline)
    return image


line_image = get_line_image(target_table)
imshow(line_image)

def get_line_image(table):
    hline = get_hline_image(table)
    vline = get_vline_image(table)
    image = cv2.add(hline, vline)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    return cv2.dilate(image, kernel, iterations=5)


line_image = get_line_image(target_table)
imshow(line_image)

def remove_line(table, line):
    return cv2.subtract(table, line)


line_image = get_line_image(target_table)
image_without_line = remove_line(target_table, line_image)
imshow(image_without_line)

def get_mask(line):
    return cv2.bitwise_not(line)


line_image = get_line_image(target_table)
mask = get_mask(line_image)
imshow(mask)

def get_mask_contour(line):
    mask = get_mask(line)
    contour, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    return contour


line_image = get_line_image(target_table)
mask_contour = get_mask_contour(line_image)

mask_contour_image = np.zeros_like(target_table)
cv2.drawContours(mask_contour_image, mask_contour, -1, 255, 2)
imshow(mask_contour_image)

def get_word_bbox(line):
    mask_contour = get_mask_contour(line)
    bbox = []
    for contour in mask_contour:
        x, y, w, h = cv2.boundingRect(contour)
        bbox.append([x, y, w, h])
    return bbox


line_image = get_line_image(target_table)
word_bbox = get_word_bbox(line_image)

word_bbox_image = np.zeros_like(table_a)
for x, y, w, h in word_bbox:
    word_bbox_image = cv2.rectangle(
        word_bbox_image, (x, y), (x + w, y + h), 255, 2
    )
imshow(word_bbox_image)

def get_word_bbox(line):
    mask_contour = get_mask_contour(line)
    bbox = []
    for contour in mask_contour:
        x, y, w, h = cv2.boundingRect(contour)
        if y > 0 and w > 10 and h > 10:
            bbox.append([x, y, w, h])
    return bbox


line_image = get_line_image(target_table)
word_bbox = get_word_bbox(line_image)

def get_bbox_rows(line):

    word_bbox = get_word_bbox(line)
    word_bbox = list(reversed(word_bbox))

    bbox_hmean = np.mean([h for x, y, w, h in word_bbox if w > 10 and h > 10])
    bbox_hmean_div2 = bbox_hmean / 2

    rows = []
    cur_row = [word_bbox[0]]
    for bbox in word_bbox[1:]:
        cur_y = bbox[1]
        pre_y = cur_row[-1][1]
        dif_y = abs(cur_y - pre_y)
        if dif_y <= bbox_hmean_div2:
            cur_row.append(bbox)
        else:
            rows.append(sorted(cur_row, key=lambda x: x[0]))
            cur_row = [bbox]

    rows.append(sorted(cur_row, key=lambda x: x[0]))

    return rows

def get_cropped_bbox_rows(line):
    rows = get_bbox_rows(line)
    cropped_rows = np.zeros_like(rows)
    for ri, row in enumerate(rows):
        for ci, col in enumerate(row):
            word = table_a[col[1] : col[1] + col[3], col[0] : col[0] + col[2]]
            if np.max(word) == 0:
                cropped_rows[ri, ci] = [col[0], col[1], 1, 1]
            else:
                ys = np.any(word, axis=1)
                xs = np.any(word, axis=0)
                ymin, ymax = np.where(ys)[0][[0, -1]]
                xmin, xmax = np.where(xs)[0][[0, -1]]
                cropped_rows[ri, ci] = [
                    col[0] + xmin,
                    col[1] + ymin,
                    xmax - xmin + 1,
                    ymax - ymin + 1,
                ]
    return cropped_rows

line_image = get_line_image(target_table)
bbox_rows = get_bbox_rows(line_image)

fig, axes = plt.subplots(len(bbox_rows), len(bbox_rows[0]), figsize=(20, 6))

for ri, row in enumerate(bbox_rows):
    for ci, col in enumerate(row):
        ax = axes[ri, ci]
        ax.set_axis_off()
        word = target_table[col[1] : col[1] + col[3], col[0] : col[0] + col[2]]
        padding = np.full(4, 5)
        padded_word = cv2.copyMakeBorder(
            word,
            *padding,
            cv2.BORDER_CONSTANT,
            (0, 0, 0),
        )
        ax.imshow(padded_word, cmap=plt.cm.gray)
        cv2.imwrite(
            "data/pism/image/table-1-{}-{:02d}{:02d}.{}".format(
                target_table_name, ri, ci, image_ext
            ),
            padded_word,
        )

import tesserocr
from tesserocr import OEM, PSM, PyTessBaseAPI

line_image = get_line_image(target_table)
bbox_rows = get_bbox_rows(line_image)

fig, axes = plt.subplots(len(bbox_rows), len(bbox_rows[0]), figsize=(20, 6))

data = []

with PyTessBaseAPI(
    path="/usr/share/tesseract-ocr/5/tessdata",
    oem=OEM.LSTM_ONLY,
    psm=PSM.SINGLE_LINE,
) as api:

    api.SetVariable("tessedit_char_whitelist", "nN0123456789-——+/")

    for ri, row in enumerate(bbox_rows):
        data_row = []
        for ci, col in enumerate(row):

            ax = axes[ri, ci]
            ax.set_axis_off()

            word = target_table[
                col[1] : col[1] + col[3], col[0] : col[0] + col[2]
            ]

            padding = np.full(4, 5)
            padded_word = cv2.copyMakeBorder(
                word,
                *padding,
                cv2.BORDER_CONSTANT,
                (0, 0, 0),
            )

            if np.any(word):

                api.SetImage(Image.fromarray(padded_word))
                text = api.GetUTF8Text().strip()

                if not text:
                    num_dilate = cv2.dilate(padded_word, None, iterations=1)
                    api.SetImage(Image.fromarray(num_dilate))
                    text = api.GetUTF8Text().strip()

                if not text:
                    kernel = np.zeros((5, 5), np.uint8)
                    kernel[2] = 1
                    bar_dilate = cv2.dilate(padded_word, kernel, iterations=3)
                    api.SetImage(Image.fromarray(bar_dilate))
                    text = api.GetUTF8Text().strip()
            else:
                text = ""

            if text and text in "-——":
                text = "-"

            data_row.append(text)

            ax.set_title(text, fontsize=8)
            ax.imshow(padded_word, cmap=plt.cm.gray)

        data.append(data_row)

plt.tight_layout()
plt.show()

import pandas as pd

pd.options.display.max_columns = 40
pd.options.display.max_rows = 100
df = pd.DataFrame(data).drop(0, axis=1)
df

import tesserocr
from tesserocr import OEM, PSM, PyTessBaseAPI


def parse_bbox(api, bbox):

    if not np.any(bbox):
        return ""

    padding = np.full(4, 5)
    bbox = cv2.copyMakeBorder(
        bbox,
        *padding,
        cv2.BORDER_CONSTANT,
        (0, 0, 0),
    )

    api.SetVariable("tessedit_char_whitelist", "+/")
    api.SetImage(Image.fromarray(bbox))
    text = api.GetUTF8Text().strip()

    if not text:
        tmp = cv2.dilate(bbox, None, iterations=1)
        api.SetImage(Image.fromarray(tmp))
        text = api.GetUTF8Text().strip()

    if not text:
        api.SetVariable("tessedit_char_whitelist", "0123456789")
        api.SetImage(Image.fromarray(bbox))
        text = api.GetUTF8Text().strip()

        if not text:
            kernel = np.zeros((3, 3), np.uint8)
            kernel[1] = 1
            kernel[:, 1] = 1
            tmp = cv2.dilate(bbox, kernel, iterations=1)
            api.SetImage(Image.fromarray(tmp))
            text = api.GetUTF8Text().strip()

        if not text:
            tmp = cv2.morphologyEx(bbox, cv2.MORPH_CLOSE, None)
            api.SetImage(Image.fromarray(tmp))
            text = api.GetUTF8Text().strip()

    if not text:
        api.SetVariable("tessedit_char_whitelist", "N")
        kernel = np.zeros((5, 5), np.uint8)
        kernel[:, 2] = 1
        tmp = cv2.dilate(bbox, kernel, iterations=1)
        api.SetImage(Image.fromarray(tmp))
        text = api.GetUTF8Text().strip()

    if not text:
        api.SetVariable("tessedit_char_whitelist", "n")
        kernel = np.zeros((3, 3), np.uint8)
        kernel[:, 1] = 1
        kernel[1] = 1
        tmp = cv2.dilate(bbox, kernel=kernel, iterations=1)
        tmp = cv2.erode(tmp, None, iterations=1)
        api.SetImage(Image.fromarray(tmp))
        text = api.GetUTF8Text().strip()

    if not text:
        api.SetVariable("tessedit_char_whitelist", "-——")
        kernel = np.zeros((5, 5), np.uint8)
        kernel[2] = 1
        tmp = cv2.erode(bbox, kernel, iterations=2)
        tmp = cv2.dilate(bbox, kernel, iterations=4)
        api.SetImage(Image.fromarray(tmp))
        text = api.GetUTF8Text().strip()

    if not text:
        return ""

    if text in "-——":
        text = "-"

    return text


def parse_table_image(table, debug=False):

    line_image = get_line_image(table)
    bbox_rows = get_bbox_rows(line_image)

    if debug:
        fig, axes = plt.subplots(
            len(bbox_rows), len(bbox_rows[0]), figsize=(20, 6)
        )

    data = []

    api = PyTessBaseAPI(
        path="/usr/share/tesseract-ocr/5/tessdata",
        oem=OEM.LSTM_ONLY,
        psm=PSM.SINGLE_LINE,
    )

    for ri, row in enumerate(bbox_rows):
        data_row = []
        for ci, col in enumerate(row):
            bbox = table[col[1] : col[1] + col[3], col[0] : col[0] + col[2]]
            text = parse_bbox(api, bbox)
            data_row.append(text)
            if debug:
                ax = axes[ri, ci]
                ax.set_axis_off()
                ax.set_title(text, fontsize=13)
                ax.imshow(bbox, cmap=plt.cm.gray)
        data.append(data_row)

    if debug:
        plt.tight_layout()
        plt.show()
    return data


shift_table = pd.DataFrame()
for team, table in zip(["c", "b", "a"], tables):
    print("チーム", team.upper())
    data = parse_table_image(table, True)
    df = pd.DataFrame(data).drop(0, axis=1)
    shift_table = pd.concat([shift_table, df], ignore_index=True)
    df.to_csv("data/pism/53-2-231-table-1-{}.csv".format(team), index=False)
    display(df)

チーム C

チーム B

チーム A

shift_table[[1, 32, 33, 34, 35]] = shift_table[[1, 32, 33, 34, 35]].astype(int)
shift_table = shift_table.sort_values(by=1, ignore_index=True)
shift_table

answer = pd.read_csv(
    "data/pism/53-2-231-table-1.csv",
    header=None,
    skiprows=1,
    names=[i for i in range(1, 36)],
)
answer

answer.compare(shift_table)