File size: 5,538 Bytes
d520909 |
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 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 |
"""
Layout Detection Base Interface
Defines the abstract interface for document layout detection.
"""
from abc import ABC, abstractmethod
from typing import List, Optional, Dict, Any
from dataclasses import dataclass, field
from pydantic import BaseModel, Field
import numpy as np
from ..schemas.core import BoundingBox, LayoutRegion, LayoutType, OCRRegion
class LayoutConfig(BaseModel):
"""Configuration for layout detection."""
# Detection method
method: str = Field(
default="rule_based",
description="Detection method: rule_based, paddle_structure, layoutlm"
)
# Confidence thresholds
min_confidence: float = Field(
default=0.5,
ge=0.0,
le=1.0,
description="Minimum confidence for detected regions"
)
# Region detection settings
detect_tables: bool = Field(default=True, description="Detect table regions")
detect_figures: bool = Field(default=True, description="Detect figure regions")
detect_headers: bool = Field(default=True, description="Detect header/footer")
detect_titles: bool = Field(default=True, description="Detect title/heading")
detect_lists: bool = Field(default=True, description="Detect list structures")
# Merging settings
merge_threshold: float = Field(
default=0.7,
ge=0.0,
le=1.0,
description="IoU threshold for merging overlapping regions"
)
# GPU settings
use_gpu: bool = Field(default=True, description="Use GPU acceleration")
gpu_id: int = Field(default=0, ge=0, description="GPU device ID")
# Table detection specific
table_min_rows: int = Field(default=2, ge=1, description="Minimum rows for table")
table_min_cols: int = Field(default=2, ge=1, description="Minimum columns for table")
# Title/heading detection
title_max_lines: int = Field(default=3, description="Max lines for title")
heading_font_ratio: float = Field(
default=1.2,
description="Font size ratio vs body text for headings"
)
@dataclass
class LayoutResult:
"""Result of layout detection for a page."""
page: int
regions: List[LayoutRegion] = field(default_factory=list)
image_width: int = 0
image_height: int = 0
processing_time_ms: float = 0.0
# Error handling
success: bool = True
error: Optional[str] = None
def get_regions_by_type(self, layout_type: LayoutType) -> List[LayoutRegion]:
"""Get regions of a specific type."""
return [r for r in self.regions if r.type == layout_type]
def get_tables(self) -> List[LayoutRegion]:
"""Get table regions."""
return self.get_regions_by_type(LayoutType.TABLE)
def get_figures(self) -> List[LayoutRegion]:
"""Get figure regions."""
return self.get_regions_by_type(LayoutType.FIGURE)
def get_text_regions(self) -> List[LayoutRegion]:
"""Get text-based regions (paragraph, title, heading, list)."""
text_types = {
LayoutType.TEXT,
LayoutType.TITLE,
LayoutType.HEADING,
LayoutType.PARAGRAPH,
LayoutType.LIST,
}
return [r for r in self.regions if r.type in text_types]
class LayoutDetector(ABC):
"""
Abstract base class for layout detectors.
"""
def __init__(self, config: Optional[LayoutConfig] = None):
"""
Initialize layout detector.
Args:
config: Layout detection configuration
"""
self.config = config or LayoutConfig()
self._initialized = False
@abstractmethod
def initialize(self):
"""Initialize the detector (load models, etc.)."""
pass
@abstractmethod
def detect(
self,
image: np.ndarray,
page_number: int = 0,
ocr_regions: Optional[List[OCRRegion]] = None,
) -> LayoutResult:
"""
Detect layout regions in an image.
Args:
image: Image as numpy array (RGB, HWC format)
page_number: Page number
ocr_regions: Optional OCR regions for text-aware detection
Returns:
LayoutResult with detected regions
"""
pass
def detect_batch(
self,
images: List[np.ndarray],
page_numbers: Optional[List[int]] = None,
ocr_results: Optional[List[List[OCRRegion]]] = None,
) -> List[LayoutResult]:
"""
Detect layout in multiple images.
Args:
images: List of images
page_numbers: Optional page numbers
ocr_results: Optional OCR regions for each page
Returns:
List of LayoutResult
"""
if page_numbers is None:
page_numbers = list(range(len(images)))
if ocr_results is None:
ocr_results = [None] * len(images)
results = []
for img, page_num, ocr in zip(images, page_numbers, ocr_results):
results.append(self.detect(img, page_num, ocr))
return results
@property
def name(self) -> str:
"""Return detector name."""
return self.__class__.__name__
@property
def is_initialized(self) -> bool:
"""Check if detector is initialized."""
return self._initialized
def __enter__(self):
"""Context manager entry."""
if not self._initialized:
self.initialize()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Context manager exit."""
pass
|