crazylemonade commited on
Commit
6b1b2d4
·
verified ·
1 Parent(s): 1bb000f

Delete main.py

Browse files
Files changed (1) hide show
  1. main.py +0 -1360
main.py DELETED
@@ -1,1360 +0,0 @@
1
- """
2
- OpenSkill OCR Service — v4.0
3
- FastAPI application for Hugging Face Docker Space (CPU / pipeline backend)
4
-
5
- ═══════════════════════════════════════════════════════════════════════════════
6
- ARCHITECTURE (v4.0 — OCR-only, AI-first)
7
- ═══════════════════════════════════════════════════════════════════════════════
8
-
9
- This service is an extraction layer only. It does NOT:
10
- - classify documents
11
- - extract named entities
12
- - validate fields
13
- - generate summaries
14
- - perform board/marksheet/JEE-specific logic
15
-
16
- All document understanding is delegated to the AI layer downstream.
17
-
18
- PATH A — Fast OCR (images: jpg / png / webp / bmp / heic / heif / avif)
19
- Engine : rapidocr-onnxruntime ≥ 1.3.22
20
- Models : Bundled in pip wheel — zero first-use download, ~50 MB
21
- Resize : images capped at MAX_OCR_SIDE px (default 1600) before inference
22
- Target : 1–4 s (acceptable < 8 s)
23
- Fallback: if confidence < FAST_CONFIDENCE_THRESHOLD → MinerU fallback
24
-
25
- PATH B — Full pipeline (PDFs, multi-page, layout-sensitive docs)
26
- Engine : MinerU magic-pdf pipeline backend
27
- Models : opendatalab/PDF-Extract-Kit-1.0 (downloaded at build time)
28
- Target : 5–20 s (acceptable < 30 s)
29
-
30
- ═══════════════════════════════════════════════════════════════════════════════
31
- RESPONSE FORMAT (v4.0)
32
- ═══════════════════════════════════════════════════════════════════════════════
33
-
34
- {
35
- "success": true,
36
- "filename": "scan.jpg",
37
- "engine": "rapidocr",
38
- "confidence": 0.91,
39
- "text": "...",
40
- "markdown": "...",
41
- "pageCount": 1,
42
- "cached": false,
43
- "processingTimeMs": 1840,
44
- "timings": {
45
- "uploadMs": 12,
46
- "hashMs": 4,
47
- "memCheckMs": 8,
48
- "decodeMs": 55,
49
- "resizeMs": 18,
50
- "detectMs": 610,
51
- "recognizeMs": 980,
52
- "postProcessMs": 14,
53
- "totalMs": 1840
54
- },
55
- "metadata": {
56
- "imgW": 3024,
57
- "imgH": 4032,
58
- "imgWResized": 1200,
59
- "imgHResized": 1600,
60
- "textBlocks": 47,
61
- "passesUsed": 1,
62
- "backend": "rapidocr"
63
- }
64
- }
65
-
66
- ═══════════════════════════════════════════════════════════════════════════════
67
- API ENDPOINTS
68
- ═══════════════════════════════════════════════════════════════════════════════
69
-
70
- GET /health Liveness (always fast)
71
- GET /status Node status: memory, uptime, cache, engine state
72
- GET /warmup Pre-load both OCR engines (also called at startup)
73
- GET /diagnostics Full environment + model inventory
74
- POST /benchmark Multi-size RapidOCR timing benchmark (small/medium/large)
75
- POST /extract Single file — PDF or image — with SHA256 cache
76
- POST /batch Up to 8 files, sequential, per-file error isolation
77
- """
78
-
79
- import hashlib
80
- import io
81
- import os
82
- import re
83
- import shutil
84
- import sys
85
- import tempfile
86
- import threading
87
- import time
88
- import traceback
89
- import logging
90
- from importlib.metadata import version as pkg_version
91
- from typing import Any, Optional
92
-
93
- import fitz # PyMuPDF
94
- import numpy as np
95
- from PIL import Image
96
-
97
- from fastapi import FastAPI, File, UploadFile
98
- from fastapi.middleware.cors import CORSMiddleware
99
- from fastapi.responses import JSONResponse
100
-
101
- # ── Logging ───────────────────────────────────────────────────────────────────
102
- logging.basicConfig(
103
- level=logging.INFO,
104
- format="%(asctime)s %(levelname)-8s %(name)s %(message)s",
105
- )
106
- logger = logging.getLogger("ocr-service")
107
-
108
- # ── Start time ────────────────────────────────────────────────────────────────
109
- _START_TIME: float = time.time()
110
-
111
- # ── Upload / batch limits ─────────────────────────────────────────────────��───
112
- MAX_UPLOAD_BYTES = 30 * 1024 * 1024 # 30 MB
113
- BATCH_MAX_FILES = 8
114
-
115
- # ── File type sets ────────────────────────────────────────────────────────────
116
- PDF_EXTENSIONS = {"pdf"}
117
- NATIVE_IMAGE_EXTENSIONS = {"jpg", "jpeg", "png"}
118
- PILLOW_IMAGE_EXTENSIONS = {"webp", "bmp", "tiff", "tif", "gif", "heic", "heif", "avif"}
119
- IMAGE_EXTENSIONS = NATIVE_IMAGE_EXTENSIONS | PILLOW_IMAGE_EXTENSIONS
120
- OFFICE_EXTENSIONS = {"docx", "pptx", "xlsx"}
121
- ALLOWED_EXTENSIONS = PDF_EXTENSIONS | IMAGE_EXTENSIONS | OFFICE_EXTENSIONS
122
-
123
- # ── OCR tuning ────────────────────────────────────────────────────────────────
124
- FAST_CONFIDENCE_THRESHOLD = 0.65 # below this → MinerU fallback
125
- MAX_OCR_SIDE = 1600 # pixels — longest side cap before OCR
126
-
127
- # ── Memory safety ─────────────────────────────────────────────────────────────
128
- BYTES_PER_OCR_PAGE = 100 * 1024 * 1024
129
- IMAGE_MEMORY_FACTOR = 4
130
- MEM_SAFETY_FLOOR_MB = 1024
131
-
132
- # ── SHA256 extraction cache ───────────────────────────────────────────────────
133
- _cache: dict[str, dict[str, Any]] = {}
134
- _cache_lock = threading.Lock()
135
-
136
- # ── Active-request counter ────────────────────────────────────────────────────
137
- _active_requests: int = 0
138
- _active_lock = threading.Lock()
139
-
140
- # ── Engine state ──────────────────────────────────────────────────────────────
141
- _rapidocr_engine: Any = None
142
- _rapidocr_lock = threading.Lock()
143
- _rapidocr_load_ms: int = 0
144
- _rapidocr_ready: bool = False
145
-
146
- _pipeline_ready: bool = False
147
- _pipeline_lock = threading.Lock()
148
- _pipeline_load_ms: int = 0
149
-
150
- # ── Startup issues ────────────────────────────────────────────────────────────
151
- _startup_issues: list[str] = []
152
- _startup_done: bool = False
153
-
154
-
155
- # ═════════════════════════════════════════════════════════════════════════════
156
- # Structured error
157
- # ═════════════════════════════════════════════════════════════════════════════
158
- class ExtractionError(Exception):
159
- def __init__(
160
- self,
161
- stage: str,
162
- code: str,
163
- message: str,
164
- http_status: int = 422,
165
- root_cause: str = "",
166
- recommendation: str = "",
167
- ) -> None:
168
- self.stage = stage
169
- self.code = code
170
- self.message = message
171
- self.http_status = http_status
172
- self.root_cause = root_cause or message
173
- self.recommendation = recommendation
174
- super().__init__(message)
175
-
176
- def to_dict(self) -> dict[str, Any]:
177
- return {
178
- "success": False,
179
- "stage": self.stage,
180
- "errorCode": self.code,
181
- "rootCause": self.root_cause,
182
- "recommendation": self.recommendation,
183
- "message": self.message,
184
- }
185
-
186
-
187
- def _err(
188
- stage: str,
189
- code: str,
190
- msg: str,
191
- status: int = 422,
192
- root_cause: str = "",
193
- recommendation: str = "",
194
- ) -> ExtractionError:
195
- return ExtractionError(stage, code, msg, status, root_cause, recommendation)
196
-
197
-
198
- # ═════════════════════════════════════════════════════════════════════════════
199
- # Active-request helpers
200
- # ═════════════════════════════════════════════════════════════════════════════
201
- def _inc_active() -> None:
202
- global _active_requests
203
- with _active_lock:
204
- _active_requests += 1
205
-
206
-
207
- def _dec_active() -> None:
208
- global _active_requests
209
- with _active_lock:
210
- _active_requests = max(0, _active_requests - 1)
211
-
212
-
213
- # ═════════════════════════════════════════════════════════════════════════════
214
- # Engine loaders
215
- # ═════════════════════════════════════════════════════════════════════════════
216
- def _ensure_rapidocr() -> Any:
217
- """Load the RapidOCR engine once; return the singleton on every subsequent call."""
218
- global _rapidocr_engine, _rapidocr_ready, _rapidocr_load_ms
219
- if _rapidocr_ready:
220
- return _rapidocr_engine
221
- with _rapidocr_lock:
222
- if _rapidocr_ready:
223
- return _rapidocr_engine
224
- t0 = time.perf_counter()
225
- try:
226
- from rapidocr_onnxruntime import RapidOCR
227
- _rapidocr_engine = RapidOCR(
228
- det_limit_side_len=MAX_OCR_SIDE,
229
- det_limit_type="max",
230
- )
231
- _rapidocr_load_ms = int((time.perf_counter() - t0) * 1000)
232
- _rapidocr_ready = True
233
- logger.info("RapidOCR engine ready load_ms=%d", _rapidocr_load_ms)
234
- except Exception as exc:
235
- raise _err(
236
- "model_load", "RAPIDOCR_LOAD_FAILED",
237
- f"RapidOCR failed to load: {exc}", 503,
238
- root_cause=str(exc),
239
- recommendation="Check that rapidocr-onnxruntime is installed.",
240
- ) from exc
241
- return _rapidocr_engine
242
-
243
-
244
- def _ensure_pipeline() -> None:
245
- """Import and verify the MinerU pipeline once."""
246
- global _pipeline_ready, _pipeline_load_ms
247
- if _pipeline_ready:
248
- return
249
- with _pipeline_lock:
250
- if _pipeline_ready:
251
- return
252
- config_path = os.path.expanduser("~/magic-pdf.json")
253
- if not os.path.exists(config_path):
254
- raise _err(
255
- "model_load", "CONFIG_MISSING",
256
- f"magic-pdf.json not found at {config_path}.", 503,
257
- root_cause="download_models.py did not run or /root was wiped.",
258
- recommendation="Check Docker build log for download_models.py output.",
259
- )
260
- t0 = time.perf_counter()
261
- try:
262
- from magic_pdf.data.dataset import PymuDocDataset, ImageDataset # noqa
263
- from magic_pdf.data.data_reader_writer import ( # noqa
264
- FileBasedDataReader, FileBasedDataWriter)
265
- except ImportError as exc:
266
- raise _err(
267
- "model_load", "IMPORT_FAILED",
268
- f"magic_pdf not importable: {exc}", 503,
269
- root_cause=str(exc),
270
- recommendation="Check that magic-pdf[full]==1.3.12 is installed.",
271
- ) from exc
272
- _pipeline_load_ms = int((time.perf_counter() - t0) * 1000)
273
- _pipeline_ready = True
274
- logger.info("MinerU pipeline ready load_ms=%d", _pipeline_load_ms)
275
-
276
-
277
- # ═════════════════════════════════════════════════════════════════════════════
278
- # FastAPI app
279
- # ═════════════════════════════════════════════════════════════════════════════
280
- app = FastAPI(
281
- title="OpenSkill OCR Service",
282
- description="OCR-only text extraction. Document understanding is handled by the AI layer.",
283
- version="4.0.0",
284
- )
285
-
286
- app.add_middleware(
287
- CORSMiddleware,
288
- allow_origins=["*"],
289
- allow_methods=["GET", "POST"],
290
- allow_headers=["*"],
291
- )
292
-
293
-
294
- # ─────────────────────────────────────────────────────────────────────────────
295
- # Startup — pre-load RapidOCR so first request has zero cold-start cost
296
- # ─────────────────────────────────────────────────────────────────────────────
297
- @app.on_event("startup")
298
- async def startup_warmup() -> None:
299
- """
300
- Pre-load the RapidOCR engine at container start.
301
-
302
- Without this, the first /extract request pays 600–2 500 ms for ONNX model
303
- loading on top of normal inference time. Loading here moves that cost to
304
- startup where it is invisible to the user.
305
- """
306
- global _startup_done
307
- issues: list[str] = []
308
-
309
- # ── Dependency smoke-check ────────────────────────────────────────────────
310
- checks = [
311
- ("cv2", lambda: __import__("cv2").__version__),
312
- ("torch", lambda: __import__("torch").__version__),
313
- ("rapidocr", lambda: pkg_version("rapidocr-onnxruntime")),
314
- ("magic_pdf", lambda: __import__("magic_pdf").__version__),
315
- ]
316
- for name, fn in checks:
317
- try:
318
- ver = fn()
319
- logger.info("startup ✓ %-12s %s", name, ver)
320
- except Exception as exc:
321
- msg = f"{name} unavailable: {exc}"
322
- issues.append(msg)
323
- logger.critical("startup FAIL %s", msg)
324
-
325
- if not os.path.exists(os.path.expanduser("~/magic-pdf.json")):
326
- issues.append("magic-pdf.json missing")
327
- if not os.path.isdir("/app/models/PDF-Extract-Kit-1.0/models"):
328
- issues.append("Models directory missing: /app/models/PDF-Extract-Kit-1.0/models")
329
-
330
- # ── Pre-load RapidOCR ─────────────────────────────────────────────────────
331
- try:
332
- _ensure_rapidocr()
333
- logger.info("startup: RapidOCR pre-loaded load_ms=%d", _rapidocr_load_ms)
334
- except Exception as exc:
335
- msg = f"RapidOCR warmup failed: {exc}"
336
- issues.append(msg)
337
- logger.error("startup: %s", msg)
338
-
339
- _startup_issues.extend(issues)
340
- _startup_done = True
341
- if issues:
342
- logger.error("Startup completed with %d issue(s): %s", len(issues), issues)
343
- else:
344
- logger.info("Startup complete — all systems ready.")
345
-
346
-
347
- # ═════════════════════════════════════════════════════════════════════════════
348
- # GET /health
349
- # ═════════════════════════════════════════════════════════════════════════════
350
- @app.get("/health")
351
- def health() -> dict[str, Any]:
352
- return {"status": "healthy", "version": "4.0.0"}
353
-
354
-
355
- # ═════════════════════════════════════════════════════════════════════════════
356
- # GET /status
357
- # ═════════════════════════════════════════════════════════════════════════════
358
- @app.get("/status")
359
- def status() -> dict[str, Any]:
360
- used_mb, total_mb = _mem_mb()
361
- return {
362
- "status": "healthy" if not _startup_issues else "degraded",
363
- "version": "4.0.0",
364
- "architecture": "ocr-only",
365
- "engines": {
366
- "rapidocr": {
367
- "ready": _rapidocr_ready,
368
- "loadMs": _rapidocr_load_ms,
369
- "purpose": "images (1–4 s)",
370
- },
371
- "mineru": {
372
- "ready": _pipeline_ready,
373
- "loadMs": _pipeline_load_ms,
374
- "purpose": "PDFs + fallback",
375
- },
376
- },
377
- "config": {
378
- "maxOcrSidePx": MAX_OCR_SIDE,
379
- "confidenceThreshold": FAST_CONFIDENCE_THRESHOLD,
380
- "maxUploadMb": MAX_UPLOAD_BYTES // (1024 * 1024),
381
- },
382
- "startupIssues": _startup_issues,
383
- "uptimeSeconds": int(time.time() - _START_TIME),
384
- "memoryUsedMB": used_mb,
385
- "memoryTotalMB": total_mb,
386
- "activeRequests": _active_requests,
387
- "cacheEntries": len(_cache),
388
- }
389
-
390
-
391
- # ═════════════════════════════════════════════════════════════════════════════
392
- # GET /warmup
393
- # ═════════════════════════════════════════════════════════════════════════════
394
- @app.get("/warmup")
395
- def warmup() -> dict[str, Any]:
396
- """Explicitly pre-load engines. Idempotent — safe to call repeatedly."""
397
- results: dict[str, Any] = {}
398
- t0 = time.perf_counter()
399
- try:
400
- _ensure_rapidocr()
401
- results["rapidocr"] = {"status": "ready", "loadMs": _rapidocr_load_ms}
402
- except Exception as exc:
403
- results["rapidocr"] = {"status": "failed", "error": str(exc)}
404
- try:
405
- _ensure_pipeline()
406
- results["mineru"] = {"status": "ready", "loadMs": _pipeline_load_ms}
407
- except Exception as exc:
408
- results["mineru"] = {"status": "failed", "error": str(exc)}
409
- results["totalElapsedMs"] = int((time.perf_counter() - t0) * 1000)
410
- results["allReady"] = _rapidocr_ready and _pipeline_ready
411
- return results
412
-
413
-
414
- # ═════════════════════════════════════════════════════════════════════════════
415
- # GET /diagnostics
416
- # ═══════════════════��═════════════════════════════════════════════════════════
417
- @app.get("/diagnostics")
418
- def diagnostics() -> dict[str, Any]:
419
- import platform
420
- pkgs: dict[str, str] = {}
421
- for name in (
422
- "magic-pdf", "rapidocr-onnxruntime", "torch", "torchvision",
423
- "ultralytics", "doclayout-yolo", "rapid-table", "onnxruntime",
424
- "opencv-python-headless", "Pillow", "fastapi", "uvicorn",
425
- ):
426
- try:
427
- pkgs[name] = pkg_version(name)
428
- except Exception:
429
- pkgs[name] = "not found"
430
-
431
- models_root = "/app/models/PDF-Extract-Kit-1.0/models"
432
- model_files: dict[str, str] = {}
433
- for rel in [
434
- "OCR/paddleocr_torch/ch_PP-OCRv5_det_infer.pth",
435
- "OCR/paddleocr_torch/ch_PP-OCRv5_rec_infer.pth",
436
- "Layout/YOLO/doclayout_yolo_docstructbench_imgsz1280_2501.pt",
437
- ]:
438
- full = os.path.join(models_root, rel)
439
- model_files[rel] = (
440
- f"{os.path.getsize(full) / (1024 * 1024):.1f} MB"
441
- if os.path.isfile(full) else "MISSING"
442
- )
443
-
444
- used_mb, total_mb = _mem_mb()
445
- return {
446
- "python": platform.python_version(),
447
- "packages": pkgs,
448
- "modelFiles": model_files,
449
- "memory": {"usedMB": used_mb, "totalMB": total_mb},
450
- "engines": {
451
- "rapidocr": {"ready": _rapidocr_ready, "loadMs": _rapidocr_load_ms},
452
- "mineru": {"ready": _pipeline_ready, "loadMs": _pipeline_load_ms},
453
- },
454
- "config": {
455
- "maxOcrSidePx": MAX_OCR_SIDE,
456
- "confidenceThreshold": FAST_CONFIDENCE_THRESHOLD,
457
- },
458
- "uptime": int(time.time() - _START_TIME),
459
- "cacheEntries": len(_cache),
460
- }
461
-
462
-
463
- # ═════════════════════════════════════════════════════════════════════════════
464
- # GET /benchmark
465
- # Runs RapidOCR on three synthetic images (small / medium / large) and returns
466
- # full stage timings for each. Use this to measure the resize optimisation.
467
- # ═════════════════════════════════════════════════════════════════════════════
468
- @app.get("/benchmark")
469
- async def benchmark() -> JSONResponse:
470
- import cv2
471
-
472
- def _make_test_image(width: int, height: int) -> "np.ndarray":
473
- img = np.ones((height, width, 3), dtype=np.uint8) * 255
474
- lines = [
475
- "184 ENGLISH LNG & LIT. 073 020 093",
476
- "085 HINDI COURSE-B 075 020 095",
477
- "041 MATHEMATICS STD 063 020 083",
478
- "086 SCIENCE 065 020 085",
479
- "087 SOCIAL SCIENCE 057 020 077",
480
- "Roll No: 28169763 Name: TEST STUDENT",
481
- "Total: 433 / 500 Percentage: 86.6%",
482
- ]
483
- line_h = max(20, height // (len(lines) + 2))
484
- scale = max(0.5, min(1.5, width / 900))
485
- for i, text in enumerate(lines):
486
- y = line_h * (i + 1)
487
- if y < height - 10:
488
- cv2.putText(img, text, (20, y),
489
- cv2.FONT_HERSHEY_SIMPLEX, scale, (0, 0, 0), 2)
490
- return img
491
-
492
- SIZES = [
493
- ("small", 800, 1200),
494
- ("medium", 1600, 2400),
495
- ("large", 3000, 4000),
496
- ]
497
- results: dict[str, Any] = {}
498
-
499
- engine = _ensure_rapidocr()
500
-
501
- for label, w, h in SIZES:
502
- img = _make_test_image(w, h)
503
- orig_h, orig_w = img.shape[:2]
504
-
505
- # Resize
506
- t_resize = time.perf_counter()
507
- img_resized, was_resized = _resize_for_ocr(img)
508
- resize_ms = int((time.perf_counter() - t_resize) * 1000)
509
- new_h, new_w = img_resized.shape[:2]
510
-
511
- # OCR
512
- t_ocr = time.perf_counter()
513
- ocr_result, elapse = engine(img_resized)
514
- ocr_ms = int((time.perf_counter() - t_ocr) * 1000)
515
-
516
- det_ms, rec_ms = _split_elapse(elapse, ocr_ms)
517
- texts = [item[1] for item in (ocr_result or []) if len(item) > 1]
518
- scores = [item[2] for item in (ocr_result or []) if len(item) > 2 and item[2] is not None]
519
- conf = round(sum(scores) / len(scores), 4) if scores else 0.0
520
-
521
- results[label] = {
522
- "originalDimensions": f"{orig_w}×{orig_h}",
523
- "resizedDimensions": f"{new_w}×{new_h}",
524
- "wasResized": was_resized,
525
- "resizeMs": resize_ms,
526
- "detectMs": det_ms,
527
- "recognizeMs": rec_ms,
528
- "ocrTotalMs": ocr_ms,
529
- "textBlocks": len(texts),
530
- "confidence": conf,
531
- }
532
-
533
- used_mb, total_mb = _mem_mb()
534
- return JSONResponse(content={
535
- "results": results,
536
- "memory": {"usedMB": used_mb, "totalMB": total_mb},
537
- "maxOcrSide": MAX_OCR_SIDE,
538
- "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
539
- })
540
-
541
-
542
- # ═════════════════════════════════════════════════════════════════════════════
543
- # POST /extract
544
- # ═════════════════════════════════════════════════════════════════════════════
545
- @app.post("/extract")
546
- async def extract(file: UploadFile = File(...)) -> JSONResponse:
547
- t_upload_start = time.perf_counter()
548
- try:
549
- raw, filename, ext = await _read_upload(file)
550
- upload_ms = int((time.perf_counter() - t_upload_start) * 1000)
551
- result = _run_extraction(raw, filename, ext, upload_ms=upload_ms)
552
- return JSONResponse(content=result)
553
- except ExtractionError as exc:
554
- logger.warning("/extract [%s/%s]: %s", exc.stage, exc.code, exc.message)
555
- return JSONResponse(status_code=exc.http_status, content=exc.to_dict())
556
- except Exception as exc:
557
- logger.exception("/extract unhandled error")
558
- return JSONResponse(
559
- status_code=500,
560
- content={
561
- "success": False,
562
- "stage": "unknown",
563
- "errorCode": "INTERNAL_ERROR",
564
- "rootCause": str(exc),
565
- "recommendation": "Check HF Space logs for full traceback.",
566
- "message": str(exc),
567
- "traceback": traceback.format_exc()[-3000:],
568
- },
569
- )
570
-
571
-
572
- # ═════════════════════════════════════════════════════════════════════════════
573
- # POST /batch
574
- # ═════════════════════════════════════════════════════════════════════════════
575
- @app.post("/batch")
576
- async def batch(files: list[UploadFile] = File(...)) -> JSONResponse:
577
- candidates = files[:BATCH_MAX_FILES]
578
- results: list[dict[str, Any]] = []
579
- for upload in candidates:
580
- t0 = time.perf_counter()
581
- try:
582
- raw, filename, ext = await _read_upload(upload)
583
- result = _run_extraction(
584
- raw, filename, ext,
585
- upload_ms=int((time.perf_counter() - t0) * 1000),
586
- )
587
- except ExtractionError as exc:
588
- result = exc.to_dict()
589
- result["filename"] = _sanitize_filename(upload.filename or "upload")
590
- except Exception as exc:
591
- fname = _sanitize_filename(upload.filename or "upload")
592
- logger.exception("Batch item failed: %s", fname)
593
- result = {
594
- "success": False,
595
- "filename": fname,
596
- "stage": "unknown",
597
- "errorCode": "INTERNAL_ERROR",
598
- "rootCause": str(exc),
599
- "recommendation": "Check HF Space logs.",
600
- "message": str(exc),
601
- }
602
- results.append(result)
603
- return JSONResponse(content={
604
- "success": True,
605
- "processed": len(results),
606
- "results": results,
607
- })
608
-
609
-
610
- # ═════════════════════════════════════════════════════════════════════════════
611
- # Upload reader
612
- # ═════════════════════════════════════════════════════════════════════════════
613
- async def _read_upload(upload: UploadFile) -> tuple[bytes, str, str]:
614
- filename = _sanitize_filename(upload.filename or "upload")
615
- ext = filename.rsplit(".", 1)[-1].lower() if "." in filename else ""
616
-
617
- if ext not in ALLOWED_EXTENSIONS:
618
- raise _err(
619
- "validation", "UNSUPPORTED_TYPE",
620
- f"Unsupported file type '.{ext}'. "
621
- f"Supported: {sorted(ALLOWED_EXTENSIONS)}",
622
- 415,
623
- root_cause=f"Extension '{ext}' is not in the allowed set.",
624
- recommendation="Convert to PDF, JPG, PNG, or WEBP before uploading.",
625
- )
626
- raw = await upload.read(MAX_UPLOAD_BYTES + 1)
627
- if len(raw) > MAX_UPLOAD_BYTES:
628
- raise _err(
629
- "upload", "FILE_TOO_LARGE",
630
- f"'{filename}' exceeds {MAX_UPLOAD_BYTES // 1024 // 1024} MB.", 413,
631
- root_cause=f"File is {len(raw) // 1024 // 1024} MB.",
632
- recommendation="Compress or split the file.",
633
- )
634
- if len(raw) == 0:
635
- raise _err("upload", "EMPTY_FILE", f"'{filename}' is empty.", 400,
636
- root_cause="Zero bytes received.",
637
- recommendation="Check the file before uploading.")
638
- return raw, filename, ext
639
-
640
-
641
- # ═════════════════════════════════════════════════════════════════════════════
642
- # Extraction dispatcher
643
- # ═════════════════════════════════════════════════════════════════════════════
644
- def _run_extraction(
645
- raw: bytes, filename: str, ext: str, upload_ms: int = 0
646
- ) -> dict[str, Any]:
647
- logger.info("request_received file=%s size=%d ext=%s", filename, len(raw), ext)
648
-
649
- # ── Hash + cache lookup ───────────────────────────────────────────────────
650
- t_hash = time.perf_counter()
651
- file_hash = hashlib.sha256(raw).hexdigest()
652
- hash_ms = int((time.perf_counter() - t_hash) * 1000)
653
- logger.info("cache_lookup sha256=%.12s… hash_ms=%d", file_hash, hash_ms)
654
-
655
- with _cache_lock:
656
- cached = _cache.get(file_hash)
657
- if cached is not None:
658
- logger.info("cache_hit sha256=%.12s… file=%s", file_hash, filename)
659
- out = {**cached}
660
- out["cached"] = True
661
- out["processingTimeMs"] = 0
662
- out["timings"] = {**cached.get("timings", {}), "totalMs": 0}
663
- return out
664
-
665
- logger.info("cache_miss sha256=%.12s…", file_hash)
666
-
667
- # ── Memory safety ─────────────────────────────────────────────────────────
668
- t_mem = time.perf_counter()
669
- _assert_memory_safe(raw, ext)
670
- mem_check_ms = int((time.perf_counter() - t_mem) * 1000)
671
-
672
- _inc_active()
673
- work_dir = tempfile.mkdtemp(prefix="ocr_")
674
- t0 = time.perf_counter()
675
- try:
676
- if ext in PDF_EXTENSIONS:
677
- logger.info("engine_selected engine=mineru file=%s", filename)
678
- _ensure_pipeline()
679
- result = _process_pdf(raw, filename, work_dir, upload_ms=upload_ms)
680
- elif ext in OFFICE_EXTENSIONS:
681
- logger.info("engine_selected engine=office_text file=%s ext=%s", filename, ext)
682
- result = _process_office(raw, filename, ext, upload_ms=upload_ms)
683
- else:
684
- logger.info("engine_selected engine=rapidocr file=%s", filename)
685
- result = _process_image(raw, filename, ext, work_dir, upload_ms=upload_ms)
686
-
687
- total_ms = int((time.perf_counter() - t0) * 1000)
688
- result["timings"]["uploadMs"] = upload_ms
689
- result["timings"]["hashMs"] = hash_ms
690
- result["timings"]["memCheckMs"] = mem_check_ms
691
- result["timings"]["totalMs"] = total_ms
692
- result["processingTimeMs"] = total_ms
693
- result["cached"] = False
694
-
695
- # Store in cache (strip per-request fields that change on replay)
696
- entry = {k: v for k, v in result.items()
697
- if k not in ("cached", "processingTimeMs", "timings")}
698
- entry["timings"] = {k: v for k, v in result["timings"].items()
699
- if k not in ("totalMs", "hashMs", "memCheckMs", "uploadMs")}
700
- with _cache_lock:
701
- _cache[file_hash] = entry
702
-
703
- logger.info(
704
- "response_sent file=%s engine=%s conf=%.3f total_ms=%d",
705
- filename, result.get("engine", "?"), result.get("confidence", 0), total_ms,
706
- )
707
- return result
708
-
709
- except ExtractionError:
710
- raise
711
- except Exception as exc:
712
- logger.exception("extraction_failed file=%s", filename)
713
- raise _err(
714
- "unknown", "INTERNAL_ERROR", f"Unexpected error: {exc}", 500,
715
- root_cause=str(exc),
716
- recommendation="Check HF Space logs for full traceback.",
717
- ) from exc
718
- finally:
719
- _dec_active()
720
- shutil.rmtree(work_dir, ignore_errors=True)
721
-
722
-
723
- # ═════════════════════════════════════════════════════════════════════════════
724
- # Image processor — RapidOCR fast path + MinerU fallback
725
- # ═════════════════════════════════════════════════════════════════════════════
726
- def _process_image(
727
- raw: bytes, filename: str, ext: str, work_dir: str, upload_ms: int = 0
728
- ) -> dict[str, Any]:
729
- import cv2
730
-
731
- # ── Decode ────────────────────────────────────────────────────��───────────
732
- t_decode = time.perf_counter()
733
- img_bgr = _decode_image_to_bgr(raw, ext)
734
- decode_ms = int((time.perf_counter() - t_decode) * 1000)
735
- orig_h, orig_w = img_bgr.shape[:2]
736
- logger.info("image_decoded file=%s dims=%dx%d decode_ms=%d",
737
- filename, orig_w, orig_h, decode_ms)
738
-
739
- # ── Resize ────────────────────────────────────────────────────────────────
740
- t_resize = time.perf_counter()
741
- img_ocr, was_resized = _resize_for_ocr(img_bgr)
742
- resize_ms = int((time.perf_counter() - t_resize) * 1000)
743
- new_h, new_w = img_ocr.shape[:2]
744
- logger.info("image_resized file=%s original=%dx%d resized=%dx%d"
745
- " was_resized=%s resize_ms=%d",
746
- filename, orig_w, orig_h, new_w, new_h, was_resized, resize_ms)
747
-
748
- # ── RapidOCR ──────────────────────────────────────────────────────────────
749
- logger.info("ocr_started file=%s engine=rapidocr dims=%dx%d",
750
- filename, new_w, new_h)
751
- t_ocr = time.perf_counter()
752
- try:
753
- engine = _ensure_rapidocr()
754
- ocr_result, elapse = engine(img_ocr)
755
- except ExtractionError:
756
- raise
757
- except Exception as exc:
758
- raise _err(
759
- "ocr", "OCR_ENGINE_FAILED", f"RapidOCR failed: {exc}", 500,
760
- root_cause=str(exc),
761
- recommendation="Check rapidocr-onnxruntime in Dockerfile Layer 1.",
762
- ) from exc
763
- ocr_ms = int((time.perf_counter() - t_ocr) * 1000)
764
- det_ms, rec_ms = _split_elapse(elapse, ocr_ms)
765
- logger.info("ocr_finished file=%s engine=rapidocr ocr_ms=%d"
766
- " det_ms=%d rec_ms=%d", filename, ocr_ms, det_ms, rec_ms)
767
-
768
- # ── Parse output ──────────────────────────────────────────────────────────
769
- t_post = time.perf_counter()
770
- plain_text, confidence = _parse_rapidocr_output(ocr_result)
771
- post_ms = int((time.perf_counter() - t_post) * 1000)
772
- logger.info("post_process file=%s conf=%.3f text_len=%d blocks=%d post_ms=%d",
773
- filename, confidence, len(plain_text),
774
- len(ocr_result) if ocr_result else 0, post_ms)
775
-
776
- # ── MinerU fallback if confidence is low ──────────────────────────────────
777
- passes_used = 1
778
- engine_name = "rapidocr"
779
- if confidence < FAST_CONFIDENCE_THRESHOLD and plain_text.strip():
780
- logger.info(
781
- "fallback_triggered conf=%.3f < %.2f file=%s trying mineru",
782
- confidence, FAST_CONFIDENCE_THRESHOLD, filename,
783
- )
784
- try:
785
- _ensure_pipeline()
786
- mr = _process_image_mineru(raw, filename, ext, work_dir)
787
- if len(mr.get("text", "")) > len(plain_text) * 0.8:
788
- mr["engine"] = "mineru_fallback"
789
- mr["metadata"]["passesUsed"] = 2
790
- mr["timings"]["pass1RapidOCRMs"] = ocr_ms
791
- mr["timings"]["decodeMs"] = decode_ms
792
- mr["timings"]["resizeMs"] = resize_ms
793
- logger.info("fallback_used file=%s mineru result accepted", filename)
794
- return mr
795
- except Exception as exc:
796
- logger.warning("fallback_failed file=%s error=%s using rapidocr result", filename, exc)
797
- passes_used = 2
798
- else:
799
- logger.info("fallback_not_needed conf=%.3f file=%s", confidence, filename)
800
-
801
- return {
802
- "success": True,
803
- "filename": filename,
804
- "engine": engine_name,
805
- "confidence": confidence,
806
- "text": plain_text,
807
- "markdown": plain_text,
808
- "pageCount": 1,
809
- "timings": {
810
- "uploadMs": upload_ms,
811
- "hashMs": 0,
812
- "memCheckMs": 0,
813
- "decodeMs": decode_ms,
814
- "resizeMs": resize_ms,
815
- "detectMs": det_ms,
816
- "recognizeMs": rec_ms,
817
- "postProcessMs": post_ms,
818
- "totalMs": 0,
819
- },
820
- "metadata": {
821
- "imgW": orig_w,
822
- "imgH": orig_h,
823
- "imgWResized": new_w,
824
- "imgHResized": new_h,
825
- "wasResized": was_resized,
826
- "textBlocks": len(ocr_result) if ocr_result else 0,
827
- "passesUsed": passes_used,
828
- "backend": "rapidocr",
829
- },
830
- }
831
-
832
-
833
- def _process_image_mineru(
834
- raw: bytes, filename: str, ext: str, work_dir: str
835
- ) -> dict[str, Any]:
836
- from magic_pdf.data.data_reader_writer import (
837
- FileBasedDataReader, FileBasedDataWriter)
838
- from magic_pdf.data.dataset import ImageDataset
839
- from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
840
-
841
- images_dir = os.path.join(work_dir, "images_mineru")
842
- os.makedirs(images_dir, exist_ok=True)
843
-
844
- if ext in PILLOW_IMAGE_EXTENSIONS:
845
- raw = _convert_to_png(raw, ext)
846
- save_ext = "png"
847
- else:
848
- save_ext = ext
849
-
850
- img_path = os.path.join(work_dir, f"input_mineru.{save_ext}")
851
- with open(img_path, "wb") as fh:
852
- fh.write(raw)
853
-
854
- t_ocr = time.perf_counter()
855
- try:
856
- reader = FileBasedDataReader(work_dir)
857
- image_bytes = reader.read(f"input_mineru.{save_ext}")
858
- ds = ImageDataset(image_bytes)
859
- infer_result = ds.apply(doc_analyze, ocr=True)
860
- pipe_result = infer_result.pipe_ocr_mode(FileBasedDataWriter(images_dir))
861
- except Exception as exc:
862
- raise _err(
863
- "ocr", "OCR_PIPELINE_FAILED",
864
- f"MinerU image pipeline failed: {exc}", 500,
865
- root_cause=str(exc),
866
- recommendation="Check magic-pdf installation and model files.",
867
- ) from exc
868
- ocr_ms = int((time.perf_counter() - t_ocr) * 1000)
869
-
870
- t_md = time.perf_counter()
871
- try:
872
- markdown = pipe_result.get_markdown(images_dir)
873
- except Exception as exc:
874
- raise _err("markdown", "MARKDOWN_FAILED", f"get_markdown failed: {exc}") from exc
875
- md_ms = int((time.perf_counter() - t_md) * 1000)
876
-
877
- plain_text = _markdown_to_plain(markdown)
878
-
879
- return {
880
- "success": True,
881
- "filename": filename,
882
- "engine": "mineru",
883
- "confidence": 0.85,
884
- "text": plain_text,
885
- "markdown": markdown,
886
- "pageCount": 1,
887
- "timings": {
888
- "uploadMs": 0,
889
- "hashMs": 0,
890
- "memCheckMs": 0,
891
- "decodeMs": 0,
892
- "resizeMs": 0,
893
- "detectMs": 0,
894
- "recognizeMs": ocr_ms,
895
- "postProcessMs": md_ms,
896
- "totalMs": 0,
897
- },
898
- "metadata": {
899
- "imgW": 0, "imgH": 0,
900
- "imgWResized": 0, "imgHResized": 0,
901
- "wasResized": False,
902
- "textBlocks": 0,
903
- "passesUsed": 1,
904
- "backend": "pipeline",
905
- },
906
- }
907
-
908
-
909
- # ═════════════════════════════════════════════════════════════════════════════
910
- # Office document processor — DOCX / PPTX / XLSX (text extraction, no OCR)
911
- # No image rendering or OCR is performed. Text is read directly from the
912
- # structured XML inside the Office Open XML container.
913
- # ═════════════════════════════════════════════════════════════════════════════
914
- def _process_office(
915
- raw: bytes, filename: str, ext: str, upload_ms: int = 0
916
- ) -> dict[str, Any]:
917
- t0 = time.perf_counter()
918
- logger.info("ocr_started file=%s engine=office_text ext=%s", filename, ext)
919
-
920
- try:
921
- if ext == "docx":
922
- plain_text, page_count = _extract_docx(raw)
923
- elif ext == "pptx":
924
- plain_text, page_count = _extract_pptx(raw)
925
- elif ext == "xlsx":
926
- plain_text, page_count = _extract_xlsx(raw)
927
- else:
928
- raise _err("decode", "UNSUPPORTED_OFFICE_TYPE",
929
- f"Unrecognised office extension: {ext}", 415)
930
- except ExtractionError:
931
- raise
932
- except Exception as exc:
933
- raise _err(
934
- "ocr", "OFFICE_EXTRACT_FAILED",
935
- f"Could not extract text from {ext.upper()}: {exc}", 422,
936
- root_cause=str(exc),
937
- recommendation=f"Ensure the file is a valid, non-password-protected {ext.upper()}.",
938
- ) from exc
939
-
940
- extract_ms = int((time.perf_counter() - t0) * 1000)
941
- logger.info("ocr_finished file=%s engine=office_text extract_ms=%d text_len=%d",
942
- filename, extract_ms, len(plain_text))
943
-
944
- return {
945
- "success": True,
946
- "filename": filename,
947
- "engine": f"office_text_{ext}",
948
- "confidence": 1.0,
949
- "text": plain_text,
950
- "markdown": plain_text,
951
- "pageCount": page_count,
952
- "timings": {
953
- "uploadMs": upload_ms,
954
- "hashMs": 0,
955
- "memCheckMs": 0,
956
- "decodeMs": 0,
957
- "resizeMs": 0,
958
- "detectMs": 0,
959
- "recognizeMs": extract_ms,
960
- "postProcessMs": 0,
961
- "totalMs": 0,
962
- },
963
- "metadata": {
964
- "imgW": 0, "imgH": 0,
965
- "imgWResized": 0, "imgHResized": 0,
966
- "wasResized": False,
967
- "textBlocks": plain_text.count("\n") + 1,
968
- "passesUsed": 1,
969
- "backend": f"office_text_{ext}",
970
- },
971
- }
972
-
973
-
974
- def _extract_docx(raw: bytes) -> tuple[str, int]:
975
- """Extract plain text from a DOCX file. Returns (text, page_estimate)."""
976
- try:
977
- import docx as _docx
978
- except ImportError as exc:
979
- raise _err("decode", "DOCX_DEPS_MISSING",
980
- "python-docx is not installed.", 503,
981
- recommendation="Add python-docx to Dockerfile Layer 1.") from exc
982
- doc = _docx.Document(io.BytesIO(raw))
983
- paragraphs = [p.text for p in doc.paragraphs if p.text.strip()]
984
- # Tables
985
- for table in doc.tables:
986
- for row in table.rows:
987
- row_text = " | ".join(
988
- cell.text.strip() for cell in row.cells if cell.text.strip()
989
- )
990
- if row_text:
991
- paragraphs.append(row_text)
992
- text = "\n".join(paragraphs)
993
- # Rough page estimate: ~3 000 chars per page
994
- pages = max(1, len(text) // 3000)
995
- return text, pages
996
-
997
-
998
- def _extract_pptx(raw: bytes) -> tuple[str, int]:
999
- """Extract plain text from a PPTX file. Returns (text, slide_count)."""
1000
- try:
1001
- from pptx import Presentation as _Presentation
1002
- except ImportError as exc:
1003
- raise _err("decode", "PPTX_DEPS_MISSING",
1004
- "python-pptx is not installed.", 503,
1005
- recommendation="Add python-pptx to Dockerfile Layer 1.") from exc
1006
- prs = _Presentation(io.BytesIO(raw))
1007
- lines: list[str] = []
1008
- for slide_num, slide in enumerate(prs.slides, 1):
1009
- lines.append(f"--- Slide {slide_num} ---")
1010
- for shape in slide.shapes:
1011
- if hasattr(shape, "text") and shape.text.strip():
1012
- lines.append(shape.text.strip())
1013
- return "\n".join(lines), len(prs.slides)
1014
-
1015
-
1016
- def _extract_xlsx(raw: bytes) -> tuple[str, int]:
1017
- """Extract plain text from an XLSX file. Returns (text, sheet_count)."""
1018
- try:
1019
- import openpyxl as _openpyxl
1020
- except ImportError as exc:
1021
- raise _err("decode", "XLSX_DEPS_MISSING",
1022
- "openpyxl is not installed.", 503,
1023
- recommendation="Add openpyxl to Dockerfile Layer 1.") from exc
1024
- wb = _openpyxl.load_workbook(io.BytesIO(raw), read_only=True, data_only=True)
1025
- lines: list[str] = []
1026
- for sheet in wb.worksheets:
1027
- lines.append(f"--- Sheet: {sheet.title} ---")
1028
- for row in sheet.iter_rows(values_only=True):
1029
- row_text = " | ".join(
1030
- str(cell) for cell in row if cell is not None and str(cell).strip()
1031
- )
1032
- if row_text:
1033
- lines.append(row_text)
1034
- wb.close()
1035
- return "\n".join(lines), len(wb.worksheets)
1036
-
1037
-
1038
- # ═════════════════════════════════════════════════════════════════════════════
1039
- # PDF processor — MinerU
1040
- # ═════════════════════════════════════════════════════════════════════════════
1041
- def _process_pdf(
1042
- raw: bytes, filename: str, work_dir: str, upload_ms: int = 0
1043
- ) -> dict[str, Any]:
1044
- from magic_pdf.data.data_reader_writer import FileBasedDataWriter
1045
- from magic_pdf.data.dataset import PymuDocDataset
1046
- from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
1047
- from magic_pdf.config.enums import SupportedPdfParseMethod
1048
-
1049
- images_dir = os.path.join(work_dir, "images")
1050
- os.makedirs(images_dir, exist_ok=True)
1051
- page_count = _pdf_page_count(raw)
1052
-
1053
- logger.info("pdf_classify file=%s pages=%d", filename, page_count)
1054
- t_classify = time.perf_counter()
1055
- try:
1056
- ds = PymuDocDataset(raw)
1057
- method = ds.classify()
1058
- except Exception as exc:
1059
- raise _err(
1060
- "decode", "PDF_PARSE_FAILED", f"Could not parse PDF: {exc}", 422,
1061
- root_cause=str(exc),
1062
- recommendation="Ensure the file is a valid, non-encrypted PDF.",
1063
- ) from exc
1064
- classify_ms = int((time.perf_counter() - t_classify) * 1000)
1065
-
1066
- logger.info("ocr_started file=%s engine=mineru method=%s", filename, method)
1067
- t_ocr = time.perf_counter()
1068
- try:
1069
- image_writer = FileBasedDataWriter(images_dir)
1070
- if method == SupportedPdfParseMethod.TXT:
1071
- infer_result = ds.apply(doc_analyze, ocr=False)
1072
- pipe_result = infer_result.pipe_txt_mode(image_writer)
1073
- parse_method = "txt"
1074
- else:
1075
- infer_result = ds.apply(doc_analyze, ocr=True)
1076
- pipe_result = infer_result.pipe_ocr_mode(image_writer)
1077
- parse_method = "ocr"
1078
- except Exception as exc:
1079
- raise _err(
1080
- "ocr", "OCR_PIPELINE_FAILED", f"doc_analyze/pipe failed: {exc}", 500,
1081
- root_cause=str(exc),
1082
- recommendation="Check model files in /app/models and validate.py output.",
1083
- ) from exc
1084
- ocr_ms = int((time.perf_counter() - t_ocr) * 1000)
1085
- logger.info("ocr_finished file=%s engine=mineru ocr_ms=%d", filename, ocr_ms)
1086
-
1087
- t_md = time.perf_counter()
1088
- try:
1089
- markdown = pipe_result.get_markdown(images_dir)
1090
- except Exception as exc:
1091
- raise _err("markdown", "MARKDOWN_FAILED", f"get_markdown failed: {exc}") from exc
1092
- md_ms = int((time.perf_counter() - t_md) * 1000)
1093
-
1094
- plain_text = _markdown_to_plain(markdown)
1095
-
1096
- return {
1097
- "success": True,
1098
- "filename": filename,
1099
- "engine": "mineru",
1100
- "confidence": 0.9 if parse_method == "txt" else 0.85,
1101
- "text": plain_text,
1102
- "markdown": markdown,
1103
- "pageCount": page_count,
1104
- "timings": {
1105
- "uploadMs": upload_ms,
1106
- "hashMs": 0,
1107
- "memCheckMs": 0,
1108
- "decodeMs": classify_ms,
1109
- "resizeMs": 0,
1110
- "detectMs": 0,
1111
- "recognizeMs": ocr_ms,
1112
- "postProcessMs": md_ms,
1113
- "totalMs": 0,
1114
- },
1115
- "metadata": {
1116
- "imgW": 0, "imgH": 0,
1117
- "imgWResized": 0, "imgHResized": 0,
1118
- "wasResized": False,
1119
- "textBlocks": 0,
1120
- "passesUsed": 1,
1121
- "backend": "pipeline",
1122
- "parseMethod": parse_method,
1123
- "pages": page_count,
1124
- },
1125
- }
1126
-
1127
-
1128
- # ═════════════════════════════════════════════════════════════════════════════
1129
- # Image helpers
1130
- # ═════════════════════════════════════════════════════════════════════════════
1131
- def _resize_for_ocr(img: "np.ndarray") -> tuple["np.ndarray", bool]:
1132
- """
1133
- Resize image so the longest side is at most MAX_OCR_SIDE pixels.
1134
-
1135
- Returns (resized_img, was_resized).
1136
-
1137
- Uses cv2.INTER_AREA which is the correct algorithm for downscaling:
1138
- it averages pixels (anti-aliasing) rather than sampling individual pixels,
1139
- preserving text legibility at smaller sizes.
1140
-
1141
- No upscaling: images smaller than MAX_OCR_SIDE are returned unchanged.
1142
- """
1143
- import cv2
1144
- h, w = img.shape[:2]
1145
- longest = max(h, w)
1146
- if longest <= MAX_OCR_SIDE:
1147
- return img, False
1148
- scale = MAX_OCR_SIDE / longest
1149
- new_w = int(w * scale)
1150
- new_h = int(h * scale)
1151
- resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
1152
- return resized, True
1153
-
1154
-
1155
- def _decode_image_to_bgr(raw: bytes, ext: str) -> "np.ndarray":
1156
- import cv2
1157
- if ext in {"heic", "heif"}:
1158
- try:
1159
- from pillow_heif import register_heif_opener
1160
- register_heif_opener()
1161
- except ImportError:
1162
- raise _err(
1163
- "decode", "HEIF_NOT_SUPPORTED",
1164
- "HEIC/HEIF requires pillow-heif.", 415,
1165
- recommendation="Add pillow-heif to Dockerfile Layer 1.",
1166
- )
1167
- try:
1168
- pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
1169
- buf = io.BytesIO()
1170
- pil_img.save(buf, format="PNG")
1171
- raw = buf.getvalue()
1172
- except Exception as exc:
1173
- raise _err("decode", "HEIF_DECODE_FAILED",
1174
- f"HEIF decode error: {exc}") from exc
1175
-
1176
- arr = np.frombuffer(raw, np.uint8)
1177
- img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
1178
- if img is None:
1179
- try:
1180
- pil_img = Image.open(io.BytesIO(raw)).convert("RGB")
1181
- img = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
1182
- except Exception as exc:
1183
- raise _err(
1184
- "decode", "IMAGE_DECODE_FAILED",
1185
- f"Could not decode image: {exc}", 422,
1186
- root_cause=str(exc),
1187
- recommendation="Ensure the file is a valid, non-corrupted image.",
1188
- ) from exc
1189
- return img
1190
-
1191
-
1192
- def _convert_to_png(raw: bytes, ext: str) -> bytes:
1193
- if ext in {"heic", "heif"}:
1194
- try:
1195
- from pillow_heif import register_heif_opener
1196
- register_heif_opener()
1197
- except ImportError:
1198
- raise _err("decode", "HEIF_NOT_SUPPORTED",
1199
- "HEIC/HEIF requires pillow-heif.", 415)
1200
- try:
1201
- img = Image.open(io.BytesIO(raw)).convert("RGB")
1202
- buf = io.BytesIO()
1203
- img.save(buf, format="PNG")
1204
- return buf.getvalue()
1205
- except Exception as exc:
1206
- raise _err("decode", "IMAGE_DECODE_FAILED",
1207
- f"Pillow could not open image: {exc}", 422) from exc
1208
-
1209
-
1210
- # ═════════════════════════════════════════════════════════════════════════════
1211
- # RapidOCR output parser
1212
- # Returns (plain_text, mean_confidence)
1213
- # ══════════════════════════════════════════��══════════════════════════════════
1214
- def _parse_rapidocr_output(result: Any) -> tuple[str, float]:
1215
- if not result:
1216
- return "", 0.0
1217
-
1218
- def _avg_y(item: Any) -> float:
1219
- box = item[0]
1220
- try:
1221
- return sum(pt[1] for pt in box) / 4
1222
- except Exception:
1223
- return 0.0
1224
-
1225
- def _avg_x(item: Any) -> float:
1226
- box = item[0]
1227
- try:
1228
- return sum(pt[0] for pt in box) / 4
1229
- except Exception:
1230
- return 0.0
1231
-
1232
- sorted_items = sorted(result, key=_avg_y)
1233
-
1234
- LINE_GAP = 20
1235
- lines: list[list[Any]] = []
1236
- if sorted_items:
1237
- current: list[Any] = [sorted_items[0]]
1238
- for item in sorted_items[1:]:
1239
- if abs(_avg_y(item) - _avg_y(current[-1])) < LINE_GAP:
1240
- current.append(item)
1241
- else:
1242
- lines.append(current)
1243
- current = [item]
1244
- lines.append(current)
1245
-
1246
- text_lines: list[str] = []
1247
- for line in lines:
1248
- words = sorted(line, key=_avg_x)
1249
- text_lines.append(" ".join(str(item[1]) for item in words if len(item) > 1))
1250
-
1251
- plain_text = "\n".join(text_lines)
1252
- scores = [item[2] for item in result if len(item) > 2 and item[2] is not None]
1253
- mean_conf = float(sum(scores) / len(scores)) if scores else 0.5
1254
- return plain_text, round(mean_conf, 4)
1255
-
1256
-
1257
- def _split_elapse(elapse: Any, total_ms: int) -> tuple[int, int]:
1258
- """
1259
- Extract det_ms / rec_ms from RapidOCR's elapse return value.
1260
-
1261
- rapidocr-onnxruntime ≥ 1.3 returns a dict: {"det": s, "rec": s, "cls": s}.
1262
- Older versions return a scalar total. We handle both.
1263
- """
1264
- if isinstance(elapse, dict):
1265
- det_ms = int(elapse.get("det", 0) * 1000)
1266
- rec_ms = int(elapse.get("rec", 0) * 1000)
1267
- return det_ms, rec_ms
1268
- # Scalar fallback — measured total, no reliable split available
1269
- return 0, total_ms
1270
-
1271
-
1272
- # ═════════════════════════════════════════════════════════════════════════════
1273
- # Misc helpers
1274
- # ═════════════════════════════════════════════════════════════════════════════
1275
- def _sanitize_filename(name: str) -> str:
1276
- name = os.path.basename(name)
1277
- name = re.sub(r"[^\w.\-]", "_", name)
1278
- return name[:200] or "upload"
1279
-
1280
-
1281
- def _markdown_to_plain(markdown: str) -> str:
1282
- text = re.sub(r"!\[.*?\]\(.*?\)", "", markdown)
1283
- text = re.sub(r"\[([^\]]+)\]\([^\)]+\)", r"\1", text)
1284
- text = re.sub(r"#{1,6}\s*", "", text)
1285
- text = re.sub(r"\*{1,2}([^*]+)\*{1,2}", r"\1", text)
1286
- text = re.sub(r"`{1,3}[^`]*`{1,3}", "", text)
1287
- text = re.sub(r"\|", " ", text)
1288
- text = re.sub(r"-{3,}", "", text)
1289
- text = re.sub(r"\n{3,}", "\n\n", text)
1290
- return text.strip()
1291
-
1292
-
1293
- def _pdf_page_count(raw: bytes) -> int:
1294
- try:
1295
- doc = fitz.open(stream=raw, filetype="pdf")
1296
- count = doc.page_count
1297
- doc.close()
1298
- return count
1299
- except Exception:
1300
- return 1
1301
-
1302
-
1303
- def _mem_mb() -> tuple[int, int]:
1304
- try:
1305
- import psutil
1306
- vm = psutil.virtual_memory()
1307
- return (vm.total - vm.available) // (1024 * 1024), vm.total // (1024 * 1024)
1308
- except Exception:
1309
- pass
1310
- try:
1311
- info: dict[str, int] = {}
1312
- with open("/proc/meminfo") as f:
1313
- for line in f:
1314
- parts = line.split()
1315
- if len(parts) >= 2:
1316
- info[parts[0].rstrip(":")] = int(parts[1])
1317
- total_kb = info.get("MemTotal", 0)
1318
- avail_kb = info.get("MemAvailable", 0)
1319
- return (total_kb - avail_kb) // 1024, total_kb // 1024
1320
- except Exception:
1321
- return 0, 0
1322
-
1323
-
1324
- def _assert_memory_safe(raw: bytes, ext: str) -> None:
1325
- """
1326
- Reject requests that would likely exhaust available RAM.
1327
-
1328
- For images: estimate from raw byte count only (no Pillow decode needed —
1329
- avoids the double-decode that existed in v3.0). Raw JPEG at 3 MP ≈ 1–3 MB;
1330
- the decompressed BGR array is w*h*3 bytes. We conservatively multiply by
1331
- IMAGE_MEMORY_FACTOR to cover both the decode buffer and OCR working memory.
1332
- """
1333
- used_mb, total_mb = _mem_mb()
1334
- if total_mb == 0:
1335
- return
1336
- available_mb = total_mb - used_mb
1337
- if ext in PDF_EXTENSIONS:
1338
- page_count = max(1, _pdf_page_count(raw))
1339
- estimated_mb = (page_count * BYTES_PER_OCR_PAGE) // (1024 * 1024)
1340
- else:
1341
- # Estimate from compressed size — no Pillow decode required.
1342
- # Compressed-to-raw expansion ratio for JPEG ≈ 10–20×; we use 20× and
1343
- # multiply by IMAGE_MEMORY_FACTOR for working memory overhead.
1344
- estimated_mb = len(raw) * 20 * IMAGE_MEMORY_FACTOR // (1024 * 1024)
1345
-
1346
- free_after = available_mb - estimated_mb
1347
- logger.info(
1348
- "memory_check avail_mb=%d est_mb=%d free_after_mb=%d",
1349
- available_mb, estimated_mb, free_after,
1350
- )
1351
- if free_after < MEM_SAFETY_FLOOR_MB:
1352
- raise _err(
1353
- "validation", "LOW_MEMORY",
1354
- f"Insufficient memory. Available: {available_mb} MB, "
1355
- f"Estimated needed: {estimated_mb} MB.", 507,
1356
- root_cause=f"Container has {available_mb} MB free; "
1357
- f"pipeline needs ~{estimated_mb} MB.",
1358
- recommendation="Wait for active requests to complete, "
1359
- "or use a smaller file.",
1360
- )