Changes committed
This commit is contained in:
317
docengine/app/services/layout_service.py
Normal file
317
docengine/app/services/layout_service.py
Normal file
@@ -0,0 +1,317 @@
|
||||
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
|
||||
Reference in New Issue
Block a user