Annotations로 문서에서 데이터 추출하기

Annotations로 문서에서 데이터 추출하기 (Extract Data from Documents via Annotations)

이 쿡북에서는 Mistral의 OCR 모델로 Annotations를 활용해 구조화된 출력(structured outputs)을 얻는 방법을 배워요. 기존 비전 모델만으로는 부족할 때 OCR 모델과 결합해 더 나은 구조화된 데이터 추출을 해내는 노하우랍니다.

출처: 문서

본문

구조화된 출력과 데이터 추출을 위한 Annotations

이 쿡북에서는 Annotations의 기초를 짚어보면서, OCR 모델 기반의 구조화된 출력을 만드는 방법을 살펴볼게요.

지금 사용 중인 비전 모델이 충분히 강력하지 않을 때 이 방법을 쓰면 좋아요. OCR 모델로 비전 모델의 OCR 능력을 보강해서, 더 정확한 구조화된 데이터 추출을 얻을 수 있거든요.

Annotations이 뭔가요?

Mistral Document AI API에서는 두 가지 annotation 기능을 제공해요:

• document_annotation: 입력한 스키마(schema)를 기준으로 문서 전체에 대한 annotation을 돌려줘요.
• box_annotation: OCR 모델이 추출한 bbox(차트/그림 등)에 대한 annotation을 요구사항에 맞춰 돌려줘요. 예를 들어 그림에 대한 설명이나 캡션을 달라고 요청할 수 있어요.

Annotations에 대해 더 자세히 알고 싶다면 여기를 참고하세요.

설정 (Setup)

먼저 mistralai를 설치하고 필요한 파일을 내려받아요.

%%capture
!pip install mistralai

PDF 내려받기 (Download PDF)

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

클라이언트 만들기 (Create Client)

클라이언트를 설정해야 해요. API 키는 AI 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)

Annotations 없는 Mistral OCR

이 쿡북에서는 PDF 파일을 사용해서, annotation을 하고 문서에서 데이터를 추출해 볼게요.

먼저 PDF 파일을 base64로 인코딩하는 함수를 만들어야 해요. 파일을 Mistral 클라우드에 업로드하고 signed url을 쓰는 방법도 있답니다.

import base64

def encode_pdf(pdf_path):
    """Encode the pdf to base64."""
    try:
        with open(pdf_path, "rb") as pdf_file:
            return base64.b64encode(pdf_file.read()).decode('utf-8')
    except FileNotFoundError:
        print(f"Error: The file {pdf_path} was not found.")
        return None
    except Exception as e:  # Added general exception handling
        print(f"Error: {e}")
        return None

이제 함수가 준비됐으니, PDF를 인코딩하고 OCR 모델을 호출해 볼게요.

import requests
import os
import json
# Path to your pdf
pdf_path = "mistral7b.pdf"
# Getting the base64 string
base64_pdf = encode_pdf(pdf_path)
# Call the OCR API
pdf_response = client.ocr.process(
    model="mistral-ocr-4-0",
    document={
        "type": "document_url",
        "document_url": f"data:application/pdf;base64,{base64_pdf}"
    },
    include_image_base64=True,
    extract_header=True,
	extract_footer=True,
	confidence_scores_granularity="word"
)
# Convert response to JSON format
response_dict = json.loads(pdf_response.model_dump_json())

print(json.dumps(response_dict, indent=4)) 

OCR4가 나온 이후부터는 결과를 아래 방식으로 확인할 수 있어요. confidence_scores_granularity="word"로 설정하면, 단어들이 OCR 신뢰도(confidence)에 따라 색으로 구분돼요:

• Red: < 60%
• Orange: 60–80%
• Yellow: 80–95%
• White: 95–100%

단어 위에 마우스를 올리면 정확한 신뢰도 점수를 볼 수 있어요.

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 confidence_color(score_pct: float) -> str | None:
    """Map a 0–100 confidence score to a display color. None means default text."""
    if score_pct < 60:
        return "#d32f2f"  # red
    if score_pct < 80:
        return "#ef6c00"  # orange
    if score_pct < 95:
        return "#f9a825"  # yellow
    return None  # 95–100%: default text color

def apply_confidence_colors(markdown_str: str, word_scores) -> str:
    """Wrap each OCR word in a colored span based on its confidence score."""
    if not word_scores:
        return markdown_str

    colored = markdown_str
    for word in sorted(word_scores, key=lambda w: w.start_index, reverse=True):
        start = word.start_index
        end = start + len(word.text)
        text = colored[start:end]
        if text != word.text:
            continue

        score_pct = word.confidence * 100
        color = confidence_color(score_pct)
        if color is None:
            continue
        colored = (
            colored[:start]
            + f'<span style="color: {color};" title="{score_pct:.1f}%">{text}</span>'
            + colored[end:]
        )
    return colored

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

    Words are color-coded by OCR confidence when word-level scores are available.

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

    Returns:
        Combined markdown string with embedded images and confidence coloring
    """
    markdowns: list[str] = []
    for page in ocr_response.pages:
        page_md = page.markdown
        if page.confidence_scores and page.confidence_scores.word_confidence_scores:
            page_md = apply_confidence_colors(
                page_md, page.confidence_scores.word_confidence_scores
            )

        image_data = {img.id: img.image_base64 for img in page.images}
        markdowns.append(replace_images_in_markdown(page_md, image_data))

    return "\n\n".join(markdowns)
# Display combined markdowns and images
display(Markdown(get_combined_markdown(pdf_response)))

Annotations 있는 Mistral OCR

먼저 Annotation 포맷을 만들어야 해요. 이때는 pydantic을 사용하는 걸 권장합니다. 이 예시에서는 각 bbox의 이미지 타입과 설명을 추출하고, 문서 전체의 언어(language), 저자(authors), 요약(summary)도 함께 뽑아 볼게요.

from pydantic import BaseModel, Field
from enum import Enum

class ImageType(str, Enum):
    GRAPH = "graph"
    TEXT = "text"
    TABLE = "table"
    IMAGE = "image"

class Image(BaseModel):
    image_type: ImageType = Field(..., description="The type of the image. Must be one of 'graph', 'text', 'table' or 'image'.")
    description: str = Field(..., description="A description of the image.")

class Document(BaseModel):
    language: str = Field(..., description="The language of the document in ISO 639-1 code format (e.g., 'en', 'fr').")
    summary: str = Field(..., description="A summary of the document.")
    authors: list[str] = Field(..., description="A list of authors who contributed to the document.")

Annotations용 pydantic 모델을 만들었으니, 이제 OCR 엔드포인트를 호출할 수 있어요. 목표는 문서와 감지된 bbox/이미지에서 정보를 annotation하고 추출하는 거예요.

from mistralai.extra import response_format_from_pydantic_model
# OCR Call with Annotations
annotations_response = client.ocr.process(
    model="mistral-ocr-latest",
    pages=list(range(8)), # Document Annotations has a limit of 8 pages, we recommend spliting your documents when using it; bbox annotations does not have the same limit
    document={
        "type": "document_url",
        "document_url": f"data:application/pdf;base64,{base64_pdf}"
    },
    bbox_annotation_format=response_format_from_pydantic_model(Image),
    document_annotation_format=response_format_from_pydantic_model(Document),
    include_image_base64=False # We are not interested on retrieving the bbox images in this example, only their annotations
  )
# Convert response to JSON format
response_dict = json.loads(annotations_response.model_dump_json())

print(json.dumps(response_dict, indent=4))

Annotations만 출력해 볼게요!

print("Document Annotation:\n", annotations_response.document_annotation)
print("\nBBox/Images:")
for page in annotations_response.pages:
    for image in page.images:
        print("\nImage", image.id)
        print("Location:")
        print(" - top_left_x:", image.top_left_x)
        print(" - top_left_y:", image.top_left_y)
        print(" - bottom_right_x:", image.bottom_right_x)
        print(" - bottom_right_y:", image.bottom_right_y)
        print("BBox/Image Annotation:\n", image.image_annotation)

Annotation 포함한 전체 문서 (Full Document with Annotation)

참고 삼아, 같은 작업을 bbox 이미지를 포함해 해볼게요.

# OCR Call with Annotations
annotations_response = client.ocr.process(
    model="mistral-ocr-latest",
    pages=list(range(8)), # Document Annotations has a limit of 8 pages, we recommend spliting your documents when using it; bbox annotations does not have the same limit
    document={
        "type": "document_url",
        "document_url": f"data:application/pdf;base64,{base64_pdf}"
    },
    bbox_annotation_format=response_format_from_pydantic_model(Image),
    document_annotation_format=response_format_from_pydantic_model(Document),
    include_image_base64=True
  )

이제 OCR 내용과 annotation을 굵게(bold) 표시한 전체 문서를 보여줄게요:

• 문서 시작 부분에 Document Annotation.
• 추출된 각 bbox/이미지 아래에 BBox Annotation.

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

    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 and their annotation
    """
    for img_name, data in images_dict.items():
        markdown_str = markdown_str.replace(
            f"![{img_name}]({img_name})", f"![{img_name}]({data['image']})\n\n**{data['annotation']}**"
        )
    return markdown_str

def get_combined_markdown_annotated(ocr_response: OCRResponse) -> str:
    """
    Combine OCR text, annotation 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 and their annotation
    """
    markdowns: list[str] = ["**" + ocr_response.document_annotation + "**"]
    # Extract images from page
    for page in ocr_response.pages:
        image_data = {}
        for img in page.images:
            image_data[img.id] = {"image":img.image_base64, "annotation": img.image_annotation}
        # Replace image placeholders with actual images
        markdowns.append(replace_images_in_markdown_annotated(page.markdown, image_data))

    return "\n\n".join(markdowns)
# Display combined markdowns and images
display(Markdown(get_combined_markdown_annotated(annotations_response)))

다른 예시 (Other Examples)

#@title PDF Financial Document
# Create the annotations formats
class ImageType(str, Enum):
    GRAPH = "graph"
    TEXT = "text"
    TABLE = "table"
    IMAGE = "image"

class Image(BaseModel):
    image_type: ImageType = Field(..., description="The type of the image. Must be one of 'graph', 'text', 'table' or 'image'.")
    description: str = Field(..., description="A description of the image.")

class Document(BaseModel):
    languages: list[str] = Field(..., description="The list of languages present in the document in ISO 639-1 code format (e.g., 'en', 'fr').")
    summary: str = Field(..., description="A summary of the document.")
# OCR Call with Annotations
annotations_response = client.ocr.process(
    model="mistral-ocr-latest",
    pages=list(range(8)), # Document Annotations has a limit of 8 pages, we recommend spliting your documents when using it; bbox annotations does not have the same limit
    document={
        "type": "document_url",
        "document_url": "https://upload.wikimedia.org/wikipedia/foundation/f/f6/WMF_Mid-Year-Financials_08-09-FINAL.pdf"
    },
    bbox_annotation_format=response_format_from_pydantic_model(Image),
    document_annotation_format=response_format_from_pydantic_model(Document),
    include_image_base64=True
  )
# Display combined markdowns and images
display(Markdown(get_combined_markdown_annotated(annotations_response)))

더 알아보기 (Learn more)

• Mistral OCR Annotations 문서
• Mistral OCR 개요