OCR 탐색과 간단한 구조화 출력

OCR 탐색과 간단한 구조화 출력 (deprecated) — OCR Exploration and Simple Structured Outputs

Mistral OCR의 기초를 탐색하고, 이를 기존 모델과 결합해 구조화된 출력(structured outputs)을 얻는 방법을 배우는 문서예요. 더 나은 결과를 위해 새 Annotations 기능을 쓰는 것을 권장하지만, 기본 개념을 이해하는 데 유용해요.

출처: 문서

본문

이 쿡북에서는 OCR의 기초를 살펴보고, OCR 모델 기반의 구조화된 출력을 얻기 위해 기존 모델과 함께 활용하는 방법을 배워요 (더 나은 결과를 위해 새 Annotations 기능 사용을 권장합니다). 현재 비전 모델이 충분히 강력하지 않을 때, OCR 모델로 비전·OCR 능력을 보강해 더 나은 구조화된 데이터 추출을 얻고 싶을 때 쓰면 좋아요.

사용 모델:

  • Mistral OCR
  • Pixtral 12B & Ministral 8B

더 최신의 구조화 출력 가이드는 [Data Extraction에 대한 Annotations 쿡북]을 참고하세요.

설정 (Setup)

먼저 mistralai를 설치하고 필요한 파일을 다운로드해요.

%%capture
!pip install mistralai

PDF와 이미지 파일 다운로드

%%capture
!wget https://raw.githubusercontent.com/mistralai/cookbook/refs/heads/main/mistral/ocr/mistral7b.pdf
!wget https://raw.githubusercontent.com/mistralai/cookbook/refs/heads/main/mistral/ocr/receipt.png

PDF로 Mistral OCR 사용하기

클라이언트를 설정해요. API 키는 [Studio]에서 만들 수 있어요.

# Initialize Mistral client with API key
from mistralai.client import Mistral

api_key = "API_KEY" # Replace with your API key
client = Mistral(api_key=api_key)

OCR을 적용할 수 있는 파일 종류는 두 가지가 있어요:

  • PDF 파일
  • 이미지 파일

PDF 파일부터 시작해 볼게요.

# Import required libraries
from pathlib import Path
from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk
import json

# Verify PDF file exists
pdf_file = Path("mistral7b.pdf")
assert pdf_file.is_file()

# Upload PDF file to Mistral's OCR service
uploaded_file = client.files.upload(
    file={
        "file_name": pdf_file.stem,
        "content": pdf_file.read_bytes(),
    },
    purpose="ocr",
)

# Get URL for the uploaded file
signed_url = client.files.get_signed_url(file_id=uploaded_file.id, expiry=1)

# Process PDF with OCR, including embedded images
pdf_response = client.ocr.process(
    document=DocumentURLChunk(document_url=signed_url.url),
    model="mistral-ocr-latest",
    include_image_base64=True
)

# Convert response to JSON format
response_dict = json.loads(pdf_response.model_dump_json())

print(json.dumps(response_dict, indent=4)[0:1000]) # check the first 1000 characters

결과는 다음과 같이 확인해요.

from mistralai.client.models import OCRResponse
from IPython.display import Markdown, display

def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str:
    """
    Replace image placeholders in markdown with base64-encoded images.

    Args:
        markdown_str: Markdown text containing image placeholders
        images_dict: Dictionary mapping image IDs to base64 strings

    Returns:
        Markdown text with images replaced by base64 data
    """
    for img_name, base64_str in images_dict.items():
        markdown_str = markdown_str.replace(
            f"![{img_name}]({img_name})", f"![{img_name}]({base64_str})"
        )
    return markdown_str

def get_combined_markdown(ocr_response: OCRResponse) -> str:
    """
    Combine OCR text and images into a single markdown document.

    Args:
        ocr_response: Response from OCR processing containing text and images

    Returns:
        Combined markdown string with embedded images
    """
    markdowns: list[str] = []
    # Extract images from page
    for page in ocr_response.pages:
        image_data = {}
        for img in page.images:
            image_data[img.id] = img.image_base64
        # Replace image placeholders with actual images
        markdowns.append(replace_images_in_markdown(page.markdown, image_data))

    return "\n\n".join(markdowns)

# Display combined markdowns and images
display(Markdown(get_combined_markdown(pdf_response)))

이미지로 Mistral OCR 사용하기

PDF 외에도 이미지 파일을 처리할 수 있어요.

import base64

# Verify image exists
image_file = Path("receipt.png")
assert image_file.is_file()

# Encode image as base64 for API
encoded = base64.b64encode(image_file.read_bytes()).decode()
base64_data_url = f"data:image/jpeg;base64,{encoded}"

# Process image with OCR
image_response = client.ocr.process(
    document=ImageURLChunk(image_url=base64_data_url),
    model="mistral-ocr-latest"
)

# Convert response to JSON
response_dict = json.loads(image_response.model_dump_json())
json_string = json.dumps(response_dict, indent=4)
print(json_string)

OCR 결과에서 구조화 데이터 추출하기

OCR 결과는 다른 모델로 추가 처리할 수 있어요. 목표는 이 결과에서 구조화된 데이터를 추출하는 거예요. 이를 위해 OCR 모델과 함께 pixtral-12b-latest 모델을 사용해 더 좋고 품질 높은 답변을 얻을 수 있어요.

# Get OCR results for processing
image_ocr_markdown = image_response.pages[0].markdown

# Get structured response from model
chat_response = client.chat.complete(
    model="pixtral-12b-latest",
    messages=[
        {
            "role": "user",
            "content": [
                ImageURLChunk(image_url=base64_data_url),
                TextChunk(
                    text=(
                        f"This is image's OCR in markdown:\n\n{image_ocr_markdown}\n.\n"
                        "Convert this into a sensible structured json response. "
                        "The output should be strictly be json with no extra commentary"
                    )
                ),
            ],
        }
    ],
    response_format={"type": "json_object"},
    temperature=0,
)

# Parse and return JSON response
response_dict = json.loads(chat_response.choices[0].message.content)
print(json.dumps(response_dict, indent=4))

위 예시에서는 이미 비전 작업이 가능한 모델을 활용했어요. 하지만 텍스트 전용 모델로도 구조화된 출력을 만들 수 있어요. 이 경우 메시지에 이미지를 포함하지 않는다는 점에 주목하세요:

# Get OCR results for processing
image_ocr_markdown = image_response.pages[0].markdown

# Get structured response from model
chat_response = client.chat.complete(
    model="ministral-8b-latest",
    messages=[
        {
            "role": "user",
            "content": [
                TextChunk(
                    text=(
                        f"This is image's OCR in markdown:\n\n{image_ocr_markdown}\n.\n"
                        "Convert this into a sensible structured json response. "
                        "The output should be strictly be json with no extra commentary"
                    )
                ),
            ],
        }
    ],
    response_format={"type": "json_object"},
    temperature=0,
)

# Parse and return JSON response
response_dict = json.loads(chat_response.choices[0].message.content)
print(json.dumps(response_dict, indent=4))

모두 합치기 — Mistral OCR + 커스텀 구조화 출력

image_path 파일을 받아 특정 형식의 JSON 구조화 출력을 반환하는 간단한 함수를 설계해 볼게요. 이 예시에서는 임의로 다음과 같은 출력을 원한다고 정했어요:

class StructuredOCR:
    file_name: str # can be any string
    topics: list[str] # must be a list of strings
    languages: str # string
    ocr_contents: dict # any dictionary, can be freely defined by the model

[커스텀 구조화 출력]을 사용해 구현할게요.

from enum import Enum
from pathlib import Path
from pydantic import BaseModel
import base64

class StructuredOCR(BaseModel):
    file_name: str
    topics: list[str]
    languages: str
    ocr_contents: dict

def structured_ocr(image_path: str) -> StructuredOCR:
    """
    Process an image using OCR and extract structured data.

    Args:
        image_path: Path to the image file to process

    Returns:
        StructuredOCR object containing the extracted data

    Raises:
        AssertionError: If the image file does not exist
    """
    # Validate input file
    image_file = Path(image_path)
    assert image_file.is_file(), "The provided image path does not exist."

    # Read and encode the image file
    encoded_image = base64.b64encode(image_file.read_bytes()).decode()
    base64_data_url = f"data:image/jpeg;base64,{encoded_image}"

    # Process the image using OCR
    image_response = client.ocr.process(
        document=ImageURLChunk(image_url=base64_data_url),
        model="mistral-ocr-latest"
    )
    image_ocr_markdown = image_response.pages[0].markdown

    # Parse the OCR result into a structured JSON response
    chat_response = client.chat.parse(
        model="pixtral-12b-latest",
        messages=[
            {
                "role": "user",
                "content": [
                    ImageURLChunk(image_url=base64_data_url),
                    TextChunk(text=(
                        f"This is the image's OCR in markdown:\n{image_ocr_markdown}\n.\n"
                        "Convert this into a structured JSON response "
                        "with the OCR contents in a sensible dictionnary."
                    )
                    )
                ]
            }
        ],
        response_format=StructuredOCR,
        temperature=0
    )

    return chat_response.choices[0].message.parsed

이제 OCR 모델로 처리된 어떤 이미지에서든 구조화된 출력을 추출할 수 있어요.

# Example usage
image_path = "receipt.png" # Path to sample receipt image
structured_response = structured_ocr(image_path) # Process image and extract data

# Parse and return JSON response
response_dict = json.loads(structured_response.model_dump_json())
print(json.dumps(response_dict, indent=4))

원본 이미지는 아래와 같은 방법으로 비교 확인할 수 있어요.

from PIL import Image

image = Image.open(image_path)
image.resize((image.width // 5, image.height // 5))

더 알아보기 (Learn more)