historical-ocr / ocr_processing.py
milwright's picture
modularize + nest scripts; reduce technical debt
94e74f0
# Standard library imports
import os
import hashlib
import tempfile
import logging
import time
from datetime import datetime
from pathlib import Path
# Configure logging
logging.basicConfig(level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
# Third-party imports
import streamlit as st
# Local application imports
from structured_ocr import StructuredOCR
# Import from updated utils directory
from utils.image_utils import clean_ocr_result
# Temporarily retain old utils imports until they are fully migrated
from utils import generate_cache_key, timing, format_timestamp, create_descriptive_filename, extract_subject_tags
import preprocessing
from error_handler import handle_ocr_error, check_file_size
from image_segmentation import segment_image_for_ocr, process_segmented_image
@st.cache_data(ttl=24*3600, max_entries=20, show_spinner=False)
def process_file_cached(file_path, file_type, use_vision, file_size_mb, cache_key, preprocessing_options_hash=None, custom_prompt=None):
"""
Cached version of OCR processing to reuse results
Args:
file_path: Path to the file to process
file_type: Type of file (pdf or image)
use_vision: Whether to use vision model
file_size_mb: File size in MB
cache_key: Cache key for the file
preprocessing_options_hash: Hash of preprocessing options
custom_prompt: Custom prompt to use for OCR
Returns:
dict: OCR result
"""
# Initialize OCR processor
processor = StructuredOCR()
# Process the file
with timing(f"OCR processing of {file_type} file"):
result = processor.process_file(
file_path,
file_type=file_type,
use_vision=use_vision,
file_size_mb=file_size_mb,
custom_prompt=custom_prompt
)
return result
def process_file(uploaded_file, use_vision=True, preprocessing_options=None, progress_reporter=None,
pdf_dpi=150, max_pages=3, pdf_rotation=0, custom_prompt=None, perf_mode="Quality",
use_segmentation=False):
"""
Process the uploaded file and return the OCR results
Args:
uploaded_file: The uploaded file to process
use_vision: Whether to use vision model
preprocessing_options: Dictionary of preprocessing options
progress_reporter: ProgressReporter instance for UI updates
pdf_dpi: DPI for PDF conversion
max_pages: Maximum number of pages to process
pdf_rotation: PDF rotation value
custom_prompt: Custom prompt for OCR
perf_mode: Performance mode (Quality or Speed)
Returns:
dict: OCR result
"""
if preprocessing_options is None:
preprocessing_options = {}
# Create a container for progress indicators if not provided
if progress_reporter is None:
from ui.ui_components import ProgressReporter
progress_reporter = ProgressReporter(st.empty()).setup()
# Initialize temporary file paths list
temp_file_paths = []
# Also track temporary files in session state for reliable cleanup
if 'temp_file_paths' not in st.session_state:
st.session_state.temp_file_paths = []
try:
# Check if file size exceeds maximum allowed size
is_valid, file_size_mb, error_message = check_file_size(uploaded_file.getvalue())
if not is_valid:
progress_reporter.complete(success=False)
st.error(error_message)
return {
"file_name": uploaded_file.name,
"topics": ["Document"],
"languages": ["English"],
"error": error_message,
"ocr_contents": {
"error": error_message,
"partial_text": "Document could not be processed due to size limitations."
}
}
# Update progress
progress_reporter.update(10, "Initializing OCR processor...")
# Determine file type from extension
file_ext = Path(uploaded_file.name).suffix.lower()
file_type = "pdf" if file_ext == ".pdf" else "image"
file_bytes = uploaded_file.getvalue()
# For PDFs, we need to handle differently
if file_type == "pdf":
progress_reporter.update(20, "Preparing PDF document...")
# Create a temporary file for processing
temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=file_ext).name
with open(temp_path, 'wb') as f:
f.write(file_bytes)
temp_file_paths.append(temp_path)
# Track temp files in session state for reliable cleanup
if temp_path not in st.session_state.temp_file_paths:
st.session_state.temp_file_paths.append(temp_path)
logger.info(f"Added temp file to session state: {temp_path}")
# Generate cache key
cache_key = generate_cache_key(
file_bytes,
file_type,
use_vision,
preprocessing_options,
pdf_rotation,
custom_prompt
)
# Use the document type information from preprocessing options
doc_type = preprocessing_options.get("document_type", "standard")
modified_custom_prompt = custom_prompt
# Enhance the prompt with document-type specific instructions
# Check for letterhead/marginalia document types with specialized handling
try:
from utils.helpers.letterhead_handler import get_letterhead_prompt, is_likely_letterhead
# Extract text density features if available
features = None
if 'text_density' in preprocessing_options:
features = preprocessing_options['text_density']
# Check if this looks like a letterhead document
if is_likely_letterhead(temp_path, features):
# Get specialized letterhead prompt
letterhead_prompt = get_letterhead_prompt(temp_path, features)
if letterhead_prompt:
logger.info(f"Using specialized letterhead prompt for document")
modified_custom_prompt = letterhead_prompt
# Set document type for tracking
preprocessing_options["document_type"] = "letterhead"
doc_type = "letterhead"
except ImportError:
logger.debug("Letterhead handler not available")
# Add document-type specific instructions based on preprocessing options
if doc_type == "handwritten" and not modified_custom_prompt:
modified_custom_prompt = "This is a handwritten document. Please carefully transcribe all handwritten text, preserving line breaks and original formatting."
elif doc_type == "handwritten" and "handwritten" not in modified_custom_prompt.lower():
modified_custom_prompt += " This is a handwritten document. Please carefully transcribe all handwritten text, preserving line breaks and original formatting."
elif doc_type == "newspaper" and not modified_custom_prompt:
modified_custom_prompt = "This is a newspaper or document with columns. Please extract all text content from each column, maintaining proper reading order."
elif doc_type == "newspaper" and "column" not in modified_custom_prompt.lower() and "newspaper" not in modified_custom_prompt.lower():
modified_custom_prompt += " This appears to be a newspaper or document with columns. Please extract all text content from each column."
elif doc_type == "book" and not modified_custom_prompt:
modified_custom_prompt = "This is a book page. Extract titles, headers, footnotes, and body text, preserving paragraph structure and formatting."
# Update the cache key with the modified prompt
if modified_custom_prompt != custom_prompt:
cache_key = generate_cache_key(
open(temp_path, 'rb').read(),
file_type,
use_vision,
preprocessing_options,
pdf_rotation,
modified_custom_prompt
)
progress_reporter.update(30, "Processing PDF with enhanced OCR...")
# Process with cached function if possible
try:
result = process_file_cached(temp_path, file_type, use_vision, file_size_mb, cache_key,
str(preprocessing_options), modified_custom_prompt)
progress_reporter.update(90, "Finalizing results...")
except Exception as e:
logger.warning(f"Cached processing failed: {str(e)}. Using direct processing.")
progress_reporter.update(60, f"Processing error: {str(e)}. Using enhanced PDF processor...")
# Import the enhanced PDF processor
try:
from utils.pdf_ocr import PDFOCR
# Use our specialized PDF processor
pdf_processor = PDFOCR()
# Process with the enhanced PDF processor
result = pdf_processor.process_pdf(
pdf_path=temp_path,
use_vision=use_vision,
max_pages=max_pages,
custom_prompt=modified_custom_prompt
)
logger.info("PDF successfully processed with enhanced PDF processor")
progress_reporter.update(90, "Finalizing results...")
except ImportError:
logger.warning("Enhanced PDF processor not available. Falling back to standard processing.")
progress_reporter.update(70, "Falling back to standard PDF processing...")
# If enhanced processor is not available, fall back to direct StructuredOCR processing
processor = StructuredOCR()
result = processor.process_file(
file_path=temp_path,
file_type="pdf",
use_vision=use_vision,
custom_prompt=modified_custom_prompt,
file_size_mb=file_size_mb,
max_pages=max_pages
)
progress_reporter.update(90, "Finalizing results...")
else:
# For image files
progress_reporter.update(20, "Preparing image for processing...")
# Apply preprocessing if needed
temp_path, preprocessing_applied = preprocessing.apply_preprocessing_to_file(
file_bytes,
file_ext,
preprocessing_options,
temp_file_paths
)
if preprocessing_applied:
progress_reporter.update(30, "Applied image preprocessing...")
# Apply image segmentation if requested
# This is especially helpful for complex documents with mixed text and images
if use_segmentation:
progress_reporter.update(35, "Applying image segmentation to separate text and image regions...")
try:
# Perform image segmentation with content preservation if requested
preserve_content = preprocessing_options.get("preserve_content", True)
segmentation_results = segment_image_for_ocr(
temp_path,
vision_enabled=use_vision,
preserve_content=preserve_content
)
if segmentation_results['combined_result'] is not None:
# Save the segmented result to a new temporary file
segmented_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg').name
segmentation_results['combined_result'].save(segmented_temp_path)
temp_file_paths.append(segmented_temp_path)
# Check if we have individual region images to process separately
if 'region_images' in segmentation_results and segmentation_results['region_images']:
# Process each region separately for better results
regions_count = len(segmentation_results['region_images'])
logger.info(f"Processing {regions_count} text regions individually")
progress_reporter.update(40, f"Processing {regions_count} text regions separately...")
# Initialize StructuredOCR processor
processor = StructuredOCR()
# Store individual region results
region_results = []
# Process each region individually
for idx, region_info in enumerate(segmentation_results['region_images']):
# Save region image to temp file
region_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix='.jpg').name
region_info['pil_image'].save(region_temp_path)
temp_file_paths.append(region_temp_path)
# Create region-specific prompt
region_prompt = f"This is region {idx+1} of {regions_count} from a segmented document. Extract all visible text precisely, preserving line breaks and structure."
# Process the region
try:
region_result = processor.process_file(
file_path=region_temp_path,
file_type="image",
use_vision=use_vision,
custom_prompt=region_prompt,
file_size_mb=None
)
# Store result with region info
if 'ocr_contents' in region_result and 'raw_text' in region_result['ocr_contents']:
region_results.append({
'text': region_result['ocr_contents']['raw_text'],
'coordinates': region_info['coordinates'],
'order': region_info['order']
})
except Exception as region_err:
logger.warning(f"Error processing region {idx+1}: {str(region_err)}")
# Sort regions by their order for correct reading flow
region_results.sort(key=lambda x: x['order'])
# Import the text utilities for intelligent merging
try:
from utils.text_utils import merge_region_texts
# Use intelligent merging to avoid duplication in overlapped regions
combined_text = merge_region_texts(region_results)
logger.info("Using intelligent text merging for overlapping regions")
except ImportError:
# Fallback to simple joining if import fails
combined_text = "\n\n".join([r['text'] for r in region_results if r['text'].strip()])
logger.warning("Using simple text joining (utils.text_utils not available)")
# Store combined results for later use
preprocessing_options['segmentation_data'] = {
'text_regions_coordinates': segmentation_results.get('text_regions_coordinates', []),
'regions_count': regions_count,
'segmentation_applied': True,
'combined_text': combined_text,
'region_results': region_results
}
logger.info(f"Successfully processed {len(region_results)} text regions")
# Set up the temp path to use the segmented image
temp_path = segmented_temp_path
# IMPORTANT: We've already extracted text from individual regions,
# emphasize their importance in our prompt
if custom_prompt:
# Add strong emphasis on using the already extracted text
custom_prompt += f" IMPORTANT: The document has been segmented into {regions_count} text regions that have been processed individually. The text from these regions should be given HIGHEST PRIORITY and used as the primary source of text for the document. The combined image is provided only as supplementary context."
else:
# Create explicit prompt prioritizing region text
custom_prompt = f"CRITICAL: This document has been preprocessed to highlight {regions_count} text regions that have been individually processed. The text from these regions is the PRIMARY source of content and should be prioritized over any text extracted from the combined image. Use the combined image only for context and layout understanding."
else:
# No individual regions found, use combined result
temp_path = segmented_temp_path
# Enhanced prompt based on segmentation results
regions_count = len(segmentation_results.get('text_regions_coordinates', []))
if custom_prompt:
# Add segmentation info to existing prompt
custom_prompt += f" The document has been segmented and contains approximately {regions_count} text regions that should be carefully extracted. Please focus on extracting all text from these regions."
else:
# Create new prompt focused on text extraction from segmented regions
custom_prompt = f"This document has been preprocessed to highlight {regions_count} text regions. Please carefully extract all text from these highlighted regions, preserving the reading order and structure."
# Store segmentation data in preprocessing options for later use
preprocessing_options['segmentation_data'] = {
'text_regions_coordinates': segmentation_results.get('text_regions_coordinates', []),
'regions_count': regions_count,
'segmentation_applied': True
}
logger.info(f"Image segmentation applied. Found {len(segmentation_results.get('text_regions_coordinates', []))} text regions.")
progress_reporter.update(40, f"Identified {len(segmentation_results.get('text_regions_coordinates', []))} text regions for extraction...")
else:
logger.warning("Image segmentation produced no result, using original image.")
except Exception as seg_error:
logger.warning(f"Image segmentation failed: {str(seg_error)}. Continuing with standard processing.")
# Generate cache key
cache_key = generate_cache_key(
open(temp_path, 'rb').read(),
file_type,
use_vision,
preprocessing_options,
0, # No rotation for images (handled in preprocessing)
custom_prompt
)
# Process the file using cached function if possible
progress_reporter.update(50, "Processing document with OCR...")
try:
# Use the document type from preprocessing options
doc_type = preprocessing_options.get("document_type", "standard")
modified_custom_prompt = custom_prompt
# Check for letterhead/marginalia document types with specialized handling
try:
from utils.helpers.letterhead_handler import get_letterhead_prompt, is_likely_letterhead
# Extract text density features if available
features = None
if 'text_density' in preprocessing_options:
features = preprocessing_options['text_density']
# Check if this looks like a letterhead document
if is_likely_letterhead(temp_path, features):
# Get specialized letterhead prompt
letterhead_prompt = get_letterhead_prompt(temp_path, features)
if letterhead_prompt:
logger.info(f"Using specialized letterhead prompt for document")
modified_custom_prompt = letterhead_prompt
# Set document type for tracking
preprocessing_options["document_type"] = "letterhead"
doc_type = "letterhead"
except ImportError:
logger.debug("Letterhead handler not available")
# Add document-type specific instructions based on preprocessing options
if doc_type == "handwritten" and not modified_custom_prompt:
modified_custom_prompt = "This is a handwritten document. Please carefully transcribe all handwritten text, preserving line breaks and original formatting."
elif doc_type == "handwritten" and "handwritten" not in modified_custom_prompt.lower():
modified_custom_prompt += " This is a handwritten document. Please carefully transcribe all handwritten text, preserving line breaks and original formatting."
elif doc_type == "newspaper" and not modified_custom_prompt:
modified_custom_prompt = "This is a newspaper or document with columns. Please extract all text content from each column, maintaining proper reading order."
elif doc_type == "newspaper" and "column" not in modified_custom_prompt.lower() and "newspaper" not in modified_custom_prompt.lower():
modified_custom_prompt += " This appears to be a newspaper or document with columns. Please extract all text content from each column."
elif doc_type == "book" and not modified_custom_prompt:
modified_custom_prompt = "This is a book page. Extract titles, headers, footnotes, and body text, preserving paragraph structure and formatting."
# Update the cache key with the modified prompt
if modified_custom_prompt != custom_prompt:
cache_key = generate_cache_key(
open(temp_path, 'rb').read(),
file_type,
use_vision,
preprocessing_options,
0,
modified_custom_prompt
)
result = process_file_cached(temp_path, file_type, use_vision, file_size_mb, cache_key, str(preprocessing_options), modified_custom_prompt)
progress_reporter.update(80, "Analyzing document structure...")
progress_reporter.update(90, "Finalizing results...")
except Exception as e:
logger.warning(f"Cached processing failed: {str(e)}. Retrying with direct processing.")
progress_reporter.update(60, f"Processing error: {str(e)}. Retrying...")
# If caching fails, process directly
processor = StructuredOCR()
# Apply performance mode settings
if perf_mode == "Speed":
# Use simpler processing for speed
pass # Any speed optimizations would be handled by the StructuredOCR class
# Use the document type from preprocessing options
doc_type = preprocessing_options.get("document_type", "standard")
modified_custom_prompt = custom_prompt
# Check for letterhead/marginalia document types with specialized handling
try:
from utils.helpers.letterhead_handler import get_letterhead_prompt, is_likely_letterhead
# Extract text density features if available
features = None
if 'text_density' in preprocessing_options:
features = preprocessing_options['text_density']
# Check if this looks like a letterhead document
if is_likely_letterhead(temp_path, features):
# Get specialized letterhead prompt
letterhead_prompt = get_letterhead_prompt(temp_path, features)
if letterhead_prompt:
logger.info(f"Using specialized letterhead prompt for document")
modified_custom_prompt = letterhead_prompt
# Set document type for tracking
preprocessing_options["document_type"] = "letterhead"
doc_type = "letterhead"
except ImportError:
logger.debug("Letterhead handler not available")
# Add document-type specific instructions based on preprocessing options
if doc_type == "handwritten" and not modified_custom_prompt:
modified_custom_prompt = "This is a handwritten document. Please carefully transcribe all handwritten text, preserving line breaks and original formatting."
elif doc_type == "handwritten" and "handwritten" not in modified_custom_prompt.lower():
modified_custom_prompt += " This is a handwritten document. Please carefully transcribe all handwritten text, preserving line breaks and original formatting."
elif doc_type == "newspaper" and not modified_custom_prompt:
modified_custom_prompt = "This is a newspaper or document with columns. Please extract all text content from each column, maintaining proper reading order."
elif doc_type == "newspaper" and "column" not in modified_custom_prompt.lower() and "newspaper" not in modified_custom_prompt.lower():
modified_custom_prompt += " This appears to be a newspaper or document with columns. Please extract all text content from each column."
elif doc_type == "book" and not modified_custom_prompt:
modified_custom_prompt = "This is a book page. Extract titles, headers, footnotes, and body text, preserving paragraph structure and formatting."
result = processor.process_file(
file_path=temp_path,
file_type=file_type,
use_vision=use_vision,
custom_prompt=modified_custom_prompt,
file_size_mb=file_size_mb
)
progress_reporter.update(90, "Finalizing results...")
# Add additional metadata to result
result = process_result(result, uploaded_file, preprocessing_options)
# Make sure file_type is explicitly set for PDFs
if file_type == "pdf":
result['file_type'] = "pdf"
# Check for duplicated text patterns that indicate handwritten text issues
try:
from utils.helpers.ocr_text_repair import detect_duplicate_text_issues, get_enhanced_preprocessing_options, get_handwritten_specific_prompt, clean_duplicated_text
# Check OCR output for duplication issues
if result and 'ocr_contents' in result and 'raw_text' in result['ocr_contents']:
ocr_text = result['ocr_contents']['raw_text']
has_duplication, duplication_details = detect_duplicate_text_issues(ocr_text)
# If we detect significant duplication in the output
if has_duplication and duplication_details.get('duplication_rate', 0) > 0.1:
logger.info(f"Detected text duplication issues. Reprocessing as handwritten document with enhanced settings...")
progress_reporter.update(75, "Detected duplication issues. Reprocessing with enhanced settings...")
# Save original result before reprocessing
original_result = result
# Get enhanced preprocessing options for handwritten text
enhanced_options = get_enhanced_preprocessing_options(preprocessing_options)
# Reprocess with enhanced settings and specialized prompt
handwritten_prompt = get_handwritten_specific_prompt(custom_prompt)
# Process the image with the enhanced settings
try:
# Apply enhanced preprocessing to the original image
enhanced_temp_path, _ = preprocessing.apply_preprocessing_to_file(
open(temp_path, 'rb').read(),
Path(temp_path).suffix.lower(),
enhanced_options,
temp_file_paths
)
# Process with enhanced settings
processor = StructuredOCR()
enhanced_result = processor.process_file(
file_path=enhanced_temp_path,
file_type="image",
use_vision=use_vision,
custom_prompt=handwritten_prompt,
file_size_mb=file_size_mb
)
# Check if the enhanced result is better (less duplication)
if 'ocr_contents' in enhanced_result and 'raw_text' in enhanced_result['ocr_contents']:
enhanced_text = enhanced_result['ocr_contents']['raw_text']
_, enhanced_issues = detect_duplicate_text_issues(enhanced_text)
# Use the enhanced result if it's better
if enhanced_issues.get('duplication_rate', 1.0) < duplication_details.get('duplication_rate', 1.0):
logger.info("Enhanced processing improved OCR quality. Using enhanced result.")
result = enhanced_result
# Preserve document type and preprocessing info
result['document_type'] = 'handwritten'
result['preprocessing'] = enhanced_options
else:
# If enhancement didn't help, clean up the original result
logger.info("Enhanced processing did not improve OCR quality. Cleaning original result.")
result = original_result
# Clean up duplication in the text
if 'ocr_contents' in result and 'raw_text' in result['ocr_contents']:
result['ocr_contents']['raw_text'] = clean_duplicated_text(result['ocr_contents']['raw_text'])
else:
# Fallback to original with cleaning
logger.info("Enhanced processing failed. Cleaning original result.")
result = original_result
# Clean up duplication in the text
if 'ocr_contents' in result and 'raw_text' in result['ocr_contents']:
result['ocr_contents']['raw_text'] = clean_duplicated_text(result['ocr_contents']['raw_text'])
except Exception as enh_error:
logger.warning(f"Enhanced processing failed: {str(enh_error)}. Using cleaned original.")
# Fallback to original with cleaning
result = original_result
# Clean up duplication in the text
if 'ocr_contents' in result and 'raw_text' in result['ocr_contents']:
result['ocr_contents']['raw_text'] = clean_duplicated_text(result['ocr_contents']['raw_text'])
except ImportError:
logger.debug("OCR text repair module not available")
# 🔧 ALWAYS normalize result before returning
result = clean_ocr_result(
result,
use_segmentation=use_segmentation,
vision_enabled=use_vision,
preprocessing_options=preprocessing_options
)
# Complete progress
progress_reporter.complete()
return result
except Exception as e:
# Handle errors
error_message = handle_ocr_error(e, progress_reporter)
# Return error result
return {
"file_name": uploaded_file.name,
"topics": ["Document"],
"languages": ["English"],
"error": error_message,
"ocr_contents": {
"error": f"Failed to process file: {error_message}",
"partial_text": "Document could not be processed due to an error."
}
}
finally:
# Clean up temporary files
for temp_path in temp_file_paths:
try:
if os.path.exists(temp_path):
os.unlink(temp_path)
logger.info(f"Removed temporary file: {temp_path}")
except Exception as e:
logger.warning(f"Failed to remove temporary file {temp_path}: {str(e)}")
def process_result(result, uploaded_file, preprocessing_options=None):
"""
Process OCR result to add metadata, tags, etc.
Args:
result: OCR result dictionary
uploaded_file: The uploaded file
preprocessing_options: Dictionary of preprocessing options
Returns:
dict: Processed OCR result
"""
# Add timestamp
result['timestamp'] = format_timestamp()
# Add processing time if not already present
if 'processing_time' not in result:
result['processing_time'] = 0.0
# Generate descriptive filename
file_ext = Path(uploaded_file.name).suffix.lower()
result['descriptive_file_name'] = create_descriptive_filename(
uploaded_file.name,
result,
file_ext,
preprocessing_options
)
# Extract raw text from OCR contents for tag extraction without duplicating content
raw_text = ""
if 'ocr_contents' in result:
# Try fields in order of preference
for field in ["raw_text", "content", "text", "transcript", "main_text"]:
if field in result['ocr_contents'] and result['ocr_contents'][field]:
raw_text = result['ocr_contents'][field]
break
# Extract subject tags if not already present or enhance existing ones
if 'topics' not in result or not result['topics']:
result['topics'] = extract_subject_tags(result, raw_text, preprocessing_options)
return result