Spaces:
Sleeping
Sleeping
| """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') | |
| 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 []) | |