core/models.py

"""
Core Data Models for Local-First OCR Benchmarking & Region Classification.
Standardizes bounding boxes, semantic categories, and benchmark results across all 6 models.
"""

from dataclasses import dataclass, field, asdict
from typing import List, Dict, Any, Optional
from enum import Enum


class RegionType(str, Enum):
    """Standardized 10 semantic region categories."""
    TITLE_HEADER = "Title/Header"
    PARAGRAPH = "Paragraph"
    TEXT = "Text"
    TABLE = "Table"
    TABLE_CELL = "Table cell"
    KEY_VALUE = "Key-Value"
    NUMBER_PRICE = "Number/Price"
    IMAGE = "Image"
    FOOTER = "Footer"
    OTHER = "Other"


# Standard 10 categories as list of strings
ALL_REGION_TYPES: List[str] = [rt.value for rt in RegionType]

# Aesthetic color palette for visualization (RGBA / Hex)
REGION_COLORS: Dict[str, Dict[str, Any]] = {
    RegionType.TITLE_HEADER.value: {
        "hex": "#8B5CF6",         # Purple
        "rgb": (139, 92, 246),
        "fill_rgba": (139, 92, 246, 50),
        "badge_bg": "#EDE9FE",
        "badge_text": "#5B21B6"
    },
    RegionType.PARAGRAPH.value: {
        "hex": "#10B981",         # Emerald
        "rgb": (16, 185, 129),
        "fill_rgba": (16, 185, 129, 45),
        "badge_bg": "#D1FAE5",
        "badge_text": "#065F46"
    },
    RegionType.TEXT.value: {
        "hex": "#06B6D4",         # Cyan / Teal
        "rgb": (6, 182, 212),
        "fill_rgba": (6, 182, 212, 40),
        "badge_bg": "#CFFAFE",
        "badge_text": "#155E75"
    },
    RegionType.TABLE.value: {
        "hex": "#3B82F6",         # Blue
        "rgb": (59, 130, 246),
        "fill_rgba": (59, 130, 246, 55),
        "badge_bg": "#DBEAFE",
        "badge_text": "#1E40AF"
    },
    RegionType.TABLE_CELL.value: {
        "hex": "#60A5FA",         # Light Blue
        "rgb": (96, 165, 250),
        "fill_rgba": (96, 165, 250, 40),
        "badge_bg": "#EFF6FF",
        "badge_text": "#1D4ED8"
    },
    RegionType.KEY_VALUE.value: {
        "hex": "#F59E0B",         # Amber
        "rgb": (245, 158, 11),
        "fill_rgba": (245, 158, 11, 55),
        "badge_bg": "#FEF3C7",
        "badge_text": "#92400E"
    },
    RegionType.NUMBER_PRICE.value: {
        "hex": "#F97316",         # Orange
        "rgb": (249, 115, 22),
        "fill_rgba": (249, 115, 22, 60),
        "badge_bg": "#FFEDD5",
        "badge_text": "#9A3412"
    },
    RegionType.IMAGE.value: {
        "hex": "#EC4899",         # Pink / Rose
        "rgb": (236, 72, 153),
        "fill_rgba": (236, 72, 153, 50),
        "badge_bg": "#FCE7F3",
        "badge_text": "#9D174D"
    },
    RegionType.FOOTER.value: {
        "hex": "#64748B",         # Slate
        "rgb": (100, 116, 139),
        "fill_rgba": (100, 116, 139, 45),
        "badge_bg": "#F1F5F9",
        "badge_text": "#334155"
    },
    RegionType.OTHER.value: {
        "hex": "#9CA3AF",         # Gray
        "rgb": (156, 163, 175),
        "fill_rgba": (156, 163, 175, 40),
        "badge_bg": "#F3F4F6",
        "badge_text": "#374151"
    }
}


@dataclass
class Region:
    """Represents a single detected spatial layout region on the document image."""
    box: List[int]                     # [x1, y1, x2, y2] in exact pixel coordinates
    text: str                          # Recognized text content
    region_type: str = RegionType.TEXT.value  # One of the 10 standardized categories
    confidence: Optional[float] = None # Confidence score between 0.0 and 1.0
    details: Optional[Dict[str, Any]] = None # Extra metadata (e.g. table HTML, key-value split)

    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)


@dataclass
class OCRModelOutput:
    """Standardized result output from any of the 6 OCR models."""
    model_name: str
    model_id: str
    status: str                        # 'SUCCESS' or 'ERROR'
    inference_time_seconds: Optional[float] = None
    inference_time_str: str = "N/A"
    text: Optional[str] = None
    markdown: Optional[str] = None
    json: Optional[Any] = None
    output_type: str = "markdown"
    word_count: int = 0
    regions: List[Region] = field(default_factory=list)
    region_counts: Dict[str, int] = field(default_factory=dict)
    annotated_image_path: Optional[str] = None
    annotated_image_base64: Optional[str] = None
    error: Optional[str] = None

    def to_dict(self) -> Dict[str, Any]:
        d = asdict(self)
        d["regions"] = [r.to_dict() if isinstance(r, Region) else r for r in self.regions]
        return d