From b79ca81965ef8a4ec2137d4f54bb40873880d7cd Mon Sep 17 00:00:00 2001 From: Alex-Gilbert Date: Sun, 15 Feb 2026 18:41:05 -0600 Subject: [PATCH] feat: OCR module with character-level bounding boxes Co-Authored-By: Claude Opus 4.6 --- tests/test_ocr.py | 40 ++++++++++++++ unredact/pipeline/ocr.py | 113 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 153 insertions(+) create mode 100644 tests/test_ocr.py create mode 100644 unredact/pipeline/ocr.py diff --git a/tests/test_ocr.py b/tests/test_ocr.py new file mode 100644 index 0000000..59f8d0e --- /dev/null +++ b/tests/test_ocr.py @@ -0,0 +1,40 @@ +from pathlib import Path + +from PIL import Image + +from unredact.pipeline.ocr import ocr_page, OcrChar, OcrLine +from unredact.pipeline.rasterize import rasterize_pdf + + +def test_ocr_returns_lines_with_chars(sample_pdf: Path): + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + lines = ocr_page(pages[0]) + assert len(lines) > 0 + assert all(isinstance(line, OcrLine) for line in lines) + # Each line should have characters + non_empty = [l for l in lines if l.chars] + assert len(non_empty) > 0 + + +def test_ocr_chars_have_bounding_boxes(sample_pdf: Path): + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + lines = ocr_page(pages[0]) + for line in lines: + for char in line.chars: + assert isinstance(char, OcrChar) + assert char.x >= 0 + assert char.y >= 0 + assert char.w > 0 + assert char.h > 0 + assert len(char.text) == 1 + + +def test_ocr_finds_known_text(sample_pdf: Path): + """The first page contains 'Jeffrey Epstein' — OCR should find it.""" + pages = rasterize_pdf(sample_pdf, first_page=1, last_page=1) + lines = ocr_page(pages[0]) + full_text = " ".join( + "".join(c.text for c in line.chars) for line in lines + ) + # Tesseract may not be perfect, but should get close + assert "Jeffrey" in full_text or "jeffrey" in full_text.lower() diff --git a/unredact/pipeline/ocr.py b/unredact/pipeline/ocr.py new file mode 100644 index 0000000..96ff909 --- /dev/null +++ b/unredact/pipeline/ocr.py @@ -0,0 +1,113 @@ +from dataclasses import dataclass + +import pytesseract +from PIL import Image + + +@dataclass +class OcrChar: + """A single OCR'd character with its bounding box.""" + text: str + x: int + y: int + w: int + h: int + conf: float # 0-100 confidence + + +@dataclass +class OcrLine: + """A line of OCR'd text.""" + chars: list[OcrChar] + x: int + y: int + w: int + h: int + + @property + def text(self) -> str: + return "".join(c.text for c in self.chars) + + +def ocr_page(image: Image.Image) -> list[OcrLine]: + """Run Tesseract OCR on an image and return character-level results. + + Uses tsv output to get bounding boxes at every level. + We use word-level boxes and estimate character positions from them + (Tesseract's character-level output is unreliable). + """ + data = pytesseract.image_to_data( + image, output_type=pytesseract.Output.DICT, config="--psm 6" + ) + + # Group words into lines, then estimate character positions + lines: dict[tuple[int, int, int, int], list[dict]] = {} + n = len(data["text"]) + + for i in range(n): + text = data["text"][i].strip() + if not text: + continue + conf = float(data["conf"][i]) + if conf < 0: + continue + + line_key = ( + data["block_num"][i], + data["par_num"][i], + data["line_num"][i], + 0, + ) + word_info = { + "text": text, + "x": data["left"][i], + "y": data["top"][i], + "w": data["width"][i], + "h": data["height"][i], + "conf": conf, + } + lines.setdefault(line_key, []).append(word_info) + + result: list[OcrLine] = [] + for _key, words in sorted(lines.items()): + chars: list[OcrChar] = [] + for wi, word in enumerate(words): + # Estimate per-character positions by dividing word box evenly + word_text = word["text"] + if not word_text: + continue + char_w = word["w"] / len(word_text) + for ci, ch in enumerate(word_text): + chars.append(OcrChar( + text=ch, + x=int(word["x"] + ci * char_w), + y=word["y"], + w=max(1, int(char_w)), + h=word["h"], + conf=word["conf"], + )) + # Add space character between words (except after last word) + if wi < len(words) - 1: + next_word = words[wi + 1] + space_x = word["x"] + word["w"] + space_w = next_word["x"] - space_x + if space_w > 0: + chars.append(OcrChar( + text=" ", + x=space_x, + y=word["y"], + w=space_w, + h=word["h"], + conf=word["conf"], + )) + + if chars: + line_x = chars[0].x + line_y = min(c.y for c in chars) + line_w = (chars[-1].x + chars[-1].w) - line_x + line_h = max(c.y + c.h for c in chars) - line_y + result.append(OcrLine( + chars=chars, x=line_x, y=line_y, w=line_w, h=line_h + )) + + return result