Spaces:
Running
Running
File size: 4,491 Bytes
b611f38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | """
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
|