[YOLO] Yolov8 이미지에서 표 영역 탐지하는 법

문서 내에는 텍스트, 이미지, 표 등 다양한 형식의 요소가 포함되어 있습니다. 특히, 표는 정보를 구조적으로 정리하여 전달하는 중요한 수단 중 하나입니다. 표의 영역을 정확하게 탐지하고 추출하는 것은 문서 처리 및 정보 추출에 있어서 핵심적인 작업입니다. 이를 위해 최신 객체 탐지 기술 중 하나의 YOLOv8 모델을 활용하여 표의 위치를 결정하고 해당 영역을 정확하게 추출하는 과정이 필요합니다. 이 글에서는 YOLOv8을 활용하여 문서 내의 표를 식별하고 추출하는 방법에 대해 살펴봅니다.

표 영역 탐지를 위한 YOLOv8 활용

표의 영역을 탐지하는 것은 문서 내에서 텍스트, 이미지, 표 등 다양한 형식의 요소를 인식하고 그 중에서 표를 식별하여 표의 행과열, 셀 등의 구조를 추출하는 것입니다. 이를 위해 YOLOv8 객체 탐지 모델을 사용하여 표를 감지하고, 해당 표를 분할하여 표 영역을 정확하게 추출합니다.

• YOLOv8을 사용하여 문서 내의 테이블의 위치를 결정합니다. 미리 학습된 모델을 사용하여 표를 포함하는 영역의 경계 상자를 예측합니다.
• 표의 위치를 결정한 후, 해당 영역을 분할하여 표 구간만을 추출합니다. 이를 통해 이미지에서 표가 있는 부분만을 따로 분리하여 추출합니다.

1. 라이브러리 설치

먼저, 필요한 라이브러리를 설치합니다.

pip install transformers ultralyticsplus==0.0.23 ultralytics==8.0.21
pip install pdf2image # pdf를 이미지로 변환하는 라이브러리

 

pdf2image 라이브러리를 사용하여 PDF 파일을 불러올 때 오류 메시지가 뜨는 경우가 있습니다.

해당 오류메시지는 Poppler를 설치하면 해결됩니다.

conda install -c conda-forge poppler

 

2. Python Code

YOLOv8 모델을 사용하여 표 영역을 탐지한 후 표 영역만을 추출하는 python code 입니다.

 

필요한 라이브러리를 import 합니다.

import numpy as np
from ultralyticsplus import YOLO, render_result
from PIL import Image
import pdf2image

 

논문과 같은 문서는 pdf 형식인 경우가 많습니다.

따라서 pdf를 이미지로 변환하는 pdf2image 라이브러리를 사용하여 이미지를 불러옵니다.

pdf_file = "C:/Users/JIYUN/Project/etc/test_data.pdf"
img = np.asarray(pdf2image.convert_from_path(pdf_file)[3]) # [3]은 pdf의 페이지수를 의미
pil_img  = Image.fromarray(img)

# 이미지 확인
pil_img.save('./tmp.png')

 

YOLOv8m 모델을 로드하고 모델의 매개변수를 설정합니다.

# load model
model = YOLO('keremberke/yolov8m-table-extraction')

# set model parameters
model.overrides['conf'] = 0.25  # NMS confidence threshold
model.overrides['iou'] = 0.45  # NMS IoU threshold
model.overrides['agnostic_nms'] = False  # NMS class-agnostic
model.overrides['max_det'] = 1000  # maximum number of detections per image

 

YOLOv8m 모델을 사용하여 이미지 내에서 표의 영역을 탐지합니다.

results = model.predict(pil_img)

print('Boxes: ', results[0].boxes)
render = render_result(model=model, image=img, result=results[0])
render.save('./tmp_render.png') # 이미지 확인

 

 

논문 내 표 영역 탐지

 

해당 영역만을 이미지에서 크롭합니다.

위 이미지에서 표시된 영역의 표만 저장됩니다.

# crop
x1, y1, x2, y2, _, _ = tuple(int(item) for item in results[0].boxes.data.numpy()[0])
cropped_image = img[y1:y2, x1:x2]
cropped_image = Image.fromarray(cropped_image)

cropped_image.save('./tmp_crop.png')

 

 

3. 전체 코드

위에서 작성된 전체 코드입니다.

import numpy as np
from ultralyticsplus import YOLO, render_result
from PIL import Image
import pdf2image

pdf_file = "C:/Users/JIYUN/Project/etc/test_data.pdf"
img = np.asarray(pdf2image.convert_from_path(pdf_file)[3])
pil_img  = Image.fromarray(img)

pil_img.save('./tmp.png')


# YOLOv8m model
model = YOLO('keremberke/yolov8m-table-extraction')

# set model parameters
model.overrides['conf'] = 0.25  # NMS confidence threshold
model.overrides['iou'] = 0.45  # NMS IoU threshold
model.overrides['agnostic_nms'] = False  # NMS class-agnostic
model.overrides['max_det'] = 1000  # maximum number of detections per image

# table detection
results = model.predict(pil_img)
print('Boxes: ', results[0].boxes)
render = render_result(model=model, image=img, result=results[0])
render.save('./tmp_render.png')

# crop
x1, y1, x2, y2, _, _ = tuple(int(item) for item in results[0].boxes.data.numpy()[0])
cropped_image = img[y1:y2, x1:x2]
cropped_image = Image.fromarray(cropped_image)

cropped_image.save('./tmp_crop.png')

 

 

결론적으로, 이 글에서는 YOLOv8 객체 탐지 모델을 활용하여 문서 내의 표를 정확하게 탐지하고 추출하는 방법을 소개하였습니다. 우리는 먼저 YOLOv8 모델을 사용하여 표를 포함하는 영역의 경계 상자를 예측하였고, 그 후 해당 영역을 분할하여 표의 구역을 추출하였습니다. 실험 결과 이 방법은 표의 위치를 정확하게 파악하고 표의 내용을 정확하게 추출하는 것을 성공하였습니다. 문서 처리 및 정보 추출에 있어서 효율적이고 정확한 방법을 제공함으로써, 다양한 응용 분야에서 활용될 수 있을 것으로 기대됩니다.