import os
import cv2
import numpy as np
import math
import time
import gc
import torch
import fitz # PyMuPDF
from PIL import Image
import gradio as gr
from paddleocr import PaddleOCR
from google import genai
from google.genai import types
from pydantic import BaseModel
import concurrent.futures
# --- 1. Structured Output Definitions ---
class BoundingBox(BaseModel):
box_2d: list[int]
text: str
# API Key configuration (cleared for security, now loaded via environment variables)
API_KEYS = []
def get_gemini_api_keys():
import os
import json
# 1. Check for GEMINI_API_KEYS (comma-separated env var)
env_keys_str = os.environ.get("GEMINI_API_KEYS")
if env_keys_str:
keys = [k.strip() for k in env_keys_str.split(",") if k.strip()]
if keys:
return keys
# 2. Check numbered env vars (GEMINI_API_KEY, GEMINI_API_KEY_2, etc.)
env_keys = []
default_key = os.environ.get("GEMINI_API_KEY")
if default_key:
env_keys.append(default_key.strip())
for i in range(1, 10):
k = os.environ.get(f"GEMINI_API_KEY_{i}") or os.environ.get(f"GEMINI_API_KEY{i}")
if k:
k_clean = k.strip()
if k_clean and k_clean not in env_keys:
env_keys.append(k_clean)
if env_keys:
return env_keys
# 3. Fallback to hardcoded keys
valid_keys = [k.strip() for k in API_KEYS if k.strip()]
if valid_keys:
return valid_keys
return []
def fetch_gemini_ocr_for_page(page_num, img_bytes, api_keys, start_key_idx, prompt, mode):
"""
Runs Gemini OCR call for a single page inside a thread.
Tries different API keys in a loop if one hits a rate limit or quota.
Handles fallback to gemini-2.5-flash.
"""
from google import genai
from google.genai import types
import time
num_keys = len(api_keys)
max_attempts = num_keys * 3 # Try each key up to 3 times
for attempt in range(max_attempts):
key_idx = (start_key_idx + attempt) % num_keys
api_key = api_keys[key_idx]
# If we have cycled through all keys at least once, sleep 5 seconds
if attempt >= num_keys and attempt % num_keys == 0:
print(f"[API] Seite {page_num+1}: Alle Keys einmal versucht. Schlafe 5 Sekunden...")
time.sleep(5)
client = genai.Client(api_key=api_key)
# Try gemini-3.1-flash-lite first, fallback to gemini-2.5-flash if needed
for model in ['gemini-3.1-flash-lite', 'gemini-2.5-flash']:
try:
response = client.models.generate_content(
model=model,
contents=[prompt, types.Part.from_bytes(data=img_bytes, mime_type='image/png')],
config=types.GenerateContentConfig(
response_mime_type="application/json",
response_schema=list[BoundingBox],
temperature=0.0
)
)
if response.parsed:
# Filter out boxes that are invalid or don't have exactly 4 values in box_2d
valid_parsed = []
for box in response.parsed:
if hasattr(box, 'box_2d') and box.box_2d and len(box.box_2d) == 4:
valid_parsed.append(box)
else:
print(f"[API] Warning: Filtered out invalid box on page {page_num+1}: {box}")
return page_num, valid_parsed, None
return page_num, response.parsed, None
except Exception as e:
error_msg = str(e)
print(f"[API] Fehler auf Seite {page_num+1} mit Key-Index {key_idx} (Modell {model}): {error_msg}")
# Check if it's a rate limit, quota issue, or server error
is_quota_or_rate = any(code in error_msg for code in ["429", "Quota", "exhausted", "ResourceExhausted", "limit"])
is_server_err = any(code in error_msg for code in ["503", "500", "502", "504", "unavailable"])
if not (is_quota_or_rate or is_server_err):
# For non-retriable errors (like bad requests), fail this page immediately
return page_num, None, e
# For rate limit or temporary server error, break the model loop to try the next key immediately
break
return page_num, None, Exception("Alle API-Schlüssel sind fehlgeschlagen oder im Limit.")
# Caching models for reuse
_PADDLE_OCR = None
def get_paddle_ocr():
global _PADDLE_OCR
if _PADDLE_OCR is None:
print("[API] Initializing PaddleOCR...")
_PADDLE_OCR = PaddleOCR(use_angle_cls=True, lang='de')
return _PADDLE_OCR
_TROCR_PROCESSOR = None
_TROCR_MODEL = None
def get_trocr():
global _TROCR_PROCESSOR, _TROCR_MODEL
if _TROCR_PROCESSOR is None or _TROCR_MODEL is None:
print("[API] Initializing TrOCR...")
import logging as transformers_logging
transformers_logging.getLogger("transformers").setLevel(transformers_logging.ERROR)
onnx_path = "trocr_onnx"
if os.path.exists(onnx_path):
from transformers import TrOCRProcessor
from optimum.onnxruntime import ORTModelForVision2Seq
_TROCR_PROCESSOR = TrOCRProcessor.from_pretrained(onnx_path)
_TROCR_MODEL = ORTModelForVision2Seq.from_pretrained(onnx_path, provider="CPUExecutionProvider")
else:
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
_TROCR_PROCESSOR = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
_TROCR_MODEL = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
return _TROCR_PROCESSOR, _TROCR_MODEL
def recursive_xy_cut(boxes_with_data):
if len(boxes_with_data) <= 1:
return boxes_with_data
gap_threshold = 5 # Mindestabstand
# 1. Horizontale Lücken berechnen
y_intervals = sorted([(b[0][1], b[0][3]) for b in boxes_with_data])
max_h_gap = 0
h_gap_y = None
max_y = y_intervals[0][1]
for i in range(1, len(y_intervals)):
if y_intervals[i][0] > max_y:
gap = y_intervals[i][0] - max_y
if gap > max_h_gap and gap > gap_threshold:
max_h_gap = gap
h_gap_y = (max_y + y_intervals[i][0]) / 2
max_y = max(max_y, y_intervals[i][1])
# 2. Vertikale Lücken berechnen
x_intervals = sorted([(b[0][0], b[0][2]) for b in boxes_with_data])
max_v_gap = 0
v_gap_x = None
max_x = x_intervals[0][1]
for i in range(1, len(x_intervals)):
if x_intervals[i][0] > max_x:
gap = x_intervals[i][0] - max_x
if gap > max_v_gap and gap > gap_threshold:
max_v_gap = gap
v_gap_x = (max_x + x_intervals[i][0]) / 2
max_x = max(max_x, x_intervals[i][1])
# 3. Entlang der GRÖSSTEN Lücke schneiden!
if max_h_gap == 0 and max_v_gap == 0:
return sorted(boxes_with_data, key=lambda b: (b[0][1], b[0][0]))
if max_v_gap > max_h_gap:
left_boxes = [b for b in boxes_with_data if (b[0][0]+b[0][2])/2 < v_gap_x]
right_boxes = [b for b in boxes_with_data if b not in left_boxes]
if len(left_boxes) > 0 and len(right_boxes) > 0:
return recursive_xy_cut(left_boxes) + recursive_xy_cut(right_boxes)
if h_gap_y is not None:
top_boxes = [b for b in boxes_with_data if (b[0][1]+b[0][3])/2 < h_gap_y]
bottom_boxes = [b for b in boxes_with_data if b not in top_boxes]
if len(top_boxes) > 0 and len(bottom_boxes) > 0:
return recursive_xy_cut(top_boxes) + recursive_xy_cut(bottom_boxes)
return sorted(boxes_with_data, key=lambda b: (b[0][1], b[0][0]))
def process_pdf(input_file_path, mode, smart_skip=True, progress=gr.Progress()):
if not input_file_path:
raise gr.Error("Bitte lade ein PDF hoch.")
print(f"[API] Processing {input_file_path} in mode: {mode} (smart_skip: {smart_skip})")
progress(0, desc="Initialisiere Datei und Modelle...")
# Load required clients and models
api_keys = None
if "Gemini" in mode:
api_keys = get_gemini_api_keys()
paddle_ocr = None
if "PaddleOCR" in mode:
paddle_ocr = get_paddle_ocr()
# Load PDF
doc = fitz.open(input_file_path)
num_pages = len(doc)
# Pre-analysis: Determine which pages need OCR and which need redaction
pages_to_ocr = []
pages_to_redact = []
for i in range(num_pages):
page = doc.load_page(i)
has_text = len(page.get_text().strip()) > 20
if has_text and smart_skip:
print(f"[API] Seite {i+1} hat bereits Text und Smart-Skip ist aktiv. Überspringe OCR.")
else:
if has_text and not smart_skip:
print(f"[API] Seite {i+1} hat bereits Text und Smart-Skip ist inaktiv. Wird später redigiert (Force-OCR).")
pages_to_redact.append(i)
pages_to_ocr.append(i)
output_filename = f"searchable_{os.path.basename(input_file_path)}"
output_path = os.path.join(os.path.dirname(input_file_path), output_filename)
# ==========================================
# MODUS 1: Schnell (Gemini Full-Page) - Parallel
# ==========================================
if mode == "Schnell (Gemini Full-Page)":
pages_img_bytes = {}
progress(0.05, desc="Rendere PDF-Seiten...")
for i in pages_to_ocr:
page = doc.load_page(i)
zoom = 150 / 72
mat = fitz.Matrix(zoom, zoom)
pix = page.get_pixmap(matrix=mat)
pages_img_bytes[i] = pix.tobytes("png")
results = {}
errors = {}
num_keys = len(api_keys)
progress(0.1, desc=f"Starte parallele Gemini Semantic Analyse (mit {num_keys} API-Schlüsseln)...")
prompt = """Du bist ein extrem präzises OCR-System für mathematische Vorlesungsskripte.
Extrahiere absolut JEDEN Text (sowohl handgeschrieben als auch Maschinenschrift / gedruckten Text).
Verpasse kein einziges mathematisches Symbol, keinen Bruch und keinen Index.
WICHTIG FÜR FORMELN: Wandle ALLE mathematischen Formeln zwingend in eine saubere, einzeilige und logisch lesbare Schreibweise um!
- Nutze Klammern und Schrägstriche für Brüche: (A)/(B)
- Nutze '^' für Exponenten und '_' für Indizes: x^(SV), q_BM
- Nutze korrekte Unicode-Sonderzeichen für alles andere: Wurzeln (√), Integrale (∫), Summen (∑), griechische Buchstaben (α, β, γ, μ) etc.
- ACHTUNG BEI EINHEITEN: Wenn Einheiten in eckigen Klammeln [...] neben einer Formel stehen, behalte die eckigen Klammern UNBEDINGT bei! Füge KEIN Multiplikationszeichen '*' dazwischen ein. Einheiten sind reine Beschriftungen, keine Faktoren!
- Versuche NICHT, das optische 2D-Layout von Formeln mit mehrzeiligen Leerzeichen nachzuahmen!
Fasse zusammenhängende Sätze, Absätze oder komplette mathematische Formeln in EINER GEMEINSAMEN BoundingBox zusammen.
Zerstückele Formeln oder Brüche NICHT in Einzelteile! Eine komplette Formel = Eine BoundingBox.
Ignoriere Hintergrundmuster wie Punktraster komplett.
Gib für jeden Textblock/jede Formel eine BoundingBox zurück. box_2d ist [ymin, xmin, ymax, xmax] von 0 bis 1000.
Speichere den erkannten Text bzw. die Formel im Feld 'text' der BoundingBox.
WARNUNG: Es ist strengstens verboten, als Wert für das Feld 'text' einfach nur das Platzhalterwort 'text' einzutragen! Schreibe dort immer den tatsächlich erkannten Text hinein."""
if pages_to_ocr:
max_workers = min(len(pages_to_ocr), 12)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = []
for idx, i in enumerate(pages_to_ocr):
# Pass the api_keys list and start key index to allow rotation
futures.append(executor.submit(fetch_gemini_ocr_for_page, i, pages_img_bytes[i], api_keys, idx, prompt, mode))
completed = 0
for future in concurrent.futures.as_completed(futures):
p_num, parsed, err = future.result()
results[p_num] = parsed
errors[p_num] = err
completed += 1
progress(completed / len(pages_to_ocr), desc=f"Gemini Semantic Analyse: {completed} von {len(pages_to_ocr)} Seiten abgeschlossen...")
progress(0.95, desc="Generiere durchsuchbares PDF...")
for page_num in pages_to_ocr:
page = doc.load_page(page_num)
if page_num in pages_to_redact:
print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
try:
traces = page.get_texttrace()
has_visible_text = any(t.get("type") != 3 for t in traces)
except Exception:
has_visible_text = len(page.get_text().strip()) > 0
if has_visible_text:
pix = page.get_pixmap(dpi=150)
img_bytes = pix.tobytes("png")
page.add_redact_annot(page.rect)
page.apply_redactions(images=1, graphics=1, text=0)
page.insert_image(page.rect, stream=img_bytes)
else:
page.add_redact_annot(page.rect)
page.apply_redactions(images=0, graphics=0, text=0)
err = errors.get(page_num)
if err:
print(f"[API] Fehler auf Seite {page_num+1}: {err}")
continue
parsed_boxes = results.get(page_num)
if parsed_boxes:
font = fitz.Font("helv")
descender = font.descender
boxes_with_data = []
for box in parsed_boxes:
if not box.box_2d or len(box.box_2d) != 4:
print(f"[API] Warning: Invalid box_2d length {len(box.box_2d) if box.box_2d else 0} on page {page_num+1}: {box}")
continue
ymin, xmin, ymax, xmax = box.box_2d
x0 = (xmin / 1000) * page.rect.width
y0 = (ymin / 1000) * page.rect.height
x1 = (xmax / 1000) * page.rect.width
y1 = (ymax / 1000) * page.rect.height
boxes_with_data.append(([x0, y0, x1, y1], box))
sorted_data = recursive_xy_cut(boxes_with_data)
for coords, box in sorted_data:
text = box.text
if not text.strip() or text.strip() in [".", "..."]:
continue
x0, y0, x1, y1 = coords
rect = fitz.Rect(x0, y0, x1, y1)
fontsize = rect.height
text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
scale_x = rect.width / text_length if text_length > 0 else 1.0
y_baseline = rect.y1 + (descender * fontsize)
point = fitz.Point(rect.x0, y_baseline)
matrix = fitz.Matrix(scale_x, 1.0)
try:
page.insert_text(point, text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(point, matrix))
except ValueError:
clean_text = text.encode("latin-1", "ignore").decode("latin-1")
if clean_text.strip():
try:
page.insert_text(point, clean_text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(point, matrix))
except Exception:
pass
gc.collect()
# ==========================================
# MODUS 2: Präzise (True Hybrid OCR) - Parallel
# ==========================================
elif mode == "Präzise (Hybrid: PaddleOCR + Gemini)":
pages_img_bytes = {}
pages_paddle_boxes = {}
for idx, i in enumerate(pages_to_ocr):
progress((idx + 1) / (len(pages_to_ocr) * 2) if pages_to_ocr else 0.5, desc=f"PaddleOCR Geometrie Analyse: Seite {i + 1} von {num_pages}...")
page = doc.load_page(i)
zoom = 150 / 72
mat = fitz.Matrix(zoom, zoom)
pix = page.get_pixmap(matrix=mat)
img_bytes = pix.tobytes("png")
pages_img_bytes[i] = img_bytes
img_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
if pix.n == 4:
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGBA2BGR)
else:
img_np = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
result = paddle_ocr.ocr(img_np)
paddle_boxes = [(line[0], line[1][0]) for line in result[0]] if result and result[0] else []
pages_paddle_boxes[i] = paddle_boxes
results = {}
errors = {}
num_keys = len(api_keys)
progress(0.5, desc=f"Starte parallele Gemini Semantic Analyse (mit {num_keys} API-Schlüsseln)...")
prompt = """Du bist ein extrem präzises OCR-System für mathematische Vorlesungsskripte.
Extrahiere absolut JEDEN Text (sowohl handgeschrieben als auch Maschinenschrift / gedruckten Text).
Verpasse kein einziges mathematisches Symbol, keinen Bruch und keinen Index.
WICHTIG FÜR FORMELN: Wandle ALLE mathematischen Formeln zwingend in eine saubere, einzeilige und logisch lesbare Schreibweise um!
- Nutze Klammern und Schrägstriche für Brüche: (A)/(B)
- Nutze '^' für Exponenten und '_' für Indizes: x^(SV), q_BM
- Nutze korrekte Unicode-Sonderzeichen für alles andere: Wurzeln (√), Integrale (∫), Summen (∑), griechische Buchstaben (α, β, γ, μ) etc.
- ACHTUNG BEI EINHEITEN: Wenn Einheiten in eckigen Klammeln [...] neben einer Formel stehen, behalte die eckigen Klammern UNBEDINGT bei! Füge KEIN Multiplikationszeichen '*' dazwischen ein. Einheiten sind reine Beschriftungen, keine Faktoren!
- Versuche NICHT, das optische 2D-Layout von Formeln mit mehrzeiligen Leerzeichen nachzuahmen!
WICHTIG FÜR DAS LAYOUT (ABSOLUT KRITISCH!):
1. NORMALE TEXTZEILEN: Du MUSST für JEDE physische Textzeile im Bild eine EIGENE, separate BoundingBox erstellen!
- Es ist STRENGSTENS VERBOTEN, mehrere Zeilen zu einem Absatz zusammenzufassen!
- Auch wenn eine Textzeile Variablen (wie f_A) enthält, ist sie eine normale Zeile und darf NICHT mit der Zeile darunter zusammengefasst werden.
2. MEHRZEILIGE BRÜCHE: NUR WIRKLICHE mehrzeilige Formeln (Zähler über Nenner) MÜSSEN in EINER gemeinsamen BoundingBox zusammengefasst werden.
Ignoriere Hintergrundmuster wie Punktraster komplett.
Gib für jeden Textblock/jede Formel eine BoundingBox zurück. box_2d ist [ymin, xmin, ymax, xmax] von 0 bis 1000.
Speichere den erkannten Text bzw. die Formel im Feld 'text' der BoundingBox.
WARNUNG: Es ist strengstens verboten, als Wert für das Feld 'text' einfach nur das Platzhalterwort 'text' einzutragen! Schreibe dort immer den tatsächlich erkannten Text hinein."""
if pages_to_ocr:
max_workers = min(len(pages_to_ocr), 12)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = []
for idx, i in enumerate(pages_to_ocr):
# Pass the api_keys list and start key index to allow rotation
futures.append(executor.submit(fetch_gemini_ocr_for_page, i, pages_img_bytes[i], api_keys, idx, prompt, mode))
completed = 0
for future in concurrent.futures.as_completed(futures):
p_num, parsed, err = future.result()
results[p_num] = parsed
errors[p_num] = err
completed += 1
progress(0.5 + (completed / len(pages_to_ocr)) * 0.5, desc=f"Gemini Semantic Analyse: {completed} von {len(pages_to_ocr)} Seiten abgeschlossen...")
progress(0.95, desc="Generiere durchsuchbares PDF...")
zoom = 150 / 72
for page_num in pages_to_ocr:
page = doc.load_page(page_num)
if page_num in pages_to_redact:
print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
try:
traces = page.get_texttrace()
has_visible_text = any(t.get("type") != 3 for t in traces)
except Exception:
has_visible_text = len(page.get_text().strip()) > 0
if has_visible_text:
pix = page.get_pixmap(dpi=150)
img_bytes = pix.tobytes("png")
page.add_redact_annot(page.rect)
page.apply_redactions(images=1, graphics=1, text=0)
page.insert_image(page.rect, stream=img_bytes)
else:
page.add_redact_annot(page.rect)
page.apply_redactions(images=0, graphics=0, text=0)
err = errors.get(page_num)
if err:
print(f"[API] Fehler auf Seite {page_num+1}: {err}")
continue
parsed_boxes = results.get(page_num)
paddle_boxes = pages_paddle_boxes.get(page_num, [])
if parsed_boxes:
boxes_with_data = []
g_rects = []
valid_g_boxes = []
for g_box in parsed_boxes:
if not g_box.box_2d or len(g_box.box_2d) != 4:
print(f"[API] Warning: Invalid box_2d length {len(g_box.box_2d) if g_box.box_2d else 0} on page {page_num+1}: {g_box}")
continue
ymin, xmin, ymax, xmax = g_box.box_2d
x0 = (xmin / 1000) * page.rect.width
y0 = (ymin / 1000) * page.rect.height
x1 = (xmax / 1000) * page.rect.width
y1 = (ymax / 1000) * page.rect.height
g_rects.append(fitz.Rect(x0, y0, x1, y1))
valid_g_boxes.append(g_box)
parsed_boxes = valid_g_boxes
assigned_p_boxes_per_g_idx = {idx: [] for idx in range(len(parsed_boxes))}
for pb_data in paddle_boxes:
p_box, p_text = pb_data
p_xmin = min(p[0] for p in p_box) / zoom
p_ymin = min(p[1] for p in p_box) / zoom
p_xmax = max(p[0] for p in p_box) / zoom
p_ymax = max(p[1] for p in p_box) / zoom
p_rect = fitz.Rect(p_xmin, p_ymin, p_xmax, p_ymax)
best_g_idx = -1
max_overlap = 0
for idx, g_rect in enumerate(g_rects):
overlap = g_rect.intersect(p_rect).get_area()
if overlap > max_overlap:
max_overlap = overlap
best_g_idx = idx
if best_g_idx != -1 and max_overlap > 0.1 * p_rect.get_area():
assigned_p_boxes_per_g_idx[best_g_idx].append(pb_data)
for idx, g_box in enumerate(parsed_boxes):
g_text = g_box.text
assigned_p_boxes = assigned_p_boxes_per_g_idx[idx]
if assigned_p_boxes:
assigned_p_boxes.sort(key=lambda b: min(p[1] for p in b[0]))
clustered_p_boxes = []
for pb_tuple in assigned_p_boxes:
pb, pt = pb_tuple
y_center = (min(p[1] for p in pb) + max(p[1] for p in pb)) / 2
added_to_cluster = False
for cluster in clustered_p_boxes:
if abs(y_center - cluster['y_center']) < (10 / zoom):
cluster['boxes'].append(pb_tuple)
all_y = [min(p[1] for b in cluster['boxes'] for p in b[0]), max(p[1] for b in cluster['boxes'] for p in b[0])]
cluster['y_center'] = sum(all_y) / 2
added_to_cluster = True
break
if not added_to_cluster:
clustered_p_boxes.append({'y_center': y_center, 'boxes': [pb_tuple]})
for cluster in clustered_p_boxes:
cluster['boxes'].sort(key=lambda b: min(p[0] for p in b[0]))
math_chars = sum(1 for c in g_text if c in ['=', '/', '^', '[', ']'])
is_formula = (math_chars >= 4 and "=" in g_text)
if not is_formula:
g_words = g_text.split()
word_idx = 0
for c_idx, cluster in enumerate(clustered_p_boxes):
cluster_boxes = cluster['boxes']
cluster_word_count = sum(max(1, len(pt.split())) for pb, pt in cluster_boxes)
chunk = g_words[word_idx : word_idx + cluster_word_count]
line_text = " ".join(chunk)
word_idx += cluster_word_count
if c_idx == len(clustered_p_boxes) - 1 and word_idx < len(g_words):
if line_text: line_text += " "
line_text += " ".join(g_words[word_idx:])
if not line_text.strip(): continue
total_dx, total_dy = 0, 0
all_points = []
for pb, pt in cluster_boxes:
total_dx += pb[1][0] - pb[0][0]
total_dy += pb[1][1] - pb[0][1]
for p in pb: all_points.append((p[0]/zoom, p[1]/zoom))
angle_rad = math.atan2(total_dy, total_dx) if (total_dx != 0 or total_dy != 0) else 0
angle_deg = math.degrees(angle_rad)
cos_a, sin_a = math.cos(-angle_rad), math.sin(-angle_rad)
local_points = [(px * cos_a - py * sin_a, px * sin_a + py * cos_a) for px, py in all_points]
min_lx, max_lx = min(p[0] for p in local_points), max(p[0] for p in local_points)
min_ly, max_ly = min(p[1] for p in local_points), max(p[1] for p in local_points)
cos_inv, sin_inv = math.cos(angle_rad), math.sin(angle_rad)
merged_box = [[lx * cos_inv - ly * sin_inv, lx * sin_inv + ly * cos_inv] for lx, ly in [(min_lx, min_ly), (max_lx, min_ly), (max_lx, max_ly), (min_lx, max_ly)]]
p0, p1, p2, p3 = merged_box
dx_up, dy_up = p0[0] - p3[0], p0[1] - p3[1]
font = fitz.Font("helv")
pdf_baseline = fitz.Point(p3[0] + dx_up * -font.descender, p3[1] + dy_up * -font.descender)
boxes_with_data.append(([min(p[0] for p in merged_box), min(p[1] for p in merged_box), max(p[0] for p in merged_box), max(p[1] for p in merged_box)], (line_text, pdf_baseline, math.hypot(p1[0]-p0[0], p1[1]-p0[1]), math.hypot(dx_up, dy_up), angle_deg)))
else:
total_dx, total_dy = 0, 0
all_points = []
for pb, pt in assigned_p_boxes:
total_dx += pb[1][0] - pb[0][0]
total_dy += pb[1][1] - pb[0][1]
for p in pb: all_points.append((p[0]/zoom, p[1]/zoom))
angle_rad = math.atan2(total_dy, total_dx) if (total_dx != 0 or total_dy != 0) else 0
angle_deg = math.degrees(angle_rad)
cos_a, sin_a = math.cos(-angle_rad), math.sin(-angle_rad)
local_points = [(px * cos_a - py * sin_a, px * sin_a + py * cos_a) for px, py in all_points]
min_lx, max_lx = min(p[0] for p in local_points), max(p[0] for p in local_points)
min_ly, max_ly = min(p[1] for p in local_points), max(p[1] for p in local_points)
cos_inv, sin_inv = math.cos(angle_rad), math.sin(angle_rad)
merged_box = [[lx * cos_inv - ly * sin_inv, lx * sin_inv + ly * cos_inv] for lx, ly in [(min_lx, min_ly), (max_lx, min_ly), (max_lx, max_ly), (min_lx, max_ly)]]
p0, p1, p2, p3 = merged_box
dx_up, dy_up = p0[0] - p3[0], p0[1] - p3[1]
font = fitz.Font("helv")
pdf_baseline = fitz.Point(p3[0] + dx_up * -font.descender, p3[1] + dy_up * -font.descender)
boxes_with_data.append(([min(p[0] for p in merged_box), min(p[1] for p in merged_box), max(p[0] for p in merged_box), max(p[1] for p in merged_box)], (g_text.replace('\n', ' '), pdf_baseline, math.hypot(p1[0]-p0[0], p1[1]-p0[1]), math.hypot(dx_up, dy_up), angle_deg)))
else:
ymin, xmin, ymax, xmax = g_box.box_2d
x0, y0, x1, y1 = (xmin / 1000) * page.rect.width, (ymin / 1000) * page.rect.height, (xmax / 1000) * page.rect.width, (ymax / 1000) * page.rect.height
boxes_with_data.append(([x0, y0, x1, y1], (g_text.replace('\n', ' '), fitz.Point(x0, y1 - (y1-y0)*0.2), x1-x0, y1-y0, 0)))
sorted_data = recursive_xy_cut(boxes_with_data)
for coords, data in sorted_data:
text, pdf_baseline, box_width_pdf, box_height_pdf, angle_deg = data
fontsize = box_height_pdf
text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
matrix = fitz.Matrix(box_width_pdf / text_length if text_length > 0 else 1.0, 1.0) * fitz.Matrix(-angle_deg)
try:
page.insert_text(pdf_baseline, text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(pdf_baseline, matrix))
except Exception: pass
gc.collect()
# ==========================================
# MODUS 3: Lokal Deep (TrOCR) - Sequentiell
# ==========================================
elif mode == "Lokal Deep (PaddleOCR + TrOCR)":
trocr_processor, trocr_model = get_trocr()
for idx, page_num in enumerate(pages_to_ocr):
progress(idx / len(pages_to_ocr) if pages_to_ocr else 1.0, desc=f"Verarbeite Seite {page_num + 1} von {num_pages}...")
page = doc.load_page(page_num)
zoom = 3.0
mat = fitz.Matrix(zoom, zoom)
pix = page.get_pixmap(matrix=mat)
img_np = cv2.cvtColor(np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n), cv2.COLOR_RGBA2BGR if pix.n == 4 else cv2.COLOR_RGB2BGR)
if page_num in pages_to_redact:
print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
try:
traces = page.get_texttrace()
has_visible_text = any(t.get("type") != 3 for t in traces)
except Exception:
has_visible_text = len(page.get_text().strip()) > 0
if has_visible_text:
pix = page.get_pixmap(dpi=150)
img_bytes = pix.tobytes("png")
page.add_redact_annot(page.rect)
page.apply_redactions(images=1, graphics=1, text=0)
page.insert_image(page.rect, stream=img_bytes)
else:
page.add_redact_annot(page.rect)
page.apply_redactions(images=0, graphics=0, text=0)
result = paddle_ocr.ocr(img_np)
if not result or not result[0]: continue
crops, valid_boxes = [], []
for line in [l for l in result[0] if l]:
box = line[0]
x_coords, y_coords = [int(p[0]) for p in box], [int(p[1]) for p in box]
crop_img = img_np[max(0, min(y_coords) - 2):min(img_np.shape[0], max(y_coords) + 2), max(0, min(x_coords) - 2):min(img_np.shape[1], max(x_coords) + 2)]
if crop_img.size > 0:
crops.append(Image.fromarray(cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB)))
valid_boxes.append((box, line[1][0]))
results = []
for b_idx in range(0, len(crops), 4):
try:
batch_texts = trocr_processor.batch_decode(trocr_model.generate(trocr_processor(crops[b_idx:b_idx+4], return_tensors="pt").pixel_values, max_new_tokens=30), skip_special_tokens=True)
results.extend(batch_texts)
except Exception: results.extend([vb[1] for vb in valid_boxes[b_idx:b_idx+4]])
boxes_with_data = [([min(p[0] for p in vb[0]), min(p[1] for p in vb[0]), max(p[0] for p in vb[0]), max(p[1] for p in vb[0])], (vb[0], results[idx])) for idx, vb in enumerate(valid_boxes) if results[idx].strip()]
for coords, (box, text) in recursive_xy_cut(boxes_with_data):
p0, p1, p3 = box[0], box[1], box[3]
angle_deg = math.degrees(math.atan2(p1[1]-p0[1], p1[0]-p0[0]))
font = fitz.Font("helv")
base_x, base_y = (p3[0] + (p0[0]-p3[0]) * -font.descender) / zoom, (p3[1] + (p0[1]-p3[1]) * -font.descender) / zoom
fontsize = math.hypot(p0[0]-p3[0], p0[1]-p3[1]) / zoom
matrix = fitz.Matrix((math.hypot(p1[0]-p0[0], p1[1]-p0[1]) / zoom) / fitz.get_text_length(text, fontname="helv", fontsize=fontsize), 1.0) * fitz.Matrix(-angle_deg)
try:
page.insert_text(fitz.Point(base_x, base_y), text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(fitz.Point(base_x, base_y), matrix))
except Exception: pass
# ==========================================
# MODUS 4: Lokal Schnell (PaddleOCR) - Sequentiell
# ==========================================
elif mode == "Lokal Schnell (PaddleOCR)":
for idx, page_num in enumerate(pages_to_ocr):
progress(idx / len(pages_to_ocr) if pages_to_ocr else 1.0, desc=f"Verarbeite Seite {page_num + 1} von {num_pages} (Lokal Schnell)...")
page = doc.load_page(page_num)
zoom = 3.0
mat = fitz.Matrix(zoom, zoom)
pix = page.get_pixmap(matrix=mat)
img_np = cv2.cvtColor(np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n), cv2.COLOR_RGBA2BGR if pix.n == 4 else cv2.COLOR_RGB2BGR)
if page_num in pages_to_redact:
print(f"[API] Bereinige alten Textlayer auf Seite {page_num+1} (Force-OCR)...")
try:
traces = page.get_texttrace()
has_visible_text = any(t.get("type") != 3 for t in traces)
except Exception:
has_visible_text = len(page.get_text().strip()) > 0
if has_visible_text:
pix = page.get_pixmap(dpi=150)
img_bytes = pix.tobytes("png")
page.add_redact_annot(page.rect)
page.apply_redactions(images=1, graphics=1, text=0)
page.insert_image(page.rect, stream=img_bytes)
else:
page.add_redact_annot(page.rect)
page.apply_redactions(images=0, graphics=0, text=0)
result = paddle_ocr.ocr(img_np)
if not result or not result[0]: continue
page_data = result[0]
boxes_with_data = []
for line in page_data:
if not line: continue
box = line[0]
text = line[1][0]
if not text.strip(): continue
xmin = min(p[0] for p in box)
ymin = min(p[1] for p in box)
xmax = max(p[0] for p in box)
ymax = max(p[1] for p in box)
boxes_with_data.append(([xmin, ymin, xmax, ymax], (box, text)))
sorted_data = recursive_xy_cut(boxes_with_data)
for coords, (box, text) in sorted_data:
p0, p1, p3 = box[0], box[1], box[3]
angle_deg = math.degrees(math.atan2(p1[1]-p0[1], p1[0]-p0[0]))
font = fitz.Font("helv")
base_x, base_y = (p3[0] + (p0[0]-p3[0]) * -font.descender) / zoom, (p3[1] + (p0[1]-p3[1]) * -font.descender) / zoom
fontsize = math.hypot(p0[0]-p3[0], p0[1]-p3[1]) / zoom
text_length = fitz.get_text_length(text, fontname="helv", fontsize=fontsize)
scale_x = (math.hypot(p1[0]-p0[0], p1[1]-p0[1]) / zoom) / text_length if text_length > 0 else 1.0
matrix = fitz.Matrix(scale_x, 1.0) * fitz.Matrix(-angle_deg)
try:
page.insert_text(fitz.Point(base_x, base_y), text, fontsize=fontsize, fontname="helv", render_mode=3, morph=(fitz.Point(base_x, base_y), matrix))
except Exception: pass
doc.save(output_path)
doc.close()
progress(1.0, desc="Fertig! PDF erfolgreich generiert.")
print(f"[API] Searchable PDF saved to {output_path}")
# Copy to history directory for persistent Web UI history
try:
history_dir = "output_files"
os.makedirs(history_dir, exist_ok=True)
dest_filename = f"ocr_{os.path.basename(input_file_path)}"
if not dest_filename.lower().endswith(".pdf"):
dest_filename += ".pdf"
dest_path = os.path.join(history_dir, dest_filename)
import shutil
shutil.copy2(output_path, dest_path)
print(f"[API] Copied output to history (overwritten if existed): {dest_path}")
except Exception as e:
print(f"[API] Warning: Failed to copy file to history: {e}")
return output_path
HISTORY_DIR = "output_files"
def get_history_files():
if not os.path.exists(HISTORY_DIR):
return []
import glob
# get all pdfs and zips
files = glob.glob(os.path.join(HISTORY_DIR, "*.pdf")) + glob.glob(os.path.join(HISTORY_DIR, "*.zip"))
# sort by modification time descending
files.sort(key=os.path.getmtime, reverse=True)
return files
def get_history_html():
files = get_history_files()
# Filter out zip files to prevent duplicate entries
files = [f for f in files if not f.endswith(".zip")]
if not files:
return "Keine verarbeiteten Dateien vorhanden.
"
html = """
| Dateiname |
Größe |
Aktion |
"""
for path in files:
filename = os.path.basename(path)
try:
size_bytes = os.path.getsize(path)
size_mb = size_bytes / (1024 * 1024)
size_str = f"{size_mb:.1f} MB"
except Exception:
size_str = "unbekannt"
# Point download url directly to the /file= endpoint served by Gradio (requires absolute path)
import urllib.parse
abs_path = os.path.abspath(path).replace(os.sep, '/')
encoded_path = urllib.parse.quote(abs_path, safe="/")
space_id = os.environ.get("SPACE_ID")
if space_id:
subdomain = space_id.lower().replace("/", "-")
base_url = f"https://{subdomain}.hf.space/gradio_api"
else:
base_url = ""
download_url = f"{base_url}/file={encoded_path}"
html += f"""
| 📄 {filename} |
{size_str} |
📥 Herunterladen
|
"""
html += """
"""
return html
def clear_history():
if os.path.exists(HISTORY_DIR):
import shutil
for f in os.listdir(HISTORY_DIR):
try:
path = os.path.join(HISTORY_DIR, f)
if os.path.isfile(path):
os.remove(path)
elif os.path.isdir(path):
shutil.rmtree(path)
except Exception:
pass
return []
def clear_history_and_get_html():
clear_history()
return get_history_html()
def download_all_as_zip():
if not os.path.exists(HISTORY_DIR):
return None
import zipfile
zip_path = os.path.join(HISTORY_DIR, "all_ocr_files.zip")
if os.path.exists(zip_path):
try:
os.remove(zip_path)
except Exception:
pass
files = [f for f in os.listdir(HISTORY_DIR) if f.lower().endswith(".pdf")]
if not files:
return None
with zipfile.ZipFile(zip_path, "w") as zipf:
for f in files:
zipf.write(os.path.join(HISTORY_DIR, f), arcname=f)
return zip_path
custom_css = """
body, .gradio-container {
background-color: #0f172a !important;
}
.feedback {
border-radius: 16px !important;
border: 1px solid rgba(255, 255, 255, 0.1) !important;
background: rgba(30, 41, 59, 0.7) !important;
backdrop-filter: blur(12px) !important;
box-shadow: 0 8px 32px 0 rgba(0, 0, 0, 0.37) !important;
}
button.primary {
background: linear-gradient(90deg, #007AFF 0%, #00C6FF 100%) !important;
border: none !important;
color: white !important;
border-radius: 8px !important;
font-weight: bold !important;
transition: all 0.3s ease !important;
}
button.primary:hover {
transform: translateY(-2px) !important;
box-shadow: 0 4px 20px rgba(0, 122, 255, 0.5) !important;
}
"""
# Gradio Interface
with gr.Blocks(title="OCR app API & Web Interface", theme=gr.themes.Default(primary_hue="blue", neutral_hue="slate"), css=custom_css) as demo:
gr.Markdown("# 📄 Multi-Mode OCR API & Web UI")
gr.Markdown("Wähle eine PDF-Datei und einen Modus, um ein durchsuchbares PDF zu generieren. Dieser Space kann auch programmgesteuert aufgerufen werden.")
with gr.Row():
with gr.Column():
file_input = gr.File(label="PDF Datei hochladen", file_types=[".pdf"])
mode_input = gr.Radio(
choices=["Schnell (Gemini Full-Page)", "Präzise (Hybrid: PaddleOCR + Gemini)", "Lokal Schnell (PaddleOCR)", "Lokal Deep (PaddleOCR + TrOCR)"],
value="Schnell (Gemini Full-Page)",
label="OCR Modus"
)
smart_skip_input = gr.Checkbox(value=True, label="Bereits durchsuchbare Seiten überspringen (Smart-Skip)")
btn = gr.Button("🚀 OCR starten", variant="primary")
with gr.Column():
file_output = gr.File(label="Durchsuchbares PDF herunterladen", interactive=False)
# Verlauf/History Section
with gr.Accordion("📋 Verlauf / Abgeschlossene Dateien", open=True):
gr.Markdown("Hier siehst du alle fertig verarbeiteten OCR-Dokumente der aktuellen Sitzung. Klicke auf ein PDF, um es herunterzuladen.")
history_files = gr.HTML(
value=get_history_html
)
with gr.Row():
zip_btn = gr.Button("📦 Als ZIP herunterladen")
refresh_btn = gr.Button("🔄 Verlauf aktualisieren")
clear_btn = gr.Button("🗑️ Verlauf leeren", variant="stop")
zip_output = gr.File(label="Erstelltes ZIP-Archiv", visible=False, interactive=False)
# Dummy compatibility button to reserve fn_index: 0 for older clients
compat_btn = gr.Button(visible=False)
compat_btn.click(
fn=process_pdf,
inputs=[file_input, mode_input],
outputs=file_output,
api_name="process_pdf"
)
btn.click(
fn=process_pdf,
inputs=[file_input, mode_input, smart_skip_input],
outputs=file_output,
api_name="process_pdf_v2"
).then(
fn=get_history_html,
inputs=[],
outputs=history_files
)
zip_btn.click(
fn=download_all_as_zip,
inputs=[],
outputs=zip_output
).then(
fn=lambda: gr.update(visible=True),
inputs=[],
outputs=zip_output
)
refresh_btn.click(
fn=get_history_html,
inputs=[],
outputs=history_files
)
clear_btn.click(
fn=clear_history_and_get_html,
inputs=[],
outputs=history_files
)
if __name__ == "__main__":
demo.queue().launch(server_name="0.0.0.0", server_port=7860, allowed_paths=[os.path.abspath("output_files")])