214 lines
No EOL
7.1 KiB
Python
214 lines
No EOL
7.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Common utilities for DocTags processing
|
|
"""
|
|
|
|
import os
|
|
import subprocess
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple, Dict, List
|
|
import pdf2image
|
|
from pdf2image.pdf2image import pdfinfo_from_path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Configuration constants
|
|
DEFAULT_DPI = 200
|
|
DEFAULT_GRID_SIZE = 500
|
|
MAX_WIDTH = 1200
|
|
RESULTS_DIR_NAME = "results"
|
|
|
|
def get_project_root() -> Path:
|
|
"""Get the project root directory."""
|
|
# If running from backend/page_treatment/, go up to root
|
|
current_file = Path(__file__)
|
|
if current_file.parent.name == 'page_treatment':
|
|
return current_file.parent.parent.parent
|
|
elif current_file.parent.name == 'backend':
|
|
return current_file.parent.parent
|
|
else:
|
|
return Path.cwd()
|
|
|
|
def ensure_results_folder(custom_path: Optional[str] = None) -> Path:
|
|
"""Create and return the results folder path."""
|
|
if custom_path:
|
|
results_dir = Path(custom_path)
|
|
else:
|
|
results_dir = get_project_root() / RESULTS_DIR_NAME
|
|
|
|
if not results_dir.exists():
|
|
results_dir.mkdir(parents=True)
|
|
logger.info(f"Created results directory: {results_dir}")
|
|
|
|
return results_dir
|
|
|
|
def count_pdf_pages(pdf_path: str) -> int:
|
|
"""Count the number of pages in a PDF file."""
|
|
if not os.path.exists(pdf_path):
|
|
logger.error(f"PDF file not found: {pdf_path}")
|
|
return 0
|
|
|
|
try:
|
|
info = pdfinfo_from_path(pdf_path)
|
|
return info["Pages"]
|
|
except Exception as e:
|
|
logger.warning(f"pdfinfo failed: {e}, trying fallback method")
|
|
try:
|
|
# Fallback: convert first page to check
|
|
images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=1, last_page=1)
|
|
if not images:
|
|
return 0
|
|
|
|
# Binary search for last page
|
|
low, high = 1, 1000
|
|
while low < high:
|
|
mid = (low + high + 1) // 2
|
|
try:
|
|
images = pdf2image.convert_from_path(pdf_path, dpi=72, first_page=mid, last_page=mid)
|
|
if images:
|
|
low = mid
|
|
else:
|
|
high = mid - 1
|
|
except:
|
|
high = mid - 1
|
|
|
|
return low
|
|
except Exception as e2:
|
|
logger.error(f"Error counting PDF pages: {e2}")
|
|
return 0
|
|
|
|
def load_pdf_page(pdf_path: str, page_num: int = 1, dpi: int = DEFAULT_DPI) -> Optional[object]:
|
|
"""Load a specific page from PDF as an image."""
|
|
if not os.path.exists(pdf_path):
|
|
raise FileNotFoundError(f"PDF file not found: {pdf_path}")
|
|
|
|
logger.info(f"Converting PDF page {page_num} to image (DPI: {dpi})...")
|
|
try:
|
|
pdf_images = pdf2image.convert_from_path(
|
|
pdf_path,
|
|
dpi=dpi,
|
|
first_page=page_num,
|
|
last_page=page_num
|
|
)
|
|
if not pdf_images:
|
|
raise Exception(f"Could not extract page {page_num} from PDF")
|
|
return pdf_images[0]
|
|
except Exception as e:
|
|
raise Exception(f"Error converting PDF to image: {e}")
|
|
|
|
def normalize_coordinates(elements: List[Dict], image_width: int, image_height: int,
|
|
grid_size: int = DEFAULT_GRID_SIZE) -> List[Dict]:
|
|
"""
|
|
Normalize coordinates from DocTags grid to actual image dimensions.
|
|
|
|
Args:
|
|
elements: List of elements with x1, y1, x2, y2 coordinates
|
|
image_width: Width of the image in pixels
|
|
image_height: Height of the image in pixels
|
|
grid_size: The grid size used in DocTags (default 500)
|
|
|
|
Returns:
|
|
List of elements with normalized coordinates
|
|
"""
|
|
normalized = []
|
|
for element in elements:
|
|
new_element = element.copy()
|
|
new_element['x1'] = int(element['x1'] * image_width / grid_size)
|
|
new_element['y1'] = int(element['y1'] * image_height / grid_size)
|
|
new_element['x2'] = int(element['x2'] * image_width / grid_size)
|
|
new_element['y2'] = int(element['y2'] * image_height / grid_size)
|
|
normalized.append(new_element)
|
|
return normalized
|
|
|
|
def auto_adjust_coordinates(elements: List[Dict], image_width: int, image_height: int) -> List[Dict]:
|
|
"""
|
|
Automatically adjust coordinates based on image dimensions.
|
|
"""
|
|
if not elements:
|
|
return elements
|
|
|
|
# Find maximum coordinates
|
|
max_x = max([el['x2'] for el in elements])
|
|
max_y = max([el['y2'] for el in elements])
|
|
|
|
# Check if coordinates are in normalized grid (0-500 range)
|
|
if max_x <= DEFAULT_GRID_SIZE and max_y <= DEFAULT_GRID_SIZE:
|
|
logger.info(f"Detected normalized coordinates (0-{DEFAULT_GRID_SIZE} grid)")
|
|
return normalize_coordinates(elements, image_width, image_height)
|
|
|
|
# Calculate scaling factors
|
|
x_scale = calculate_scale_factor(max_x, image_width)
|
|
y_scale = calculate_scale_factor(max_y, image_height)
|
|
|
|
# Apply scaling
|
|
adjusted = []
|
|
for el in elements:
|
|
adjusted_el = el.copy()
|
|
adjusted_el['x1'] = int(el['x1'] * x_scale)
|
|
adjusted_el['y1'] = int(el['y1'] * y_scale)
|
|
adjusted_el['x2'] = int(el['x2'] * x_scale)
|
|
adjusted_el['y2'] = int(el['y2'] * y_scale)
|
|
adjusted.append(adjusted_el)
|
|
|
|
logger.info(f"Applied auto-scaling: X={x_scale:.3f}, Y={y_scale:.3f}")
|
|
return adjusted
|
|
|
|
def calculate_scale_factor(max_coord: float, image_size: float) -> float:
|
|
"""Calculate appropriate scaling factor."""
|
|
if max_coord <= 0:
|
|
return 1.0
|
|
|
|
# If coordinates are way off, apply aggressive scaling
|
|
if max_coord > image_size * 5 or max_coord < image_size / 5:
|
|
return image_size / max_coord
|
|
|
|
# Otherwise, apply conservative scaling
|
|
if max_coord > image_size:
|
|
return min(image_size / max_coord, 1.0)
|
|
else:
|
|
return max(image_size / max_coord, 0.5)
|
|
|
|
# In backend/utils.py, make sure this function exists:
|
|
def run_command_with_timeout(command: str, timeout: int = 300, input_text: str = "n\n") -> Tuple[bool, str, str]:
|
|
"""
|
|
Run a command with timeout and return success, stdout, stderr.
|
|
"""
|
|
try:
|
|
process = subprocess.Popen(
|
|
command,
|
|
shell=True,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
universal_newlines=True
|
|
)
|
|
|
|
stdout, stderr = process.communicate(input=input_text, timeout=timeout)
|
|
success = process.returncode == 0
|
|
|
|
return success, stdout, stderr
|
|
|
|
except subprocess.TimeoutExpired:
|
|
process.kill()
|
|
return False, "", "Command timed out"
|
|
except Exception as e:
|
|
return False, "", str(e)
|
|
|
|
def format_duration(seconds: float) -> str:
|
|
"""Format duration in seconds to human readable format."""
|
|
hours = int(seconds // 3600)
|
|
minutes = int((seconds % 3600) // 60)
|
|
secs = int(seconds % 60)
|
|
|
|
if hours > 0:
|
|
return f"{hours}:{minutes:02d}:{secs:02d}"
|
|
else:
|
|
return f"{minutes}:{secs:02d}"
|
|
|
|
def validate_coordinates(x1: int, y1: int, x2: int, y2: int,
|
|
width: int, height: int) -> bool:
|
|
"""Validate that coordinates are within bounds."""
|
|
return (0 <= x1 < x2 <= width and
|
|
0 <= y1 < y2 <= height) |