pdf를 text로 변환하기 위해서 marker-pdf , MarkItDown 등 여러가지를 써봤는데 최종적으로는 olmocr을 사용하기로 한다.
그 이유는 다음과 같다.
marker-pdf
PDF를 마크다운형식으로 굉장히 잘 가져온다고 판단했다. 하지만 표에 있는 내용을 모두 들고오지 못하는 문제점을 보였다. (아래 마크다운 참고)
#### <표 1> 작업대차 사용 중 위험요인 및 안전대책 | 감전 | 손상에의한감전재해 | 및작동여부확인 | | 재해 | | |
markitdown
사용해보신 분들은 알겠지만 속도가 매우 빠르다.
다만, 내가 어떤 부분의 설정을 잘 못한지 모르겠지만 한국어가 인식이 되지 않았다. 이상한 언어들이 출력되는 것을 보여서 일단 보류했다.
olmocr
데모페이지로 실험을 해봤는데 표도 잘 들고오고, 깔끔하게 텍스트를 추출함을 보여서 해당 라이브러리를 사용하기로 했다.
심지어 해당 라이브러리는 코드가 완전히 공개되어 있다는 사실.. !!
| 발생 형태 | 재해발생 위험요인 | 안전 대책 | |----------|-----------------|----------| | 공통 | 근로자의 불안전한 행동으로 인한 재해발생 | 관리감독자의 선임, 근로자의 안전보건에 관한 교육 실시, 안전수칙 준수 | | 추락 | 작업발판에서 해체작업 중 추락 | 작업발판 주위에 안전난간 설치 | | | 작업발판에서 작업 중 불시 이동에 의한 추락 | 작업자가 작업발판에 탑승한 채로 이동금지 | | | 작업대차(캔틸레버슬래브 해체형)의 불시 이동에 의한 추락 | 작업발판 바퀴에 구름방지장치 설치 | | | 하부작업대로 이동 중 추락 | 작업대차 바퀴에 구름방지장치 설치 | | 낙하 | 해체된 거푸집 인양 중 자재 낙하 | 작업자가 대차에 탑승한 채로 이동금지 | | | 거푸집해체중 거푸집, 작업 공구 등의 낙하 | 작업대차의 안전통로 확보 (등반이 형 수직사다리, 안전난간 설치 등) | | 붕괴 | 부재의 결함, 누락으로 인한 붕괴 | 작업시작 전 작업대차의 주요 구조부위, 연결 부위 등 점검실시 | | | 작업발판의 과부하로 인한 작업대 붕괴 | 작업 전 과부하 방지장치의 이상유무 확인 및 작업발판 상승 시 구조물에 걸림 여부 확인 | | 감전 | 이동용 전동공구 및 전선의 손상에 의한 감전재해 | 전원 인출 분전반에 누전차단기 부착 및 작동여부 확인 |
toolkit를 이용하면 편리하게 사용할 수 있다고 하는데 로컬PC에 GPU가 없는 관계로 Google Colab에서 진행했다.
패키지 설치
!pip install olmocr PyPDF2
!sudo apt-get update
!sudo apt-get install poppler-utils필요한 라이브러리 import
import torch
import base64
import urllib.request
from io import BytesIO
from PIL import Image
from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
from olmocr.data.renderpdf import render_pdf_to_base64png
from olmocr.prompts import build_finetuning_prompt
from olmocr.prompts.anchor import get_anchor_textModel Load
# Initialize the model
model = Qwen2VLForConditionalGeneration.from_pretrained("allenai/olmOCR-7B-0225-preview", torch_dtype=torch.bfloat16).eval()
processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)OCR → TEXT
이제 PDF 한 페이지를 OCR 처리하여 텍스트를 추출하는 파이프라인을 구현한다.
PDF의 특정 페이지를 이미지로 변환
render_pdf_to_base64png()함수를 사용하여 PDF 페이지를 이미지(.png)로 변환한 후, Base64로 인코딩함.target_longest_image_dim=1024옵션을 사용하여 최대 크기를 1024px로 설정.
PDF 문서의 메타데이터 활용하여 프롬프트 생성
get_anchor_text()함수로 PDF의 특정 부분(예: 첫 번째 페이지)에서 최대 4000자까지 텍스트를 추출.build_finetuning_prompt()함수를 이용해 OCR 모델에 입력할 프롬프트를 생성.
OCR 모델 입력 데이터 생성
OCR 모델이 처리할 입력 데이터를
messages리스트에 저장.프롬프트(
prompt)와 변환된 이미지(image_base64)를 포함.apply_chat_template()함수로 채팅 형식의 템플릿을 적용.Base64로 인코딩된 이미지를 디코딩하여
PIL.Image객체로 변환.
모델 입력 데이터 변환 및 GPU 적용
processor()함수를 사용해 OCR 모델이 이해할 수 있는 형식으로 변환.
텍스트 생성 및 디코딩
def ocr_pdf_page(pdf_path, page_number):
# Render page 1 to an image
image_base64 = render_pdf_to_base64png(pdf_path, page_number, target_longest_image_dim=1024)
# Build the prompt, using document metadata
anchor_text = get_anchor_text(pdf_path, 1, pdf_engine="pdfreport", target_length=4000)
prompt = build_finetuning_prompt(anchor_text)
# Build the full prompt
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
],
}
]
# Apply the chat template and processor
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
main_image = Image.open(BytesIO(base64.b64decode(image_base64)))
inputs = processor(
text=[text],
images=[main_image],
padding=True,
return_tensors="pt",
)
inputs = {key: value.to(device) for (key, value) in inputs.items()}
# Generate the output
output = model.generate(
**inputs,
temperature=0.8,
max_new_tokens=2048, # 생성 토큰 설정
num_return_sequences=1,
do_sample=True,
)
# Decode the output
prompt_length = inputs["input_ids"].shape[1]
new_tokens = output[:, prompt_length:]
text_output = processor.tokenizer.batch_decode(
new_tokens, skip_special_tokens=True
)
print(text_output)
return text_output[0]신경써야할 부분은 max_new_token이 허깅페이스 코드에는 50으로 설정되어있는데 PDF 문서 내 테텍스트 숫자가 50 이상이면 텍스트 출력이 잘리게 된다. 그러니 넉넉하게 주도록 할 것!
from PyPDF2 import PdfReader
def get_pdf_page_count(pdf_path):
"""Gets the number of pages in a PDF file using PyPDF2.
Args:
pdf_path: The path to the PDF file.
Returns:
The number of pages in the PDF, or None if an error occurs.
"""
try:
pdf = PdfReader(pdf_path)
return len(pdf.pages)
except Exception as e:
print(f"Error using PyPDF2: {e}")
return None
page_count = get_pdf_page_count(pdf_list[0])
if page_count:
print(f"The PDF has {page_count} pages.")
# output
The PDF has 13 pages.해당 PDF 가 몇 페이지로 구성되어 있는 출력하는 코드이다.
for idx , pdf_file in enumerate(pdf_list) :
FILE_PATH = pdf_file
NUM_PAGES = get_pdf_page_count(FILE_PATH)
pages = []
for page_idx in range(NUM_PAGES):
print(f"Processing page {page_idx+1} of {NUM_PAGES}")
page_json = ocr_pdf_page(FILE_PATH, page_idx+1)
pages.append(page_json)
file_name = os.path.splitext(os.path.basename(FILE_PATH))[0] # 파일 이름
save_path = './pdf2txt/' # 마크다운 파일 저장할 경로 설정
output_file = os.path.join(save_path, file_name + ".md") #전체 경로
# 파일 열기 (쓰기 모드)
with open(output_file, "w", encoding="utf-8") as f:
for page_json in pages:
try:
data = json.loads(page_json)
f.write(f"{data['natural_text']}\n\n")
except Exception as e:
f.write(f"Error processing page_json: {str(e)}\n\n")
print(f"{idx+1}/{len(pdf_list)} 문서 변환 완료 : Data saved to {output_file}")pdf_list에 pdf 을 넣어놓고, 하나씩 꺼내면서 md 확장자로 저장하는 코드이다. 근데 사실 텍스트만 추출되서 txt 확장자로 저장하는 것이 더 좋을 거 같다.
마크다운이 잘 저장되는 것을 확인할 수 있다.