【Python】バウンディングボックスから赤枠表示

初めに

pythonのopencvを用いてバウンディングボックスから赤枠表示をしていきます

用意するもの

• 赤枠表示をするpdfファイル
• csvファイル（必要な列は以下の5つです）
  1. page: pdfの何ページかを表す
  2. x0: 左上のx座標
  3. y0: 左上のy座標
  4. y0: 右下のx座標
  5. y0: 右下のy座標

書いていきます！

流れはpdfをpngに変換して、そこにcsvの情報からopencvを用いて赤枠表示をしています！
以下コード例です

import cv2
import os
import pathlib
import pandas as pd
from pdf2image import convert_from_path

"""
csvに必要な列['page', 'x0', 'y0', 'x1', 'y1']
"""

def pdf2png(pdf_path, png_DIR):
    convert_from_path(pdf_path, output_folder=png_DIR, fmt='png', output_file=pdf_path.stem)

def rectangle(x0, y0, x1, y1, image):
    if x0 <= 1: # バウンディングボックスが0~1の割合の時
        height, width, _ = image.shape
        cv2.rectangle(image, (int(x0*width), int(y0*height)), (int(x1*width), int(y1*height)), (0, 0, 255), 1)
        # 赤枠の左上に番号を表示されたい時
        # cv2.putText(image, str(i), (int(x0*width), int(y0*height)), font, 1, (0, 0, 255), 1, cv2.LINE_AA)
    else: # バウンディングボックスの値がピクセルの時
        cv2.rectangle(image, (int(x0), int(y0)), (int(x1), int(y1)), (0, 0, 255), 1)
        # 同上
        # cv2.putText(image, str(i), (int(x0), int(y0)), font, 1, (0, 0, 255), 1, cv2.LINE_AA)

def rect(csv_path, png_DIR, out_DIR):
    image_filenames = [filename for filename in os.listdir(png_DIR)]
    image_filenames.sort()
    image_list = []
    for filename in image_filenames:
        image = cv2.imread(os.path.join(png_DIR, filename))
        image_list.append(image)

    df = pd.read_csv(csv_path, encoding="ISO-8859-1")
    page = df["page"]
    x0 = df["x0"]
    y0 = df["y0"]
    x1 = df["x1"]
    y1 = df["y1"]

    page_list = page.to_list()
    x0_list = x0.to_list()
    y0_list = y0.to_list()
    x1_list = x1.to_list()
    y1_list = y1.to_list()
    font = cv2.FONT_HERSHEY_SIMPLEX

    for i in range(len(x0_list)):
        page = page_list[i]

        x0 = x0_list[i]
        y0 = y0_list[i]
        x1 = x1_list[i]
        y1 = y1_list[i]

        if i == 0:
            image = image_list[page-1]
            rectangle(x0, y0, x1, y1, image)
        elif i+1 == len(x0_list):
            rectangle(x0, y0, x1, y1, image)
            cv2.imwrite(os.path.join(out_DIR, "{}.png".format(page)), image)
        elif page == page_list[i-1]:
            rectangle(x0, y0, x1, y1, image)
        else:
            cv2.imwrite(os.path.join(out_DIR, "{}.png".format(page-1)), image)
            image = image_list[page-1]
            rectangle(x0, y0, x1, y1, image)

def main(pdf_path, csv_path):
    pdf_path = pathlib.Path(pdf_path)
    csv_path = pathlib.Path(csv_path)
    pdf_DIR = os.path.dirname(pdf_path)
    # pdfをpngとするDIR
    png_DIR = os.path.join(pdf_DIR, 'png')
    if not os.path.isdir(png_DIR):
        os.mkdir(png_DIR)
    # 赤枠表示のpngを保存するDIR
    out_DIR = os.path.join(pdf_DIR, 'rect')
    if not os.path.isdir(out_DIR):
        os.mkdir(out_DIR)

    pdf2png(pdf_path, png_DIR)
    rect(csv_path, png_DIR, out_DIR)


if __name__ == "__main__":
    # 入力のpdf
    pdf_path = "path/pdf"
    # バウンディングボックスの情報が載っているcsvファイル
    csv_path = "path/csv"

    main(pdf_path, csv_path)

”png”というディレクトリにpdfを分割し変換したpngファイルを、
"rect"というディレクトリに赤枠表示をしたpngファイルが保存されます！

以下このファイルを実行した出力例です！