318 lines
11 KiB
Python
318 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.logging_config import get_logger
|
|
from app.models.document import Document, DocumentPage
|
|
from app.repositories.document_repository import (
|
|
DocumentPageRepository,
|
|
DocumentTableRepository,
|
|
DocumentTextBlockRepository,
|
|
)
|
|
from app.storage.provider import get_storage_provider
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class LayoutService:
|
|
"""Document layout analysis service using OpenCV-based detection."""
|
|
|
|
def __init__(self, db: Session) -> None:
|
|
self.db = db
|
|
self.storage = get_storage_provider()
|
|
self.page_repo = DocumentPageRepository(db)
|
|
self.text_block_repo = DocumentTextBlockRepository(db)
|
|
self.table_repo = DocumentTableRepository(db)
|
|
|
|
def analyze_document_layout(self, document: Document) -> dict[str, Any]:
|
|
"""Analyze the layout of all pages in a document."""
|
|
layout_results: dict[str, Any] = {"pages": []}
|
|
|
|
for page in document.pages:
|
|
page_layout = self._analyze_page_layout(page)
|
|
layout_results["pages"].append(page_layout)
|
|
|
|
return layout_results
|
|
|
|
def _analyze_page_layout(self, page: DocumentPage) -> dict[str, Any]:
|
|
"""Analyze layout of a single page."""
|
|
result: dict[str, Any] = {
|
|
"page_number": page.page_number,
|
|
"width": page.width,
|
|
"height": page.height,
|
|
"tables": [],
|
|
"lines": [],
|
|
"rectangles": [],
|
|
"text_regions": [],
|
|
"image_regions": [],
|
|
}
|
|
|
|
if not page.image_path:
|
|
return result
|
|
|
|
image_path = self.storage.get_absolute_path(page.image_path)
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
logger.warning("layout_image_read_failed", page_id=str(page.id))
|
|
return result
|
|
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
|
|
# Detect lines
|
|
result["lines"] = self._detect_lines(gray)
|
|
|
|
# Detect rectangles (potential table cells/borders)
|
|
result["rectangles"] = self._detect_rectangles(gray)
|
|
|
|
# Detect tables
|
|
tables = self._detect_tables(gray, image.shape)
|
|
result["tables"] = tables
|
|
|
|
# Store detected tables in the database
|
|
for table_data in tables:
|
|
self.table_repo.create_table(
|
|
page_id=page.id,
|
|
x=table_data["x"],
|
|
y=table_data["y"],
|
|
width=table_data["width"],
|
|
height=table_data["height"],
|
|
rows=table_data["rows"],
|
|
columns=table_data["columns"],
|
|
data=table_data.get("cells"),
|
|
)
|
|
|
|
# Detect watermarks
|
|
watermark = self._detect_watermark(gray, image.shape)
|
|
if watermark:
|
|
result["watermark"] = watermark
|
|
|
|
return result
|
|
|
|
def _detect_lines(self, gray: np.ndarray) -> list[dict[str, Any]]:
|
|
"""Detect horizontal and vertical lines in the image."""
|
|
lines_detected: list[dict[str, Any]] = []
|
|
|
|
# Apply edge detection
|
|
edges = cv2.Canny(gray, 50, 150, apertureSize=3)
|
|
|
|
# Detect lines using Hough transform
|
|
lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=100, minLineLength=50, maxLineGap=10)
|
|
|
|
if lines is not None:
|
|
for line in lines:
|
|
x1, y1, x2, y2 = line[0]
|
|
length = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
|
|
|
|
# Classify as horizontal or vertical
|
|
angle = np.degrees(np.arctan2(y2 - y1, x2 - x1))
|
|
if abs(angle) < 5 or abs(angle - 180) < 5:
|
|
orientation = "horizontal"
|
|
elif abs(angle - 90) < 5 or abs(angle + 90) < 5:
|
|
orientation = "vertical"
|
|
else:
|
|
orientation = "diagonal"
|
|
|
|
lines_detected.append({
|
|
"x1": float(x1),
|
|
"y1": float(y1),
|
|
"x2": float(x2),
|
|
"y2": float(y2),
|
|
"length": float(length),
|
|
"orientation": orientation,
|
|
})
|
|
|
|
return lines_detected
|
|
|
|
def _detect_rectangles(self, gray: np.ndarray) -> list[dict[str, Any]]:
|
|
"""Detect rectangular regions in the image."""
|
|
rectangles: list[dict[str, Any]] = []
|
|
|
|
# Binary threshold
|
|
_, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
|
|
|
|
# Find contours
|
|
contours, _ = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
for contour in contours:
|
|
# Approximate the contour
|
|
peri = cv2.arcLength(contour, True)
|
|
approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
|
|
|
|
# If approximation has 4 vertices, it's likely a rectangle
|
|
if len(approx) == 4:
|
|
x, y, w, h = cv2.boundingRect(approx)
|
|
# Filter out very small or very large rectangles
|
|
area = w * h
|
|
if area > 500 and w > 10 and h > 10:
|
|
rectangles.append({
|
|
"x": float(x),
|
|
"y": float(y),
|
|
"width": float(w),
|
|
"height": float(h),
|
|
"area": float(area),
|
|
})
|
|
|
|
return rectangles
|
|
|
|
def _detect_tables(self, gray: np.ndarray, image_shape: tuple) -> list[dict[str, Any]]:
|
|
"""Detect table structures using morphological operations."""
|
|
tables: list[dict[str, Any]] = []
|
|
h, w = image_shape[:2]
|
|
|
|
# Binary threshold
|
|
_, binary = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
|
|
|
|
# Detect horizontal lines
|
|
horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (max(w // 30, 1), 1))
|
|
horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel, iterations=2)
|
|
|
|
# Detect vertical lines
|
|
vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(h // 30, 1)))
|
|
vertical = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel, iterations=2)
|
|
|
|
# Combine horizontal and vertical lines to find intersections
|
|
table_mask = cv2.add(horizontal, vertical)
|
|
|
|
# Find contours of table regions
|
|
contours, _ = cv2.findContours(table_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
for contour in contours:
|
|
x, y, cw, ch = cv2.boundingRect(contour)
|
|
area = cw * ch
|
|
|
|
# Filter: table should be reasonably sized
|
|
if area < 5000 or cw < 50 or ch < 30:
|
|
continue
|
|
|
|
# Estimate rows and columns
|
|
rows, columns = self._estimate_table_dimensions(
|
|
table_mask[y:y+ch, x:x+cw], cw, ch
|
|
)
|
|
|
|
if rows >= 1 and columns >= 1:
|
|
# Extract cell contents
|
|
cells = self._extract_table_cells(
|
|
gray[y:y+ch, x:x+cw], rows, columns, cw, ch
|
|
)
|
|
|
|
tables.append({
|
|
"x": float(x),
|
|
"y": float(y),
|
|
"width": float(cw),
|
|
"height": float(ch),
|
|
"rows": rows,
|
|
"columns": columns,
|
|
"cells": cells,
|
|
})
|
|
|
|
return tables
|
|
|
|
def _estimate_table_dimensions(
|
|
self,
|
|
table_region: np.ndarray,
|
|
width: int,
|
|
height: int,
|
|
) -> tuple[int, int]:
|
|
"""Estimate the number of rows and columns in a table region."""
|
|
# Project horizontal lines
|
|
h_projection = np.sum(table_region, axis=1)
|
|
h_peaks = self._count_peaks(h_projection, height)
|
|
|
|
# Project vertical lines
|
|
v_projection = np.sum(table_region, axis=0)
|
|
v_peaks = self._count_peaks(v_projection, width)
|
|
|
|
rows = max(1, h_peaks - 1)
|
|
columns = max(1, v_peaks - 1)
|
|
|
|
return rows, columns
|
|
|
|
def _count_peaks(self, projection: np.ndarray, total_length: int) -> int:
|
|
"""Count significant peaks in a projection array."""
|
|
if len(projection) == 0:
|
|
return 0
|
|
|
|
threshold = np.max(projection) * 0.3
|
|
above_threshold = projection > threshold
|
|
|
|
# Count transitions from below to above threshold
|
|
peaks = 0
|
|
in_peak = False
|
|
for val in above_threshold:
|
|
if val and not in_peak:
|
|
peaks += 1
|
|
in_peak = True
|
|
elif not val:
|
|
in_peak = False
|
|
|
|
return peaks
|
|
|
|
def _extract_table_cells(
|
|
self,
|
|
table_gray: np.ndarray,
|
|
rows: int,
|
|
columns: int,
|
|
width: int,
|
|
height: int,
|
|
) -> dict[str, Any]:
|
|
"""Extract cell structure data from a table region."""
|
|
cell_height = height / max(rows, 1)
|
|
cell_width = width / max(columns, 1)
|
|
|
|
cells: dict[str, Any] = {"rows": rows, "columns": columns, "data": []}
|
|
|
|
for r in range(rows):
|
|
row_data = []
|
|
for c in range(columns):
|
|
cell_x = int(c * cell_width)
|
|
cell_y = int(r * cell_height)
|
|
cell_w = int(cell_width)
|
|
cell_h = int(cell_height)
|
|
|
|
row_data.append({
|
|
"row": r,
|
|
"col": c,
|
|
"x": cell_x,
|
|
"y": cell_y,
|
|
"width": cell_w,
|
|
"height": cell_h,
|
|
})
|
|
cells["data"].append(row_data)
|
|
|
|
return cells
|
|
|
|
def _detect_watermark(self, gray: np.ndarray, image_shape: tuple) -> dict[str, Any] | None:
|
|
"""Detect potential watermark regions."""
|
|
h, w = image_shape[:2]
|
|
|
|
# Look for semi-transparent or light text in the center region
|
|
center_region = gray[h // 4 : 3 * h // 4, w // 4 : 3 * w // 4]
|
|
|
|
# Apply adaptive threshold to find light text
|
|
_, binary = cv2.threshold(center_region, 230, 255, cv2.THRESH_BINARY)
|
|
|
|
# Find contours in the center region
|
|
contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
|
|
for contour in contours:
|
|
x, y, cw, ch = cv2.boundingRect(contour)
|
|
area = cw * ch
|
|
# Watermark typically covers a significant portion of the center
|
|
center_area = (w // 2) * (h // 2)
|
|
if area > center_area * 0.1:
|
|
return {
|
|
"x": float(x + w // 4),
|
|
"y": float(y + h // 4),
|
|
"width": float(cw),
|
|
"height": float(ch),
|
|
"detected": True,
|
|
}
|
|
|
|
return None
|