레이아웃 파서와 Document Graph Parser — 표·그림을 살려 내는 법

def split_pdf(filepath, batch_size=10):
    """입력 PDF를 여러 개의 작은 PDF 파일로 분할"""
    input_pdf = pymupdf.open(filepath)
    num_pages = len(input_pdf)

    ret = []
    for start_page in range(0, num_pages, batch_size):
        end_page = min(start_page + batch_size, num_pages) - 1

        input_file_basename = os.path.splitext(filepath)[0]
        output_file = f"{input_file_basename}_{start_page:04d}_{end_page:04d}.pdf"
        with pymupdf.open() as output_pdf:
            output_pdf.insert_pdf(input_pdf, from_page=start_page, to_page=end_page)
            output_pdf.save(output_file)
            ret.append(output_file)

    input_pdf.close()
    return ret

class LayoutAnalyzer:
    def __init__(self, api_key):
        self.api_key = api_key

    def _upstage_layout_analysis(self, input_file):
        response = requests.post(
            "https://api.upstage.ai/v1/document-ai/layout-analysis",
            headers={"Authorization": f"Bearer {self.api_key}"},
            data={"ocr": False},
            files={"document": open(input_file, "rb")},
        )

        if response.status_code == 200:
            output_file = os.path.splitext(input_file)[0] + ".json"
            with open(output_file, "w") as f:
                json.dump(response.json(), f, ensure_ascii=False)
            return output_file
        else:
            raise ValueError(f"예상치 못한 상태 코드: {response.status_code}")

    def execute(self, input_file):
        return self._upstage_layout_analysis(input_file)

@staticmethod
def normalize_coordinates(coordinates, output_page_size):
    """bbox 4개 꼭짓점에서 축 정렬 사각형을 만든 뒤 0~1로 정규화"""
    x_values = [coord["x"] for coord in coordinates]
    y_values = [coord["y"] for coord in coordinates]
    x1, y1, x2, y2 = min(x_values), min(y_values), max(x_values), max(y_values)

    return (
        x1 / output_page_size[0],
        y1 / output_page_size[1],
        x2 / output_page_size[0],
        y2 / output_page_size[1],
    )


@staticmethod
def crop_image(img, coordinates, output_file):
    """정규화 좌표에 대상 이미지 크기를 곱해 실제 픽셀 좌표로 환산"""
    img_width, img_height = img.size
    x1, y1, x2, y2 = [
        int(coord * dim)
        for coord, dim in zip(coordinates, [img_width, img_height] * 2)
    ]
    cropped_img = img.crop((x1, y1, x2, y2))
    cropped_img.save(output_file)


@staticmethod
def pdf_to_image(pdf_file, page_num, dpi=300):
    """PDF 특정 페이지를 이미지로 렌더링"""
    with pymupdf.open(pdf_file) as doc:
        page = doc[page_num].get_pixmap(dpi=dpi)
        target_page_size = [page.width, page.height]
        page_img = Image.frombytes("RGB", target_page_size, page.samples)
    return page_img

# HTML에서 이미지 경로 업데이트
soup = BeautifulSoup(element["html"], "html.parser")
img_tag = soup.find("img")
if img_tag:
    relative_path = os.path.relpath(output_file, output_folder)
    img_tag["src"] = relative_path.replace("\\", "/")
element["html"] = str(soup)

html_content.append(element["html"])

# 전체 결합 후 — 발견된 태그로 변환 목록을 구성하되 br은 제외
combined_html_content = "\n".join(html_content)
soup = BeautifulSoup(combined_html_content, "html.parser")
all_tags = set([tag.name for tag in soup.find_all()])
html_tag_list = [tag for tag in list(all_tags) if tag not in ["br"]]

md_output = markdown(combined_html_content, convert=html_tag_list)

from typing import TypedDict


class GraphState(TypedDict):
    filepath: str                                    # 원본 경로
    filetype: str                                    # pdf
    page_numbers: list[int]
    batch_size: int
    split_filepaths: list[str]                       # 분할된 파일들
    analyzed_files: list[str]                        # 분석 완료 파일들
    page_elements: dict[int, dict[str, list[dict]]]  # 페이지별 element
    page_metadata: dict[int, dict]
    page_summary: dict[int, str]
    images: list[str]                                # 크롭된 이미지 경로
    image_summary: list[str]
    tables: list[str]
    table_summary: dict[int, str]
    table_markdown: dict[int, str]
    texts: list[str]
    text_summary: dict[int, str]
    language: str

from langgraph.graph import END, StateGraph
from langgraph.checkpoint.memory import MemorySaver

workflow = StateGraph(GraphState)

workflow.add_node("split_pdf_node", split_pdf_node)
workflow.add_node("layout_analyzer_node", layout_analyze_node)
workflow.add_node("page_element_extractor_node", page_element_extractor_node)
workflow.add_node("image_cropper_node", image_cropper_node)
workflow.add_node("table_cropper_node", table_cropper_node)
workflow.add_node("extract_page_text_node", extract_page_text)
workflow.add_node("page_summary_node", page_summary_node)
workflow.add_node("image_summary_node", image_summary_node)
workflow.add_node("table_summary_node", table_summary_node)
workflow.add_node("table_markdown_node", table_markdown_extractor)

workflow.add_edge("split_pdf_node", "layout_analyzer_node")
workflow.add_edge("layout_analyzer_node", "page_element_extractor_node")
workflow.add_edge("page_element_extractor_node", "image_cropper_node")
workflow.add_edge("page_element_extractor_node", "table_cropper_node")
workflow.add_edge("page_element_extractor_node", "extract_page_text_node")
workflow.add_edge("image_cropper_node", "page_summary_node")
workflow.add_edge("table_cropper_node", "page_summary_node")
workflow.add_edge("extract_page_text_node", "page_summary_node")
workflow.add_edge("page_summary_node", "image_summary_node")
workflow.add_edge("page_summary_node", "table_summary_node")
workflow.add_edge("image_summary_node", END)
workflow.add_edge("table_summary_node", "table_markdown_node")
workflow.add_edge("table_markdown_node", END)

workflow.set_entry_point("split_pdf_node")

memory = MemorySaver()
graph = workflow.compile(checkpointer=memory)

prompt = PromptTemplate.from_template(
    """Please summarize the sentence according to the following REQUEST.

REQUEST:
1. Summarize the main points in bullet points.
2. Write the summary in same language as the context.
3. DO NOT translate any technical terms.
4. DO NOT include any unnecessary information.
5. Summary must include important entities, numerical values.

CONTEXT:
{context}

SUMMARY:"
"""
)

llm = ChatOpenAI(model_name="gpt-4o-mini", temperature=0)
text_summary_chain = create_stuff_documents_chain(llm, prompt)

def create_image_summary_data_batches(state: GraphState):
    data_batches = []
    page_numbers = sorted(list(state["page_elements"].keys()))

    for page_num in page_numbers:
        text = state["text_summary"][page_num]
        for image_element in state["page_elements"][page_num]["image_elements"]:
            image_id = int(image_element["id"])
            data_batches.append(
                {
                    "image": state["images"][image_id],  # 이미지 파일 경로
                    "text": text,                        # 관련 텍스트 요약
                    "page": page_num,
                    "id": image_id,
                }
            )
    return GraphState(image_summary_data_batches=data_batches)

class MultiModal:
    def __init__(self, model, system_prompt=None, user_prompt=None):
        self.model = model
        self.system_prompt = system_prompt
        self.user_prompt = user_prompt
        self.init_prompt()

    def encode_image_from_file(self, file_path):
        with open(file_path, "rb") as image_file:
            image_content = image_file.read()
            file_ext = os.path.splitext(file_path)[1].lower()
            if file_ext in [".jpg", ".jpeg"]:
                mime_type = "image/jpeg"
            elif file_ext == ".png":
                mime_type = "image/png"
            else:
                mime_type = "image/unknown"
            b64 = base64.b64encode(image_content).decode("utf-8")
            return f"data:{mime_type};base64,{b64}"

    def create_messages(
        self, image_url, system_prompt=None, user_prompt=None, display_image=True
    ):
        encoded_image = self.encode_image(image_url)
        system_prompt = (
            system_prompt if system_prompt is not None else self.system_prompt
        )
        user_prompt = user_prompt if user_prompt is not None else self.user_prompt

        return [
            {"role": "system", "content": system_prompt},
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": user_prompt},
                    {"type": "image_url", "image_url": {"url": f"{encoded_image}"}},
                ],
            },
        ]

    def batch(
        self,
        image_urls: list[str],
        system_prompts: list[str] = [],
        user_prompts: list[str] = [],
        display_image=False,
    ):
        messages = []
        for image_url, system_prompt, user_prompt in zip(
            image_urls, system_prompts, user_prompts
        ):
            messages.append(
                self.create_messages(image_url, system_prompt, user_prompt, display_image)
            )
        response = self.model.batch(messages)
        return [r.content for r in response]

@chain
def table_markdown_extractor(data_batches):
    llm = ChatOpenAI(temperature=0, model_name="gpt-4o-mini")

    system_prompt = (
        "You are an expert in converting image of the TABLE into markdown format. "
        "Be sure to include all the information in the table. "
        "DO NOT narrate, just answer in markdown format."
    )

    image_paths, system_prompts, user_prompts = [], [], []

    for data_batch in data_batches:
        image_path = data_batch["table"]
        user_prompt_template = """DO NOT wrap your answer in code fences or any XML tags.

###

Output Format:

<table_markdown>

Output must be written in Korean.
"""
        image_paths.append(image_path)
        system_prompts.append(system_prompt)
        user_prompts.append(user_prompt_template)

    multimodal_llm = MultiModal(llm)
    return multimodal_llm.batch(
        image_paths, system_prompts, user_prompts, display_image=False
    )

{
  "metadata": {
    "type": "table",
    "table": ".../83.png",
    "markdown": "| 추진일정 | 2019 | 2020 | 2021 | 2022 | 관계부처 |\n|---|---|---|---|---|---|\n| | 상 | 하 | 상 | 하 | |",
    "page": "7",
    "source": ".../report.pdf",
    "id": "83"
  },
  "page_content": "<table>\n<title>\n주요 과제별 추진일정 검토 및 확정\n</title>\n<summary>...</summary>\n<entities>...</entities>\n<data_insights>...</data_insights>\n</table>"
}