document_redaction / tools /custom_image_analyser_engine.py
seanpedrickcase's picture
General code changes and reformatting to address code vulnerabilities highlighted by codeQL scan, and black/ruff repplied to code. Fixes/optimisation of Github Actions
f957846
raw
history blame
71.7 kB
import copy
import re
import time
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union
import botocore
import cv2
import gradio as gr
import numpy as np
import pytesseract
from pdfminer.layout import LTChar
from PIL import Image
from presidio_analyzer import AnalyzerEngine, RecognizerResult
from tools.config import (
AWS_PII_OPTION,
DEFAULT_LANGUAGE,
LOCAL_PII_OPTION,
PREPROCESS_LOCAL_OCR_IMAGES,
)
from tools.helper_functions import clean_unicode_text
from tools.load_spacy_model_custom_recognisers import custom_entities
from tools.presidio_analyzer_custom import recognizer_result_from_dict
if PREPROCESS_LOCAL_OCR_IMAGES == "True":
PREPROCESS_LOCAL_OCR_IMAGES = True
else:
PREPROCESS_LOCAL_OCR_IMAGES = False
try:
from paddleocr import PaddleOCR
except ImportError:
PaddleOCR = None
# --- Language utilities ---
def _normalize_lang(language: str) -> str:
return language.strip().lower().replace("-", "_") if language else "en"
def _tesseract_lang_code(language: str) -> str:
"""Map a user language input to a Tesseract traineddata code."""
lang = _normalize_lang(language)
mapping = {
# Common
"en": "eng",
"eng": "eng",
"fr": "fra",
"fre": "fra",
"fra": "fra",
"de": "deu",
"ger": "deu",
"deu": "deu",
"es": "spa",
"spa": "spa",
"it": "ita",
"ita": "ita",
"nl": "nld",
"dut": "nld",
"nld": "nld",
"pt": "por",
"por": "por",
"ru": "rus",
"rus": "rus",
"ar": "ara",
"ara": "ara",
# Nordics
"sv": "swe",
"swe": "swe",
"no": "nor",
"nb": "nor",
"nn": "nor",
"nor": "nor",
"fi": "fin",
"fin": "fin",
"da": "dan",
"dan": "dan",
# Eastern/Central
"pl": "pol",
"pol": "pol",
"cs": "ces",
"cz": "ces",
"ces": "ces",
"hu": "hun",
"hun": "hun",
"ro": "ron",
"rum": "ron",
"ron": "ron",
"bg": "bul",
"bul": "bul",
"el": "ell",
"gre": "ell",
"ell": "ell",
# Asian
"ja": "jpn",
"jp": "jpn",
"jpn": "jpn",
"zh": "chi_sim",
"zh_cn": "chi_sim",
"zh_hans": "chi_sim",
"chi_sim": "chi_sim",
"zh_tw": "chi_tra",
"zh_hk": "chi_tra",
"zh_tr": "chi_tra",
"chi_tra": "chi_tra",
"hi": "hin",
"hin": "hin",
"bn": "ben",
"ben": "ben",
"ur": "urd",
"urd": "urd",
"fa": "fas",
"per": "fas",
"fas": "fas",
}
return mapping.get(lang, "eng")
def _paddle_lang_code(language: str) -> str:
"""Map a user language input to a PaddleOCR language code.
PaddleOCR supports codes like: 'en', 'ch', 'chinese_cht', 'korean', 'japan', 'german', 'fr', 'it', 'es',
as well as script packs like 'arabic', 'cyrillic', 'latin'.
"""
lang = _normalize_lang(language)
mapping = {
"en": "en",
"fr": "fr",
"de": "german",
"es": "es",
"it": "it",
"pt": "pt",
"nl": "nl",
"ru": "cyrillic", # Russian is covered by cyrillic models
"uk": "cyrillic",
"bg": "cyrillic",
"sr": "cyrillic",
"ar": "arabic",
"tr": "tr",
"fa": "arabic", # fallback to arabic script pack
"zh": "ch",
"zh_cn": "ch",
"zh_tw": "chinese_cht",
"zh_hk": "chinese_cht",
"ja": "japan",
"jp": "japan",
"ko": "korean",
"hi": "latin", # fallback; dedicated Hindi not always available
}
return mapping.get(lang, "en")
@dataclass
class OCRResult:
text: str
left: int
top: int
width: int
height: int
conf: float = None
line: int = None
@dataclass
class CustomImageRecognizerResult:
entity_type: str
start: int
end: int
score: float
left: int
top: int
width: int
height: int
text: str
class ImagePreprocessor:
"""ImagePreprocessor class. Parent class for image preprocessing objects."""
def __init__(self, use_greyscale: bool = True) -> None:
self.use_greyscale = use_greyscale
def preprocess_image(self, image: Image.Image) -> Tuple[Image.Image, dict]:
return image, {}
def convert_image_to_array(self, image: Image.Image) -> np.ndarray:
if isinstance(image, np.ndarray):
img = image
else:
if self.use_greyscale:
image = image.convert("L")
img = np.asarray(image)
return img
@staticmethod
def _get_bg_color(
image: np.ndarray, is_greyscale: bool, invert: bool = False
) -> Union[int, Tuple[int, int, int]]:
# Note: Modified to expect numpy array for bincount
if invert:
image = 255 - image # Simple inversion for greyscale numpy array
if is_greyscale:
bg_color = int(np.bincount(image.flatten()).argmax())
else:
# This part would need more complex logic for color numpy arrays
# For this pipeline, we only use greyscale, so it's fine.
# A simple alternative:
from scipy import stats
bg_color = tuple(stats.mode(image.reshape(-1, 3), axis=0)[0][0])
return bg_color
@staticmethod
def _get_image_contrast(image: np.ndarray) -> Tuple[float, float]:
contrast = np.std(image)
mean_intensity = np.mean(image)
return contrast, mean_intensity
class BilateralFilter(ImagePreprocessor):
"""Applies bilateral filtering."""
def __init__(
self, diameter: int = 9, sigma_color: int = 75, sigma_space: int = 75
) -> None:
super().__init__(use_greyscale=True)
self.diameter = diameter
self.sigma_color = sigma_color
self.sigma_space = sigma_space
def preprocess_image(self, image: np.ndarray) -> Tuple[np.ndarray, dict]:
# Modified to accept and return numpy array for consistency in the pipeline
filtered_image = cv2.bilateralFilter(
image, self.diameter, self.sigma_color, self.sigma_space
)
metadata = {
"diameter": self.diameter,
"sigma_color": self.sigma_color,
"sigma_space": self.sigma_space,
}
return filtered_image, metadata
class SegmentedAdaptiveThreshold(ImagePreprocessor):
"""Applies adaptive thresholding."""
def __init__(
self,
block_size: int = 21,
contrast_threshold: int = 40,
c_low_contrast: int = 5,
c_high_contrast: int = 10,
bg_threshold: int = 127,
) -> None:
super().__init__(use_greyscale=True)
self.block_size = (
block_size if block_size % 2 == 1 else block_size + 1
) # Ensure odd
self.c_low_contrast = c_low_contrast
self.c_high_contrast = c_high_contrast
self.bg_threshold = bg_threshold
self.contrast_threshold = contrast_threshold
def preprocess_image(self, image: np.ndarray) -> Tuple[np.ndarray, dict]:
# Modified to accept and return numpy array
background_color = self._get_bg_color(image, True)
contrast, _ = self._get_image_contrast(image)
c = (
self.c_low_contrast
if contrast <= self.contrast_threshold
else self.c_high_contrast
)
if background_color < self.bg_threshold: # Dark background, light text
adaptive_threshold_image = cv2.adaptiveThreshold(
image,
255,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY_INV,
self.block_size,
-c,
)
else: # Light background, dark text
adaptive_threshold_image = cv2.adaptiveThreshold(
image,
255,
cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY,
self.block_size,
c,
)
metadata = {"C": c, "background_color": background_color, "contrast": contrast}
return adaptive_threshold_image, metadata
class ImageRescaling(ImagePreprocessor):
"""Rescales images based on their size."""
def __init__(self, target_dpi: int = 300, assumed_input_dpi: int = 96) -> None:
super().__init__(use_greyscale=True)
self.target_dpi = target_dpi
self.assumed_input_dpi = assumed_input_dpi
def preprocess_image(self, image: np.ndarray) -> Tuple[np.ndarray, dict]:
# Modified to accept and return numpy array
scale_factor = self.target_dpi / self.assumed_input_dpi
metadata = {"scale_factor": 1.0}
if scale_factor != 1.0:
width = int(image.shape[1] * scale_factor)
height = int(image.shape[0] * scale_factor)
dimensions = (width, height)
# Use better interpolation for upscaling vs downscaling
interpolation = cv2.INTER_CUBIC if scale_factor > 1.0 else cv2.INTER_AREA
rescaled_image = cv2.resize(image, dimensions, interpolation=interpolation)
metadata["scale_factor"] = scale_factor
return rescaled_image, metadata
return image, metadata
class ContrastSegmentedImageEnhancer(ImagePreprocessor):
"""Class containing all logic to perform contrastive segmentation."""
def __init__(
self,
bilateral_filter: Optional[BilateralFilter] = None,
adaptive_threshold: Optional[SegmentedAdaptiveThreshold] = None,
image_rescaling: Optional[ImageRescaling] = None,
low_contrast_threshold: int = 40,
) -> None:
super().__init__(use_greyscale=True)
self.bilateral_filter = bilateral_filter or BilateralFilter()
self.adaptive_threshold = adaptive_threshold or SegmentedAdaptiveThreshold()
self.image_rescaling = image_rescaling or ImageRescaling()
self.low_contrast_threshold = low_contrast_threshold
def _improve_contrast(self, image: np.ndarray) -> Tuple[np.ndarray, str, str]:
contrast, mean_intensity = self._get_image_contrast(image)
if contrast <= self.low_contrast_threshold:
# Using CLAHE as a generally more robust alternative
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
adjusted_image = clahe.apply(image)
adjusted_contrast, _ = self._get_image_contrast(adjusted_image)
else:
adjusted_image = image
adjusted_contrast = contrast
return adjusted_image, contrast, adjusted_contrast
def preprocess_image(
self, image: Image.Image, perform_binarization: bool = False
) -> Tuple[Image.Image, dict]:
"""
A corrected, logical pipeline for OCR preprocessing.
Order: Greyscale -> Rescale -> Denoise -> Enhance Contrast -> Binarize
I have found that binarization is not always helpful with Tesseract, and can sometimes degrade results. So it is off by default.
"""
# 1. Convert to greyscale NumPy array
image_np = self.convert_image_to_array(image)
# 2. Rescale image to optimal DPI (while still greyscale)
rescaled_image_np, scale_metadata = self.image_rescaling.preprocess_image(
image_np
)
# 3. Apply bilateral filtering for noise reduction
filtered_image_np, _ = self.bilateral_filter.preprocess_image(rescaled_image_np)
# 4. Improve contrast
adjusted_image_np, _, _ = self._improve_contrast(filtered_image_np)
# 5. Adaptive Thresholding (Binarization) - This is the final step
if perform_binarization:
final_image_np, threshold_metadata = (
self.adaptive_threshold.preprocess_image(adjusted_image_np)
)
else:
final_image_np = adjusted_image_np
threshold_metadata = {}
# Combine metadata
final_metadata = {**scale_metadata, **threshold_metadata}
# Convert final numpy array back to PIL Image for return
return Image.fromarray(final_image_np), final_metadata
def rescale_ocr_data(ocr_data, scale_factor: float):
# We loop from 0 to the number of detected words.
num_boxes = len(ocr_data["text"])
for i in range(num_boxes):
# We only want to process actual words, not empty boxes Tesseract might find
if int(ocr_data["conf"][i]) > -1: # -1 confidence is for structural elements
# Get coordinates from the processed image using the index 'i'
x_proc = ocr_data["left"][i]
y_proc = ocr_data["top"][i]
w_proc = ocr_data["width"][i]
h_proc = ocr_data["height"][i]
# Apply the inverse transformation (division)
x_orig = int(x_proc / scale_factor)
y_orig = int(y_proc / scale_factor)
w_orig = int(w_proc / scale_factor)
h_orig = int(h_proc / scale_factor)
# --- THE MAPPING STEP ---
# Update the dictionary values in-place using the same index 'i'
ocr_data["left"][i] = x_orig
ocr_data["top"][i] = y_orig
ocr_data["width"][i] = w_orig
ocr_data["height"][i] = h_orig
return ocr_data
def filter_entities_for_language(
entities: List[str], valid_language_entities: List[str], language: str
) -> List[str]:
if not valid_language_entities:
print(f"No valid entities supported for language: {language}")
# raise Warning(f"No valid entities supported for language: {language}")
if not entities:
print(f"No entities provided for language: {language}")
# raise Warning(f"No entities provided for language: {language}")
# print("entities:", entities)
# print("valid_language_entities:", valid_language_entities)
# print("language:", language)
filtered_entities = [
entity for entity in entities if entity in valid_language_entities
]
if not filtered_entities:
print(f"No relevant entities supported for language: {language}")
# raise Warning(f"No relevant entities supported for language: {language}")
if language != "en":
gr.Info(
f"Using {str(filtered_entities)} entities for local model analysis for language: {language}"
)
return filtered_entities
class CustomImageAnalyzerEngine:
def __init__(
self,
analyzer_engine: Optional[AnalyzerEngine] = None,
ocr_engine: str = "tesseract",
tesseract_config: Optional[str] = None,
paddle_kwargs: Optional[Dict[str, Any]] = None,
image_preprocessor: Optional[ImagePreprocessor] = None,
language: Optional[str] = DEFAULT_LANGUAGE,
):
"""
Initializes the CustomImageAnalyzerEngine.
:param ocr_engine: The OCR engine to use ("tesseract", "hybrid", or "paddle").
:param analyzer_engine: The Presidio AnalyzerEngine instance.
:param tesseract_config: Configuration string for Tesseract.
:param paddle_kwargs: Dictionary of keyword arguments for PaddleOCR constructor.
:param image_preprocessor: Optional image preprocessor.
:param language: Preferred OCR language (e.g., "en", "fr", "de"). Defaults to DEFAULT_LANGUAGE.
"""
if ocr_engine not in ["tesseract", "paddle", "hybrid"]:
raise ValueError(
"ocr_engine must be either 'tesseract', 'hybrid', or 'paddle'"
)
self.ocr_engine = ocr_engine
# Language setup
self.language = language or DEFAULT_LANGUAGE or "en"
self.tesseract_lang = _tesseract_lang_code(self.language)
self.paddle_lang = _paddle_lang_code(self.language)
if self.ocr_engine == "paddle" or self.ocr_engine == "hybrid":
if PaddleOCR is None:
raise ImportError(
"paddleocr is not installed. Please run 'pip install paddleocr paddlepaddle'"
)
# Default paddle configuration if none provided
if paddle_kwargs is None:
paddle_kwargs = {
"use_textline_orientation": True,
"lang": self.paddle_lang,
}
else:
# Enforce language if not explicitly provided
paddle_kwargs.setdefault("lang", self.paddle_lang)
self.paddle_ocr = PaddleOCR(**paddle_kwargs)
if not analyzer_engine:
analyzer_engine = AnalyzerEngine()
self.analyzer_engine = analyzer_engine
self.tesseract_config = tesseract_config or "--oem 3 --psm 11"
if not image_preprocessor:
image_preprocessor = ContrastSegmentedImageEnhancer()
self.image_preprocessor = image_preprocessor
def _sanitize_filename(self, text: str, max_length: int = 20) -> str:
"""
Sanitizes text for use in filenames by removing invalid characters and limiting length.
:param text: The text to sanitize
:param max_length: Maximum length of the sanitized text
:return: Sanitized text safe for filenames
"""
# Remove or replace invalid filename characters
# Windows: < > : " | ? * \ /
# Unix: / (forward slash)
from tools.secure_regex_utils import safe_sanitize_text
sanitized = safe_sanitize_text(text)
# Remove leading/trailing underscores and spaces
sanitized = sanitized.strip("_ ")
# If empty after sanitization, use a default value
if not sanitized:
sanitized = "text"
# Limit to max_length characters
if len(sanitized) > max_length:
sanitized = sanitized[:max_length]
# Ensure we don't end with an underscore if we cut in the middle
sanitized = sanitized.rstrip("_")
return sanitized
def _convert_paddle_to_tesseract_format(
self, paddle_results: List[Any]
) -> Dict[str, List]:
"""Converts PaddleOCR result format to Tesseract's dictionary format. NOTE: This attempts to create word-level bounding boxes by estimating the distance between characters in sentence-level text output. This is currently quite inaccurate, and word-level bounding boxes should not be relied upon."""
output = {
"text": [],
"left": [],
"top": [],
"width": [],
"height": [],
"conf": [],
}
# paddle_results is now a list of dictionaries with detailed information
if not paddle_results:
return output
for page_result in paddle_results:
# Extract text recognition results from the new format
rec_texts = page_result.get("rec_texts", [])
rec_scores = page_result.get("rec_scores", [])
rec_polys = page_result.get("rec_polys", [])
for line_text, line_confidence, bounding_box in zip(
rec_texts, rec_scores, rec_polys
):
# bounding_box is now a numpy array with shape (4, 2)
# Convert to list of coordinates if it's a numpy array
if hasattr(bounding_box, "tolist"):
box = bounding_box.tolist()
else:
box = bounding_box
# box is [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
x_coords = [p[0] for p in box]
y_coords = [p[1] for p in box]
line_left = int(min(x_coords))
line_top = int(min(y_coords))
line_width = int(max(x_coords) - line_left)
line_height = int(max(y_coords) - line_top)
# line_y_center = (max(y_coords) + min(y_coords)) / 2
# 2. Split the line into words
words = line_text.split()
if not words:
continue
# 3. Estimate bounding box for each word
total_chars = len(line_text)
# Avoid division by zero for empty lines
avg_char_width = line_width / total_chars if total_chars > 0 else 0
current_char_offset = 0
for word in words:
word_width = int(len(word) * avg_char_width)
word_left = line_left + int(current_char_offset * avg_char_width)
output["text"].append(word)
output["left"].append(word_left)
output["top"].append(line_top)
output["width"].append(word_width)
output["height"].append(line_height)
# Use the line's confidence for each word derived from it
output["conf"].append(int(line_confidence * 100))
# Update offset for the next word (add word length + 1 for the space)
current_char_offset += len(word) + 1
return output
def _perform_hybrid_ocr(
self,
image: Image.Image,
confidence_threshold: int = 65,
padding: int = 5,
ocr: Optional[Any] = None,
) -> Dict[str, list]:
"""
Performs OCR using Tesseract for bounding boxes and PaddleOCR for low-confidence text.
Returns data in the same dictionary format as pytesseract.image_to_data.
"""
if ocr is None:
if hasattr(self, "paddle_ocr") and self.paddle_ocr is not None:
ocr = self.paddle_ocr
else:
raise ValueError(
"No OCR object provided and 'paddle_ocr' is not initialized."
)
print("Starting hybrid OCR process...")
# 1. Get initial word-level results from Tesseract
tesseract_data = pytesseract.image_to_data(
image,
output_type=pytesseract.Output.DICT,
config=self.tesseract_config,
lang=self.tesseract_lang,
)
# tesseract_data['abs_line_id'] = tesseract_data.groupby(['block_num', 'par_num', 'line_num']).ngroup()
final_data = {
"text": [],
"left": [],
"top": [],
"width": [],
"height": [],
"conf": [],
}
num_words = len(tesseract_data["text"])
# This handles the "no text on page" case. If num_words is 0, the loop is skipped
# and an empty dictionary with empty lists is returned, which is the correct behavior.
for i in range(num_words):
text = tesseract_data["text"][i]
conf = int(tesseract_data["conf"][i])
# Skip empty text boxes or non-word elements (like page/block markers)
if not text.strip() or conf == -1:
continue
left = tesseract_data["left"][i]
top = tesseract_data["top"][i]
width = tesseract_data["width"][i]
height = tesseract_data["height"][i]
# line_number = tesseract_data['abs_line_id'][i]
# If confidence is low, use PaddleOCR for a second opinion
if conf < confidence_threshold:
img_width, img_height = image.size
crop_left = max(0, left - padding - 15)
crop_top = max(0, top - padding)
crop_right = min(img_width, left + width + padding + 15)
crop_bottom = min(img_height, top + height + padding)
# Ensure crop dimensions are valid
if crop_right <= crop_left or crop_bottom <= crop_top:
continue # Skip invalid crops
cropped_image = image.crop(
(crop_left, crop_top, crop_right, crop_bottom)
)
cropped_image_np = np.array(cropped_image)
if len(cropped_image_np.shape) == 2:
cropped_image_np = np.stack([cropped_image_np] * 3, axis=-1)
paddle_results = ocr.predict(cropped_image_np)
if paddle_results and paddle_results[0]:
rec_texts = paddle_results[0].get("rec_texts", [])
rec_scores = paddle_results[0].get("rec_scores", [])
if rec_texts and rec_scores:
new_text = " ".join(rec_texts)
new_conf = int(round(np.median(rec_scores) * 100, 0))
# Only replace if Paddle's confidence is better
if new_conf > conf:
print(
f" Re-OCR'd word: '{text}' (conf: {conf}) -> '{new_text}' (conf: {new_conf:.0f})"
)
# For exporting example image comparisons, not used here
safe_text = self._sanitize_filename(text, max_length=20)
self._sanitize_filename(new_text, max_length=20)
output_image_path = f"examples/tess_vs_paddle_examples/{conf}_conf_{safe_text}_to_{new_text}_{new_conf}.png"
cropped_image.save(output_image_path)
text = new_text
conf = new_conf
else:
print(
f" '{text}' (conf: {conf}) -> Paddle result '{new_text}' (conf: {new_conf:.0f}) was not better. Keeping original."
)
else:
# Paddle ran but found nothing, so discard the original low-confidence word
print(
f" '{text}' (conf: {conf}) -> No text found by Paddle. Discarding."
)
text = ""
else:
# Paddle found nothing, discard original word
print(
f" '{text}' (conf: {conf}) -> No text found by Paddle. Discarding."
)
text = ""
# Append the final result (either original, replaced, or skipped if empty)
if text.strip():
final_data["text"].append(clean_unicode_text(text))
final_data["left"].append(left)
final_data["top"].append(top)
final_data["width"].append(width)
final_data["height"].append(height)
final_data["conf"].append(int(conf))
# final_data['line_number'].append(int(line_number))
return final_data
def perform_ocr(
self, image: Union[str, Image.Image, np.ndarray], ocr: Optional[Any] = None
) -> List[OCRResult]:
"""
Performs OCR on the given image using the configured engine.
"""
if isinstance(image, str):
image = Image.open(image)
elif isinstance(image, np.ndarray):
image = Image.fromarray(image)
# Pre-process image - currently seems to give worse results!
if str(PREPROCESS_LOCAL_OCR_IMAGES).lower() == "true":
image, preprocessing_metadata = self.image_preprocessor.preprocess_image(
image
)
else:
preprocessing_metadata = {}
# Note: In testing I haven't seen that this necessarily improves results
if self.ocr_engine == "hybrid":
# Try hybrid with original image for cropping:
ocr_data = self._perform_hybrid_ocr(image)
elif self.ocr_engine == "tesseract":
ocr_data = pytesseract.image_to_data(
image,
output_type=pytesseract.Output.DICT,
config=self.tesseract_config,
lang=self.tesseract_lang, # Ensure the Tesseract language data (e.g., fra.traineddata) is installed on your system.
)
# ocr_data['abs_line_id'] = ocr_data.groupby(['block_num', 'par_num', 'line_num']).ngroup()
elif self.ocr_engine == "paddle":
image_np = np.array(image) # image_processed
# PaddleOCR may need an RGB image. Ensure it has 3 channels.
if len(image_np.shape) == 2:
image_np = np.stack([image_np] * 3, axis=-1)
if ocr is None:
if hasattr(self, "paddle_ocr") and self.paddle_ocr is not None:
ocr = self.paddle_ocr
else:
raise ValueError(
"No OCR object provided and 'paddle_ocr' is not initialised."
)
# ocr = PaddleOCR(use_textline_orientation=True, lang='en')
paddle_results = ocr.predict(image_np)
ocr_data = self._convert_paddle_to_tesseract_format(paddle_results)
else:
raise RuntimeError(f"Unsupported OCR engine: {self.ocr_engine}")
if preprocessing_metadata:
scale_factor = preprocessing_metadata.get("scale_factor", 1.0)
ocr_data = rescale_ocr_data(ocr_data, scale_factor)
# The rest of your processing pipeline now works for both engines
ocr_result = ocr_data
# Filter out empty strings and low confidence results
valid_indices = [
i
for i, text in enumerate(ocr_result["text"])
if text.strip() and int(ocr_result["conf"][i]) > 0
]
return [
OCRResult(
text=clean_unicode_text(ocr_result["text"][i]),
left=ocr_result["left"][i],
top=ocr_result["top"][i],
width=ocr_result["width"][i],
height=ocr_result["height"][i], # ,
# line_number=ocr_result['abs_line_id'][i]
)
for i in valid_indices
]
def analyze_text(
self,
line_level_ocr_results: List[OCRResult],
ocr_results_with_words: Dict[str, Dict],
chosen_redact_comprehend_entities: List[str],
pii_identification_method: str = LOCAL_PII_OPTION,
comprehend_client="",
custom_entities: List[str] = custom_entities,
language: Optional[str] = DEFAULT_LANGUAGE,
nlp_analyser: AnalyzerEngine = None,
**text_analyzer_kwargs,
) -> List[CustomImageRecognizerResult]:
page_text = ""
page_text_mapping = list()
all_text_line_results = list()
comprehend_query_number = 0
if not nlp_analyser:
nlp_analyser = self.analyzer_engine
# Collect all text and create mapping
for i, line_level_ocr_result in enumerate(line_level_ocr_results):
if page_text:
page_text += " "
start_pos = len(page_text)
page_text += line_level_ocr_result.text
# Note: We're not passing line_characters here since it's not needed for this use case
page_text_mapping.append((start_pos, i, line_level_ocr_result, None))
# Determine language for downstream services
aws_language = language or getattr(self, "language", None) or "en"
valid_language_entities = nlp_analyser.registry.get_supported_entities(
languages=[language]
)
if "CUSTOM" not in valid_language_entities:
valid_language_entities.append("CUSTOM")
if "CUSTOM_FUZZY" not in valid_language_entities:
valid_language_entities.append("CUSTOM_FUZZY")
# Process using either Local or AWS Comprehend
if pii_identification_method == LOCAL_PII_OPTION:
language_supported_entities = filter_entities_for_language(
custom_entities, valid_language_entities, language
)
if language_supported_entities:
text_analyzer_kwargs["entities"] = language_supported_entities
# if language != "en":
# gr.Info(f"Using {str(language_supported_entities)} entities for local model analysis for language: {language}")
else:
print(f"No relevant entities supported for language: {language}")
raise Warning(
f"No relevant entities supported for language: {language}"
)
analyzer_result = nlp_analyser.analyze(
text=page_text, language=language, **text_analyzer_kwargs
)
all_text_line_results = map_back_entity_results(
analyzer_result, page_text_mapping, all_text_line_results
)
elif pii_identification_method == AWS_PII_OPTION:
# Handle custom entities first
if custom_entities:
custom_redact_entities = [
entity
for entity in chosen_redact_comprehend_entities
if entity in custom_entities
]
if custom_redact_entities:
# Filter entities to only include those supported by the language
language_supported_entities = filter_entities_for_language(
custom_redact_entities, valid_language_entities, language
)
if language_supported_entities:
text_analyzer_kwargs["entities"] = language_supported_entities
page_analyser_result = nlp_analyser.analyze(
text=page_text, language=language, **text_analyzer_kwargs
)
all_text_line_results = map_back_entity_results(
page_analyser_result, page_text_mapping, all_text_line_results
)
# Process text in batches for AWS Comprehend
current_batch = ""
current_batch_mapping = list()
batch_char_count = 0
batch_word_count = 0
for i, text_line in enumerate(line_level_ocr_results):
words = text_line.text.split()
word_start_positions = list()
current_pos = 0
for word in words:
word_start_positions.append(current_pos)
current_pos += len(word) + 1
for word_idx, word in enumerate(words):
new_batch_char_count = len(current_batch) + len(word) + 1
if batch_word_count >= 50 or new_batch_char_count >= 200:
# Process current batch
all_text_line_results = do_aws_comprehend_call(
current_batch,
current_batch_mapping,
comprehend_client,
aws_language,
text_analyzer_kwargs.get("allow_list", []),
chosen_redact_comprehend_entities,
all_text_line_results,
)
comprehend_query_number += 1
# Reset batch
current_batch = word
batch_word_count = 1
batch_char_count = len(word)
current_batch_mapping = [
(0, i, text_line, None, word_start_positions[word_idx])
]
else:
if current_batch:
current_batch += " "
batch_char_count += 1
current_batch += word
batch_char_count += len(word)
batch_word_count += 1
if (
not current_batch_mapping
or current_batch_mapping[-1][1] != i
):
current_batch_mapping.append(
(
batch_char_count - len(word),
i,
text_line,
None,
word_start_positions[word_idx],
)
)
# Process final batch if any
if current_batch:
all_text_line_results = do_aws_comprehend_call(
current_batch,
current_batch_mapping,
comprehend_client,
aws_language,
text_analyzer_kwargs.get("allow_list", []),
chosen_redact_comprehend_entities,
all_text_line_results,
)
comprehend_query_number += 1
# Process results and create bounding boxes
combined_results = list()
for i, text_line in enumerate(line_level_ocr_results):
line_results = next(
(results for idx, results in all_text_line_results if idx == i), []
)
if line_results and i < len(ocr_results_with_words):
child_level_key = list(ocr_results_with_words.keys())[i]
ocr_results_with_words_line_level = ocr_results_with_words[
child_level_key
]
for result in line_results:
bbox_results = self.map_analyzer_results_to_bounding_boxes(
[result],
[
OCRResult(
text=text_line.text[result.start : result.end],
left=text_line.left,
top=text_line.top,
width=text_line.width,
height=text_line.height,
)
],
text_line.text,
text_analyzer_kwargs.get("allow_list", []),
ocr_results_with_words_line_level,
)
combined_results.extend(bbox_results)
return combined_results, comprehend_query_number
@staticmethod
def map_analyzer_results_to_bounding_boxes(
text_analyzer_results: List[RecognizerResult],
redaction_relevant_ocr_results: List[OCRResult],
full_text: str,
allow_list: List[str],
ocr_results_with_words_child_info: Dict[str, Dict],
) -> List[CustomImageRecognizerResult]:
redaction_bboxes = list()
for redaction_relevant_ocr_result in redaction_relevant_ocr_results:
# print("ocr_results_with_words_child_info:", ocr_results_with_words_child_info)
line_text = ocr_results_with_words_child_info["text"]
line_length = len(line_text)
redaction_text = redaction_relevant_ocr_result.text
for redaction_result in text_analyzer_results:
# Check if the redaction text is not in the allow list
if redaction_text not in allow_list:
# Adjust start and end to be within line bounds
start_in_line = max(0, redaction_result.start)
end_in_line = min(line_length, redaction_result.end)
# Get the matched text from this line
matched_text = line_text[start_in_line:end_in_line]
matched_text.split()
# Find the corresponding words in the OCR results
matching_word_boxes = list()
current_position = 0
for word_info in ocr_results_with_words_child_info.get("words", []):
word_text = word_info["text"]
word_length = len(word_text)
word_start = current_position
word_end = current_position + word_length
# Update current position for the next word
current_position += (
word_length + 1
) # +1 for the space after the word
# Check if the word's bounding box is within the start and end bounds
if word_start >= start_in_line and word_end <= (
end_in_line + 1
):
matching_word_boxes.append(word_info["bounding_box"])
# print(f"Matched word: {word_info['text']}")
if matching_word_boxes:
# Calculate the combined bounding box for all matching words
left = min(box[0] for box in matching_word_boxes)
top = min(box[1] for box in matching_word_boxes)
right = max(box[2] for box in matching_word_boxes)
bottom = max(box[3] for box in matching_word_boxes)
redaction_bboxes.append(
CustomImageRecognizerResult(
entity_type=redaction_result.entity_type,
start=start_in_line,
end=end_in_line,
score=redaction_result.score,
left=left,
top=top,
width=right - left,
height=bottom - top,
text=matched_text,
)
)
return redaction_bboxes
@staticmethod
def remove_space_boxes(ocr_result: dict) -> dict:
"""Remove OCR bboxes that are for spaces.
:param ocr_result: OCR results (raw or thresholded).
:return: OCR results with empty words removed.
"""
# Get indices of items with no text
idx = list()
for i, text in enumerate(ocr_result["text"]):
is_not_space = text.isspace() is False
if text != "" and is_not_space:
idx.append(i)
# Only retain items with text
filtered_ocr_result = {}
for key in list(ocr_result.keys()):
filtered_ocr_result[key] = [ocr_result[key][i] for i in idx]
return filtered_ocr_result
@staticmethod
def _scale_bbox_results(
ocr_result: Dict[str, List[Union[int, str]]], scale_factor: float
) -> Dict[str, float]:
"""Scale down the bounding box results based on a scale percentage.
:param ocr_result: OCR results (raw).
:param scale_percent: Scale percentage for resizing the bounding box.
:return: OCR results (scaled).
"""
scaled_results = deepcopy(ocr_result)
coordinate_keys = ["left", "top"]
dimension_keys = ["width", "height"]
for coord_key in coordinate_keys:
scaled_results[coord_key] = [
int(np.ceil((x) / (scale_factor))) for x in scaled_results[coord_key]
]
for dim_key in dimension_keys:
scaled_results[dim_key] = [
max(1, int(np.ceil(x / (scale_factor))))
for x in scaled_results[dim_key]
]
return scaled_results
@staticmethod
def estimate_x_offset(full_text: str, start: int) -> int:
# Estimate the x-offset based on character position
# This is a simple estimation and might need refinement for variable-width fonts
return int(start / len(full_text) * len(full_text))
def estimate_width(self, ocr_result: OCRResult, start: int, end: int) -> int:
# Extract the relevant text portion
relevant_text = ocr_result.text[start:end]
# If the relevant text is the same as the full text, return the full width
if relevant_text == ocr_result.text:
return ocr_result.width
# Estimate width based on the proportion of the relevant text length to the total text length
total_text_length = len(ocr_result.text)
relevant_text_length = len(relevant_text)
if total_text_length == 0:
return 0 # Avoid division by zero
# Proportion of the relevant text to the total text
proportion = relevant_text_length / total_text_length
# Estimate the width based on the proportion
estimated_width = int(proportion * ocr_result.width)
return estimated_width
def bounding_boxes_overlap(box1: List, box2: List):
"""Check if two bounding boxes overlap."""
return (
box1[0] < box2[2]
and box2[0] < box1[2]
and box1[1] < box2[3]
and box2[1] < box1[3]
)
def map_back_entity_results(
page_analyser_result: dict,
page_text_mapping: dict,
all_text_line_results: List[Tuple],
):
for entity in page_analyser_result:
entity_start = entity.start
entity_end = entity.end
# Track if the entity has been added to any line
added_to_line = False
for batch_start, line_idx, original_line, chars in page_text_mapping:
batch_end = batch_start + len(original_line.text)
# Check if the entity overlaps with the current line
if (
batch_start < entity_end and batch_end > entity_start
): # Overlap condition
relative_start = max(
0, entity_start - batch_start
) # Adjust start relative to the line
relative_end = min(
entity_end - batch_start, len(original_line.text)
) # Adjust end relative to the line
# Create a new adjusted entity
adjusted_entity = copy.deepcopy(entity)
adjusted_entity.start = relative_start
adjusted_entity.end = relative_end
# Check if this line already has an entry
existing_entry = next(
(entry for idx, entry in all_text_line_results if idx == line_idx),
None,
)
if existing_entry is None:
all_text_line_results.append((line_idx, [adjusted_entity]))
else:
existing_entry.append(
adjusted_entity
) # Append to the existing list of entities
added_to_line = True
# If the entity spans multiple lines, you may want to handle that here
if not added_to_line:
# Handle cases where the entity does not fit in any line (optional)
print(f"Entity '{entity}' does not fit in any line.")
return all_text_line_results
def map_back_comprehend_entity_results(
response: object,
current_batch_mapping: List[Tuple],
allow_list: List[str],
chosen_redact_comprehend_entities: List[str],
all_text_line_results: List[Tuple],
):
if not response or "Entities" not in response:
return all_text_line_results
for entity in response["Entities"]:
if entity.get("Type") not in chosen_redact_comprehend_entities:
continue
entity_start = entity["BeginOffset"]
entity_end = entity["EndOffset"]
# Track if the entity has been added to any line
added_to_line = False
# Find the correct line and offset within that line
for (
batch_start,
line_idx,
original_line,
chars,
line_offset,
) in current_batch_mapping:
batch_end = batch_start + len(original_line.text[line_offset:])
# Check if the entity overlaps with the current line
if (
batch_start < entity_end and batch_end > entity_start
): # Overlap condition
# Calculate the absolute position within the line
relative_start = max(0, entity_start - batch_start + line_offset)
relative_end = min(
entity_end - batch_start + line_offset, len(original_line.text)
)
result_text = original_line.text[relative_start:relative_end]
if result_text not in allow_list:
adjusted_entity = entity.copy()
adjusted_entity["BeginOffset"] = (
relative_start # Now relative to the full line
)
adjusted_entity["EndOffset"] = relative_end
recogniser_entity = recognizer_result_from_dict(adjusted_entity)
existing_entry = next(
(
entry
for idx, entry in all_text_line_results
if idx == line_idx
),
None,
)
if existing_entry is None:
all_text_line_results.append((line_idx, [recogniser_entity]))
else:
existing_entry.append(
recogniser_entity
) # Append to the existing list of entities
added_to_line = True
# Optional: Handle cases where the entity does not fit in any line
if not added_to_line:
print(f"Entity '{entity}' does not fit in any line.")
return all_text_line_results
def do_aws_comprehend_call(
current_batch: str,
current_batch_mapping: List[Tuple],
comprehend_client: botocore.client.BaseClient,
language: str,
allow_list: List[str],
chosen_redact_comprehend_entities: List[str],
all_text_line_results: List[Tuple],
):
if not current_batch:
return all_text_line_results
max_retries = 3
retry_delay = 3
for attempt in range(max_retries):
try:
response = comprehend_client.detect_pii_entities(
Text=current_batch.strip(), LanguageCode=language
)
all_text_line_results = map_back_comprehend_entity_results(
response,
current_batch_mapping,
allow_list,
chosen_redact_comprehend_entities,
all_text_line_results,
)
return all_text_line_results
except Exception as e:
if attempt == max_retries - 1:
print("AWS Comprehend calls failed due to", e)
raise
time.sleep(retry_delay)
def run_page_text_redaction(
language: str,
chosen_redact_entities: List[str],
chosen_redact_comprehend_entities: List[str],
line_level_text_results_list: List[str],
line_characters: List,
page_analyser_results: List = list(),
page_analysed_bounding_boxes: List = list(),
comprehend_client=None,
allow_list: List[str] = None,
pii_identification_method: str = LOCAL_PII_OPTION,
nlp_analyser: AnalyzerEngine = None,
score_threshold: float = 0.0,
custom_entities: List[str] = None,
comprehend_query_number: int = 0,
):
"""
This function performs text redaction on a page based on the specified language and chosen entities.
Args:
language (str): The language code for the text being processed.
chosen_redact_entities (List[str]): A list of entities to be redacted from the text.
chosen_redact_comprehend_entities (List[str]): A list of entities identified by AWS Comprehend for redaction.
line_level_text_results_list (List[str]): A list of text lines extracted from the page.
line_characters (List): A list of character-level information for each line of text.
page_analyser_results (List, optional): Results from previous page analysis. Defaults to an empty list.
page_analysed_bounding_boxes (List, optional): Bounding boxes for the analysed page. Defaults to an empty list.
comprehend_client: The AWS Comprehend client for making API calls. Defaults to None.
allow_list (List[str], optional): A list of allowed entities that should not be redacted. Defaults to None.
pii_identification_method (str, optional): The method used for PII identification. Defaults to LOCAL_PII_OPTION.
nlp_analyser (AnalyzerEngine, optional): The NLP analyzer engine used for local analysis. Defaults to None.
score_threshold (float, optional): The threshold score for entity detection. Defaults to 0.0.
custom_entities (List[str], optional): A list of custom entities for redaction. Defaults to None.
comprehend_query_number (int, optional): A counter for the number of Comprehend queries made. Defaults to 0.
"""
page_text = ""
page_text_mapping = list()
all_text_line_results = list()
comprehend_query_number = 0
# Collect all text from the page
for i, text_line in enumerate(line_level_text_results_list):
if chosen_redact_entities:
if page_text:
page_text += " "
start_pos = len(page_text)
page_text += text_line.text
page_text_mapping.append((start_pos, i, text_line, line_characters[i]))
valid_language_entities = nlp_analyser.registry.get_supported_entities(
languages=[language]
)
if "CUSTOM" not in valid_language_entities:
valid_language_entities.append("CUSTOM")
if "CUSTOM_FUZZY" not in valid_language_entities:
valid_language_entities.append("CUSTOM_FUZZY")
# Process based on identification method
if pii_identification_method == LOCAL_PII_OPTION:
if not nlp_analyser:
raise ValueError("nlp_analyser is required for Local identification method")
language_supported_entities = filter_entities_for_language(
chosen_redact_entities, valid_language_entities, language
)
page_analyser_result = nlp_analyser.analyze(
text=page_text,
language=language,
entities=language_supported_entities,
score_threshold=score_threshold,
return_decision_process=True,
allow_list=allow_list,
)
all_text_line_results = map_back_entity_results(
page_analyser_result, page_text_mapping, all_text_line_results
)
elif pii_identification_method == AWS_PII_OPTION:
# Process custom entities if any
if custom_entities:
custom_redact_entities = [
entity
for entity in chosen_redact_comprehend_entities
if entity in custom_entities
]
language_supported_entities = filter_entities_for_language(
custom_redact_entities, valid_language_entities, language
)
if language_supported_entities:
page_analyser_result = nlp_analyser.analyze(
text=page_text,
language=language,
entities=language_supported_entities,
score_threshold=score_threshold,
return_decision_process=True,
allow_list=allow_list,
)
all_text_line_results = map_back_entity_results(
page_analyser_result, page_text_mapping, all_text_line_results
)
current_batch = ""
current_batch_mapping = list()
batch_char_count = 0
batch_word_count = 0
for i, text_line in enumerate(line_level_text_results_list):
words = text_line.text.split()
word_start_positions = list()
# Calculate word start positions within the line
current_pos = 0
for word in words:
word_start_positions.append(current_pos)
current_pos += len(word) + 1 # +1 for space
for word_idx, word in enumerate(words):
new_batch_char_count = len(current_batch) + len(word) + 1
if batch_word_count >= 50 or new_batch_char_count >= 200:
# Process current batch
all_text_line_results = do_aws_comprehend_call(
current_batch,
current_batch_mapping,
comprehend_client,
language,
allow_list,
chosen_redact_comprehend_entities,
all_text_line_results,
)
comprehend_query_number += 1
# Start new batch
current_batch = word
batch_word_count = 1
batch_char_count = len(word)
current_batch_mapping = [
(
0,
i,
text_line,
line_characters[i],
word_start_positions[word_idx],
)
]
else:
if current_batch:
current_batch += " "
batch_char_count += 1
current_batch += word
batch_char_count += len(word)
batch_word_count += 1
if not current_batch_mapping or current_batch_mapping[-1][1] != i:
current_batch_mapping.append(
(
batch_char_count - len(word),
i,
text_line,
line_characters[i],
word_start_positions[
word_idx
], # Add the word's start position within its line
)
)
# Process final batch
if current_batch:
all_text_line_results = do_aws_comprehend_call(
current_batch,
current_batch_mapping,
comprehend_client,
language,
allow_list,
chosen_redact_comprehend_entities,
all_text_line_results,
)
comprehend_query_number += 1
# Process results for each line
for i, text_line in enumerate(line_level_text_results_list):
line_results = next(
(results for idx, results in all_text_line_results if idx == i), []
)
if line_results:
text_line_bounding_boxes = merge_text_bounding_boxes(
line_results, line_characters[i]
)
page_analyser_results.extend(line_results)
page_analysed_bounding_boxes.extend(text_line_bounding_boxes)
return page_analysed_bounding_boxes
def merge_text_bounding_boxes(
analyser_results: dict,
characters: List[LTChar],
combine_pixel_dist: int = 20,
vertical_padding: int = 0,
):
"""
Merge identified bounding boxes containing PII that are very close to one another
"""
analysed_bounding_boxes = list()
original_bounding_boxes = list() # List to hold original bounding boxes
if len(analyser_results) > 0 and len(characters) > 0:
# Extract bounding box coordinates for sorting
bounding_boxes = list()
for result in analyser_results:
# print("Result:", result)
char_boxes = [
char.bbox
for char in characters[result.start : result.end]
if isinstance(char, LTChar)
]
char_text = [
char._text
for char in characters[result.start : result.end]
if isinstance(char, LTChar)
]
if char_boxes:
# Calculate the bounding box that encompasses all characters
left = min(box[0] for box in char_boxes)
bottom = min(box[1] for box in char_boxes)
right = max(box[2] for box in char_boxes)
top = max(box[3] for box in char_boxes) + vertical_padding
bbox = [left, bottom, right, top]
bounding_boxes.append(
(bottom, left, result, bbox, char_text)
) # (y, x, result, bbox, text)
# Store original bounding boxes
original_bounding_boxes.append(
{
"text": "".join(char_text),
"boundingBox": bbox,
"result": copy.deepcopy(result),
}
)
# print("Original bounding boxes:", original_bounding_boxes)
# Sort the results by y-coordinate and then by x-coordinate
bounding_boxes.sort()
merged_bounding_boxes = list()
current_box = None
current_y = None
current_result = None
current_text = list()
for y, x, result, next_box, text in bounding_boxes:
if current_y is None or current_box is None:
# Initialize the first bounding box
current_box = next_box
current_y = next_box[1]
current_result = result
current_text = list(text)
else:
vertical_diff_bboxes = abs(next_box[1] - current_y)
horizontal_diff_bboxes = abs(next_box[0] - current_box[2])
if (
vertical_diff_bboxes <= 5
and horizontal_diff_bboxes <= combine_pixel_dist
):
# Merge bounding boxes
# print("Merging boxes")
merged_box = current_box.copy()
merged_result = current_result
merged_text = current_text.copy()
merged_box[2] = next_box[2] # Extend horizontally
merged_box[3] = max(current_box[3], next_box[3]) # Adjust the top
merged_result.end = max(
current_result.end, result.end
) # Extend text range
try:
if current_result.entity_type != result.entity_type:
merged_result.entity_type = (
current_result.entity_type + " - " + result.entity_type
)
else:
merged_result.entity_type = current_result.entity_type
except Exception as e:
print("Unable to combine result entity types:", e)
if current_text:
merged_text.append(" ") # Add space between texts
merged_text.extend(text)
merged_bounding_boxes.append(
{
"text": "".join(merged_text),
"boundingBox": merged_box,
"result": merged_result,
}
)
else:
# Start a new bounding box
current_box = next_box
current_y = next_box[1]
current_result = result
current_text = list(text)
# Combine original and merged bounding boxes
analysed_bounding_boxes.extend(original_bounding_boxes)
analysed_bounding_boxes.extend(merged_bounding_boxes)
# print("Analysed bounding boxes:", analysed_bounding_boxes)
return analysed_bounding_boxes
def recreate_page_line_level_ocr_results_with_page(
page_line_level_ocr_results_with_words: dict,
):
reconstructed_results = list()
# Assume all lines belong to the same page, so we can just read it from one item
# page = next(iter(page_line_level_ocr_results_with_words.values()))["page"]
page = page_line_level_ocr_results_with_words["page"]
for line_data in page_line_level_ocr_results_with_words["results"].values():
bbox = line_data["bounding_box"]
text = line_data["text"]
if line_data["line"]:
line_number = line_data["line"]
# Recreate the OCRResult
line_result = OCRResult(
text=text,
left=bbox[0],
top=bbox[1],
width=bbox[2] - bbox[0],
height=bbox[3] - bbox[1],
line=line_number,
)
reconstructed_results.append(line_result)
page_line_level_ocr_results_with_page = {
"page": page,
"results": reconstructed_results,
}
return page_line_level_ocr_results_with_page
def split_words_and_punctuation_from_line(
line_of_words: List[OCRResult],
) -> List[OCRResult]:
"""
Takes a list of OCRResult objects and splits words with trailing/leading punctuation.
For a word like "example.", it creates two new OCRResult objects for "example"
and "." and estimates their bounding boxes. Words with internal hyphens like
"high-tech" are preserved.
"""
# Punctuation that will be split off. Hyphen is not included.
new_word_list = list()
for word_result in line_of_words:
word_text = word_result.text
# This regex finds a central "core" word, and captures leading and trailing punctuation
# Handles cases like "(word)." -> group1='(', group2='word', group3='.'
match = re.match(r"([(\[{]*)(.*?)_?([.,?!:;)\}\]]*)$", word_text)
# Handle words with internal hyphens that might confuse the regex
if "-" in word_text and not match.group(2):
core_part_text = word_text
leading_punc = ""
trailing_punc = ""
elif match:
leading_punc, core_part_text, trailing_punc = match.groups()
else: # Failsafe
new_word_list.append(word_result)
continue
# If no split is needed, just add the original and continue
if not leading_punc and not trailing_punc:
new_word_list.append(word_result)
continue
# --- A split is required ---
# Estimate new bounding boxes by proportionally allocating width
original_width = word_result.width
if not word_text or original_width == 0:
continue # Failsafe
avg_char_width = original_width / len(word_text)
current_left = word_result.left
# Add leading punctuation if it exists
if leading_punc:
punc_width = avg_char_width * len(leading_punc)
new_word_list.append(
OCRResult(
text=leading_punc,
left=current_left,
top=word_result.top,
width=punc_width,
height=word_result.height,
)
)
current_left += punc_width
# Add the core part of the word
if core_part_text:
core_width = avg_char_width * len(core_part_text)
new_word_list.append(
OCRResult(
text=core_part_text,
left=current_left,
top=word_result.top,
width=core_width,
height=word_result.height,
)
)
current_left += core_width
# Add trailing punctuation if it exists
if trailing_punc:
punc_width = avg_char_width * len(trailing_punc)
new_word_list.append(
OCRResult(
text=trailing_punc,
left=current_left,
top=word_result.top,
width=punc_width,
height=word_result.height,
)
)
return new_word_list
def create_ocr_result_with_children(
combined_results: dict, i: int, current_bbox: dict, current_line: list
):
combined_results["text_line_" + str(i)] = {
"line": i,
"text": current_bbox.text,
"bounding_box": (
current_bbox.left,
current_bbox.top,
current_bbox.left + current_bbox.width,
current_bbox.top + current_bbox.height,
),
"words": [
{
"text": word.text,
"bounding_box": (
word.left,
word.top,
word.left + word.width,
word.top + word.height,
),
}
for word in current_line
],
}
return combined_results["text_line_" + str(i)]
def combine_ocr_results(
ocr_results: List[OCRResult],
x_threshold: float = 50.0,
y_threshold: float = 12.0,
page: int = 1,
):
"""
Group OCR results into lines, splitting words from punctuation.
"""
if not ocr_results:
return {"page": page, "results": []}, {"page": page, "results": {}}
lines = list()
current_line = list()
for result in sorted(ocr_results, key=lambda x: (x.top, x.left)):
if not current_line or abs(result.top - current_line[0].top) <= y_threshold:
current_line.append(result)
else:
lines.append(sorted(current_line, key=lambda x: x.left))
current_line = [result]
if current_line:
lines.append(sorted(current_line, key=lambda x: x.left))
page_line_level_ocr_results = list()
page_line_level_ocr_results_with_words = {}
line_counter = 1
for line in lines:
if not line:
continue
# Process the line to split punctuation from words
processed_line = split_words_and_punctuation_from_line(line)
# Re-calculate the line-level text and bounding box from the ORIGINAL words
line_text = " ".join([word.text for word in line])
line_left = line[0].left
line_top = min(word.top for word in line)
line_right = max(word.left + word.width for word in line)
line_bottom = max(word.top + word.height for word in line)
final_line_bbox = OCRResult(
text=line_text,
left=line_left,
top=line_top,
width=line_right - line_left,
height=line_bottom - line_top,
line=line_counter,
)
page_line_level_ocr_results.append(final_line_bbox)
# Use the PROCESSED line to create the children. Creates a result within page_line_level_ocr_results_with_words
page_line_level_ocr_results_with_words["text_line_" + str(line_counter)] = (
create_ocr_result_with_children(
page_line_level_ocr_results_with_words,
line_counter,
final_line_bbox,
processed_line, # <-- Use the new, split list of words
)
)
line_counter += 1
page_level_results_with_page = {
"page": page,
"results": page_line_level_ocr_results,
}
page_level_results_with_words = {
"page": page,
"results": page_line_level_ocr_results_with_words,
}
return page_level_results_with_page, page_level_results_with_words