Blame for ocr_support.py

"""Bounded English and Chinese OCR for image-only PDF pages, shared by both converters.

Recognition runs on the existing server. No document is sent to another API.
"""
from contextlib import contextmanager
import csv
import io
import math
import os
import re
from pathlib import Path
import shutil
import subprocess
import tempfile
import time

import pdfplumber
import PyPDF2

OCR_LANGUAGES = {
    'eng': ('eng', 'English'),
    'chi_sim': ('chi_sim+eng', 'Simplified Chinese and English'),
    'chi_tra': ('chi_tra+eng', 'Traditional Chinese and English'),
}
CJK = r'[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]'


def validate_ocr_language(language):
    if language not in OCR_LANGUAGES:
        raise OcrError('Choose English, Simplified Chinese + English, or Traditional Chinese + English for OCR.')
    return language


def normalise_cjk(text):
    # Remove single OCR-inserted spaces between CJK glyphs; retain line breaks.
    return re.sub(r'(?<=' + CJK + r') (?=' + CJK + r')', '', str(text))


MAX_OCR_PAGES = 3
MAX_RENDER_EDGE = 3600
OCR_BUDGET_SECONDS = 32


class OcrError(ValueError):
    def __init__(self, message, code=None):
        super().__init__(message)
        self.code = code


def needs_ocr(page):
    text = (page.extract_text() or '').strip()
    images = {(image['x0'], image['top'], image['x1'], image['bottom']) for image in page.images}
    image_area = sum(max(0, box[2] - box[0]) * max(0, box[3] - box[1]) for box in images)
    # A short selectable footer must not hide an otherwise scanned page.
    return len(text) < 200 and image_area > page.width * page.height * .55


def centers(indices):
    groups = []
    for value in indices:
        if not groups or value - groups[-1][-1] > 2:
            groups.append([])
        groups[-1].append(int(value))
    return [int(sum(group) / len(group)) for group in groups]


def preprocess(image):
    import cv2
    import numpy as np
    gray = np.array(image.convert('L'))
    _, ink = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    # Correct a small skew only when long ruling lines provide a clear angle.
    # Edge pixels prevent a solid chart or photo from outvoting real rules.
    edges = cv2.Canny(gray, 50, 150)
    detected = cv2.HoughLinesP(edges, 1, np.pi / 1800, 100,
                              minLineLength=gray.shape[1] // 4, maxLineGap=15)
    if detected is not None:
        angles = [math.degrees(math.atan2(y2 - y1, x2 - x1))
                  for x1, y1, x2, y2 in detected[:, 0]
                  if x2 > x1 and abs(y2 - y1) < (x2 - x1) * .12]
        angle = float(np.median(angles)) if len(angles) >= 3 else 0
        if .2 < abs(angle) < 5:
            transform = cv2.getRotationMatrix2D((gray.shape[1] / 2, gray.shape[0] / 2), angle, 1)
            gray = cv2.warpAffine(gray, transform, (gray.shape[1], gray.shape[0]),
                                  flags=cv2.INTER_CUBIC, borderValue=255)
            _, ink = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)
    horizontal = cv2.morphologyEx(ink, cv2.MORPH_OPEN,
                                  cv2.getStructuringElement(cv2.MORPH_RECT, (max(40, gray.shape[1] // 30), 1)))
    vertical = cv2.morphologyEx(ink, cv2.MORPH_OPEN,
                                cv2.getStructuringElement(cv2.MORPH_RECT, (1, max(40, gray.shape[0] // 40))))
    # Keep real rules, not the long strokes inside Chinese glyphs.
    for mask, axis, minimum in [(horizontal, 2, max(120, gray.shape[1] // 8)),
                                (vertical, 3, max(120, gray.shape[0] // 25))]:
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        retained = np.zeros_like(mask)
        for contour in contours:
            box = cv2.boundingRect(contour)
            thickness = box[3 if axis == 2 else 2]
            if box[axis] >= minimum and thickness <= max(15, gray.shape[0 if axis == 2 else 1] // 150):
                cv2.drawContours(retained, [contour], -1, 255, -1)
        mask[:] = retained
    rules = cv2.bitwise_or(horizontal, vertical)
    clean = gray.copy()
    clean[cv2.dilate(rules, np.ones((7, 7), np.uint8)) > 0] = 255
    return clean, horizontal, vertical, ink


def read_words(path, page_number, language='eng'):
    with open(path, encoding='utf-8') as stream:
        words = [row for row in csv.DictReader(stream, delimiter='\t')
                 if row['level'] == '5' and row.get('text', '').strip() and float(row['conf']) >= 0]
    length = sum(len(word['text']) for word in words)
    score = sum(float(word['conf']) * len(word['text']) for word in words) / max(length, 1)
    uncertain = sum(len(word['text']) for word in words if float(word['conf']) < 40) / max(length, 1)
    chinese = sum(len(re.findall(CJK, word['text'])) for word in words)
    minimum = 20 if language != 'eng' and chinese >= 6 else 40
    if length < minimum or score < 70 or uncertain > .2:
        raise OcrError(f'Page {page_number}: OCR could not read this scan reliably. Use a clearer, upright printed scan and select its OCR language. Handwriting and vertical Chinese are not supported.', code='unreadable')
    return words


def table_grids(horizontal, vertical, ink, words, dpi):
    """Keep real row separators; infer columns only from clear full-height gutters."""
    import cv2
    import numpy as np
    contours, _ = cv2.findContours(horizontal, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    rules = sorted([cv2.boundingRect(item) for item in contours], key=lambda box: box[1])
    groups = []
    for x, y, width, height in rules:
        if width < ink.shape[1] * .35:
            continue
        group = next((group for group in groups
                      if abs(x - group['x']) < dpi / 24 and abs(width - group['width']) < dpi / 16), None)
        if group is None:
            group = {'x': x, 'width': width, 'ys': []}
            groups.append(group)
        group['ys'].append(y + height // 2)
    tables = []
    for group in groups:
        ys = centers(sorted(group['ys']))
        if len(ys) < 3:
            continue
        spacing = float(np.median(np.diff(ys)))
        if not dpi * 6 / 72 <= spacing <= dpi * 100 / 72:
            continue
        x0, x1 = group['x'], group['x'] + group['width'] - 1
        top = ys[0]
        # Header rows often have only a bottom rule. Include an aligned header
        # immediately above that rule, while excluding a preceding paragraph.
        above = [word for word in words
                 if top - spacing * 1.2 < int(word['top']) + int(word['height']) / 2 < top
                 and x0 <= int(word['left']) <= x1]
        if above:
            last_baseline = max(int(word['top']) for word in above)
            above = sorted([word for word in above if abs(int(word['top']) - last_baseline) < dpi * 5 / 72],
                           key=lambda word: int(word['left']))
            gaps = sum(int(right['left']) - int(left['left']) - int(left['width']) > dpi * 12 / 72
                       for left, right in zip(above, above[1:]))
            if gaps >= 2:
                top = max(0, min(int(word['top']) for word in above) - int(dpi * 4 / 72))
                ys.insert(0, top)
        bottom = ys[-1]
        if bottom <= top:
            continue
        region = ink[top:bottom + 1, x0:x1 + 1].copy()
        region[horizontal[top:bottom + 1, x0:x1 + 1] > 0] = 0
        vertical_strength = np.count_nonzero(vertical[top:bottom + 1, x0:x1 + 1], axis=0)
        xs = [x0 + value for value in centers(np.where(vertical_strength > (bottom - top) * .8)[0])]
        if len(xs) < 3:
            occupied = np.count_nonzero(region, axis=0)
            blank = np.where(occupied <= 1)[0]
            runs = []
            for value in blank:
                if not runs or value > runs[-1][-1] + 1:
                    runs.append([])
                runs[-1].append(int(value))
            gutters = [x0 + int((run[0] + run[-1]) / 2) for run in runs
                       if len(run) >= dpi * 8 / 72 and run[0] > dpi * 6 / 72
                       and run[-1] < x1 - x0 - dpi * 6 / 72]
            xs = [x0, *gutters, x1]
        if not 3 <= len(xs) <= 13:
            continue
        # Every proposed column must contain recognised text, including a header.
        if not all(any(left < int(word['left']) + int(word['width']) / 2 < right
                       and top < int(word['top']) + int(word['height']) / 2 < bottom
                       for word in words) for left, right in zip(xs, xs[1:])):
            continue
        tables.append({'xs': xs, 'ys': ys})
    return tables


def financial_table_grids(words, dpi, existing):
    """Recover common borderless invoices only with a clear header and numeric rows."""
    rows = []
    for word in sorted(words, key=lambda item: (int(item['top']) + int(item['height']) / 2, int(item['left']))):
        center = int(word['top']) + int(word['height']) / 2
        if not rows or abs(center - rows[-1]['center']) > dpi * 4 / 72:
            rows.append({'center': center, 'words': []})
        rows[-1]['words'].append(word)

    def groups(row):
        result = []
        for word in sorted(row['words'], key=lambda item: int(item['left'])):
            left, right = int(word['left']), int(word['left']) + int(word['width'])
            if not result or left - result[-1]['right'] > dpi * 12 / 72:
                result.append({'left': left, 'right': right, 'text': word['text']})
            else:
                result[-1]['right'] = max(right, result[-1]['right'])
                result[-1]['text'] += ' ' + word['text']
        return result

    def bounds(row):
        return min(int(word['top']) for word in row['words']), max(int(word['top']) + int(word['height']) for word in row['words'])

    result = []
    for index, row in enumerate(rows):
        header = groups(row)
        if not 3 <= len(header) <= 8:
            continue
        keywords = re.findall(r'\b(?:quantity|qty|amount|price|total|balance|debit|credit|rate|hours)\b',
                              normalise_cjk(' '.join(group['text'] for group in header)), re.I)
        keywords += re.findall(r'数量|數量|金额|金額|单价|單價|合计|合計|总计|總計|余额|餘額|税额|稅額|工时|工時', normalise_cjk(' '.join(group['text'] for group in header)))
        if len(keywords) < 2:
            continue
        top, bottom = bounds(row)
        if any(grid['ys'][0] <= row['center'] <= grid['ys'][-1] for grid in existing):
            continue
        selected = [(row, header)]
        for candidate in rows[index + 1:]:
            candidate_top, candidate_bottom = bounds(candidate)
            if candidate_top - bottom > dpi * 36 / 72:
                break
            cells = groups(candidate)
            numeric_end = bool(cells and re.fullmatch(r'(?:[$£€¥￥]|[A-Z]{3})?\s*[+-]?\(?\d[\d,.]*\)?%?', cells[-1]['text'].strip()))
            if len(cells) < 2 or not numeric_end:
                break
            selected.append((candidate, cells))
            bottom = candidate_bottom
        if len(selected) < 2:
            continue
        full_rows = [cells for _, cells in selected if len(cells) == len(header)]
        if len(full_rows) < 2:
            continue
        xs = [min(cells[0]['left'] for cells in full_rows) - dpi * 3 / 72]
        for column in range(len(header) - 1):
            right = max(cells[column]['right'] for cells in full_rows)
            left = min(cells[column + 1]['left'] for cells in full_rows)
            if left - right < dpi * 3 / 72:
                break
            xs.append((right + left) / 2)
        if len(xs) != len(header):
            continue
        xs.append(max(cells[-1]['right'] for _, cells in selected) + dpi * 3 / 72)
        ys = [top - dpi * 3 / 72]
        for (previous, _), (following, _) in zip(selected, selected[1:]):
            ys.append((bounds(previous)[1] + bounds(following)[0]) / 2)
        ys.append(bottom + dpi * 3 / 72)
        result.append({'xs': xs, 'ys': ys, 'row_groups': [cells for _, cells in selected]})
    return result


def add_table_rules(pdf_path, grids, image_size, output):
    if not grids:
        return pdf_path
    from reportlab.pdfgen import canvas
    reader = PyPDF2.PdfReader(pdf_path)
    page = reader.pages[0]
    width, height = float(page.mediabox.width), float(page.mediabox.height)
    scale_x, scale_y = width / image_size[0], height / image_size[1]
    overlay = io.BytesIO()
    drawing = canvas.Canvas(overlay, pagesize=(width, height))
    drawing.setLineWidth(.3)
    for grid in grids:
        xs, ys = grid['xs'], grid['ys']
        for y in ys:
            drawing.line(xs[0] * scale_x, height - y * scale_y, xs[-1] * scale_x, height - y * scale_y)
        for column, x in enumerate(xs):
            for row, (top, bottom) in enumerate(zip(ys, ys[1:])):
                merged = 0 < column < len(xs) - 1 and any(group['left'] < x < group['right']
                            for group in grid.get('row_groups', [[] for _ in ys])[row])
                if not merged:
                    drawing.line(x * scale_x, height - top * scale_y, x * scale_x, height - bottom * scale_y)
    drawing.save()
    overlay.seek(0)
    page.merge_page(PyPDF2.PdfReader(overlay).pages[0])
    writer = PyPDF2.PdfWriter()
    writer.add_page(page)
    with open(output, 'wb') as stream:
        writer.write(stream)
    return output


def recognise_page(page, page_number, workdir, deadline, require_table, language='eng'):
    from PIL import Image
    executable = shutil.which('tesseract')
    if not executable:
        raise OcrError('OCR is temporarily unavailable. Try a PDF with selectable text or return later.')
    dpi = min(300, MAX_RENDER_EDGE * 72 / max(page.width, page.height))
    image = page.to_image(resolution=dpi).original
    clean, horizontal, vertical, ink = preprocess(image)
    image.close()
    prefix = Path(workdir).resolve() / f'ocr-{page_number}'
    Image.fromarray(clean).save(str(prefix) + '.png', dpi=(dpi, dpi))
    remaining = deadline - time.monotonic()
    if remaining <= 1:
        raise OcrError('OCR reached the processing time limit. Split this PDF into fewer scanned pages and try again.')
    try:
        process = subprocess.run([executable, str(prefix) + '.png', str(prefix), '-l', OCR_LANGUAGES[language][0],
                                  '--dpi', str(round(dpi)), '--psm', '3', 'pdf', 'tsv'],
                                 stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                 timeout=min(15, remaining), env={**os.environ, 'OMP_THREAD_LIMIT': '2'})
    except subprocess.TimeoutExpired as error:
        raise OcrError(f'Page {page_number}: OCR timed out. Try a clearer scan or split the PDF into fewer pages.') from error
    if process.returncode != 0:
        raise OcrError(f'Page {page_number}: OCR could not process this scan. Try a clearer, upright printed scan and check the OCR language.')
    words = read_words(str(prefix) + '.tsv', page_number, language)
    grids = table_grids(horizontal, vertical, ink, words, dpi)
    grids.extend(financial_table_grids(words, dpi, grids))
    if require_table and not grids:
        raise OcrError(f'Page {page_number}: no clear scanned table was found. Use visible row separators or clear quantity/amount headers. Complex or merged scan layouts need manual conversion.')
    return add_table_rules(str(prefix) + '.pdf', grids, (clean.shape[1], clean.shape[0]), str(prefix) + '-tables.pdf')


@contextmanager
def prepare_for_extraction(pdf_path, require_table=False, enabled=True, ocr_language='eng', fallback_pages=None):
    validate_ocr_language(ocr_language)
    deadline = time.monotonic() + OCR_BUDGET_SECONDS
    with pdfplumber.open(pdf_path) as pdf:
        scanned = [index for index, page in enumerate(pdf.pages) if needs_ocr(page)]
        if not scanned:
            yield pdf_path, 0
            return
        if not enabled:
            raise OcrError('This PDF needs OCR. Reload the converter to use scan recognition, then try again.')
        if len(scanned) > MAX_OCR_PAGES:
            raise OcrError('OCR supports up to 3 scanned pages per PDF. Use Split PDF first; text-based PDFs can still have up to 10 pages.')
        with tempfile.TemporaryDirectory(prefix='pdf-ocr-') as workdir:
            reader = PyPDF2.PdfReader(pdf_path)
            writer = PyPDF2.PdfWriter()
            for index, page in enumerate(pdf.pages):
                if index in scanned:
                    try:
                        recognised = recognise_page(page, index + 1, workdir, deadline, require_table, ocr_language)
                    except OcrError as error:
                        if fallback_pages is None or error.code != 'unreadable':
                            raise
                        fallback_pages.append(index)
                        writer.add_page(reader.pages[index])
                    else:
                        writer.add_page(PyPDF2.PdfReader(recognised).pages[0])
                else:
                    writer.add_page(reader.pages[index])
            output = str(Path(workdir) / 'recognised.pdf')
            with open(output, 'wb') as stream:
                writer.write(stream)
            yield output, len(scanned) - len(fallback_pages or [])