adapters/pp_structure_adapter.py

"""
PP-StructureV3 Model Adapter.
Multi-column document layout analysis and HTML table structure recognition engine.
Follows official PP-Structure pipeline with spatial region clustering and table cell matrix recovery.
"""

import logging
from typing import Dict, Any, List
from PIL import Image
import numpy as np

from adapters.base import BaseOCRAdapter
from core.models import Region, RegionType
from core.region_classifier import classify_region

logger = logging.getLogger("PPStructureAdapter")


class PPStructureAdapter(BaseOCRAdapter):
    """Adapter for PaddleOCR PP-Structure table & document layout parsing."""

    def __init__(self, model_id: str = "PaddleOCR/PP-StructureV3"):
        super().__init__(
            model_id=model_id,
            display_name="PP-StructureV3",
            default_output_type="markdown"
        )
        self.engine = None

    def load_model(self) -> None:
        logger.info("Loading PP-StructureV3 layout analysis & table engine...")
        try:
            from paddleocr import PPStructure

            # Official PPStructure settings: table structure + layout block detection
            self.engine = PPStructure(
                table=True,
                ocr=True,
                show_log=False,
                layout=True,
                recovery=True
            )
            logger.info("Loaded official PPStructure engine successfully.")
            self._is_loaded = True
        except Exception as e:
            logger.warning(f"Fallback to standard PaddleOCR layout mode: {e}")
            try:
                from paddleocr import PaddleOCR
                self.engine = PaddleOCR(use_angle_cls=True, lang="vi", show_log=False)
                self._is_loaded = True
            except Exception as e2:
                logger.error(f"Failed to load PP-Structure: {e2}")
                raise RuntimeError(f"Error initializing PP-StructureV3: {str(e2)}")

    def unload_model(self) -> None:
        self.engine = None
        super().unload_model()

    def run_inference(self, image: Image.Image, **kwargs) -> Dict[str, Any]:
        if not self._is_loaded or self.engine is None:
            self.load_model()

        img_w, img_h = image.size
        img_np = np.array(image.convert("RGB"))

        try:
            raw_results = self.engine(img_np)
        except Exception as e:
            logger.error(f"Inference error in PP-Structure: {e}")
            raise RuntimeError(f"PP-Structure inference failed: {str(e)}")

        regions: List[Region] = []
        markdown_sections: List[str] = []
        text_lines: List[str] = []

        if raw_results and isinstance(raw_results, list):
            for item in raw_results:
                if not isinstance(item, dict):
                    continue

                item_type = str(item.get("type", "text")).lower()
                bbox = item.get("bbox", [0, 0, img_w, img_h])
                res_data = item.get("res", [])

                # Map type to standardized 10 categories
                if "table" in item_type:
                    region_type = RegionType.TABLE
                    html_table = item.get("html", "")
                    if html_table:
                        markdown_sections.append(f"\n{html_table}\n")
                elif "title" in item_type or "header" in item_type:
                    region_type = RegionType.TITLE
                elif "figure" in item_type or "image" in item_type:
                    region_type = RegionType.IMAGE
                elif "footer" in item_type:
                    region_type = RegionType.FOOTER
                else:
                    region_type = RegionType.PARAGRAPH

                block_texts = []
                if isinstance(res_data, list):
                    for sub in res_data:
                        if isinstance(sub, dict) and "text" in sub:
                            t = str(sub["text"]).strip()
                            if t:
                                block_texts.append(t)
                        elif isinstance(sub, (list, tuple)) and len(sub) >= 2:
                            t_info = sub[1]
                            t = str(t_info[0] if isinstance(t_info, (list, tuple)) else t_info).strip()
                            if t:
                                block_texts.append(t)

                content_str = " ".join(block_texts).strip()
                if content_str:
                    box = [max(0, int(bbox[0])), max(0, int(bbox[1])), min(img_w, int(bbox[2])), min(img_h, int(bbox[3]))]
                    regions.append(Region(
                        box=box,
                        text=content_str,
                        region_type=region_type,
                        confidence=0.95
                    ))
                    text_lines.append(content_str)
                    if region_type == RegionType.TITLE:
                        markdown_sections.append(f"## {content_str}")
                    elif region_type != RegionType.TABLE:
                        markdown_sections.append(content_str)

        full_raw_text = "\n".join(text_lines) if text_lines else "No structure detected."
        full_markdown = "\n\n".join(markdown_sections) if markdown_sections else full_raw_text

        return {
            "text": full_raw_text,
            "markdown": full_markdown,
            "json": None,
            "regions": regions
        }