""" English Handwritten OCR API v3.0 Powered by Microsoft TrOCR Supports: JPG, PNG, WEBP, BMP, TIFF, PDF (all pages) Includes: /test — browser-based test UI Deploy on Hugging Face Spaces (Free Tier) """ from fastapi import FastAPI, File, UploadFile, HTTPException from fastapi.responses import JSONResponse, HTMLResponse from transformers import TrOCRProcessor, VisionEncoderDecoderModel from PIL import Image, ImageOps import torch import fitz import io import base64 import time import logging # ── Logging ─────────────────────────────────────────────────────────────────── logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ── App ──────────────────────────────────────────────────────────────────────── app = FastAPI( title="Handwritten OCR API", description=( "English handwritten text recognition using Microsoft TrOCR. " "Accepts image files (JPG/PNG/WEBP/BMP/TIFF) and PDF files. " "Visit /test for the browser-based test UI." ), version="3.0.0", ) # ── Model ───────────────────────────────────────────────────────────────────── MODEL_NAME = "microsoft/trocr-large-handwritten" PROCESSOR = None MODEL = None DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MAX_FILE_BYTES = 30 * 1024 * 1024 # 30 MB MAX_PDF_PAGES = 20 PDF_DPI = 200 @app.on_event("startup") async def load_model(): global PROCESSOR, MODEL logger.info(f"Loading TrOCR on {DEVICE} …") PROCESSOR = TrOCRProcessor.from_pretrained(MODEL_NAME) MODEL = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME).to(DEVICE) MODEL.eval() logger.info("Model ready.") # ── Core helpers ────────────────────────────────────────────────────────────── def preprocess_image(image: Image.Image) -> Image.Image: image = ImageOps.exif_transpose(image) return image.convert("RGB") def run_ocr(image: Image.Image) -> str: pv = PROCESSOR(images=image, return_tensors="pt").pixel_values.to(DEVICE) with torch.no_grad(): ids = MODEL.generate(pv, max_new_tokens=128) return PROCESSOR.batch_decode(ids, skip_special_tokens=True)[0] def pdf_to_images(data: bytes) -> list: doc = fitz.open(stream=data, filetype="pdf") total = len(doc) if total == 0: raise HTTPException(400, "PDF has no pages.") if total > MAX_PDF_PAGES: raise HTTPException(422, f"PDF has {total} pages; max allowed is {MAX_PDF_PAGES}. " "Split the PDF and call the API in batches.") matrix = fitz.Matrix(PDF_DPI / 72, PDF_DPI / 72) images = [] for page in doc: pix = page.get_pixmap(matrix=matrix, alpha=False) img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) images.append(img) doc.close() return images def ocr_images(images: list) -> list: results = [] for idx, img in enumerate(images, 1): img = preprocess_image(img) t0 = time.time() text = run_ocr(img) results.append({"page": idx, "text": text, "inference_seconds": round(time.time() - t0, 3)}) return results def detect_and_decode(data: bytes): if data[:4] == b"%PDF": return "pdf", pdf_to_images(data) try: img = Image.open(io.BytesIO(data)) img.verify() img = Image.open(io.BytesIO(data)) return "image", [img] except Exception: pass raise HTTPException(400, "Unrecognised file format. Supported: PDF, JPG, PNG, WEBP, BMP, TIFF.") def build_response(file_type, page_results, filename): if file_type == "image": return { "success": True, "file_type": "image", "text": page_results[0]["text"], "inference_seconds": page_results[0]["inference_seconds"], "filename": filename, } return { "success": True, "file_type": "pdf", "page_count": len(page_results), "pages": page_results, "full_text": "\n\n".join(p["text"] for p in page_results), "total_seconds": sum(p["inference_seconds"] for p in page_results), "filename": filename, } # ── Health ──────────────────────────────────────────────────────────────────── @app.get("/", tags=["Health"]) def root(): return { "service": "Handwritten OCR API", "version": "3.0.0", "model": MODEL_NAME, "device": DEVICE, "status": "running", "supported_formats": ["PDF", "JPG", "PNG", "WEBP", "BMP", "TIFF"], "endpoints": { "GET /test": "Browser test UI", "POST /ocr/file": "Upload PDF or image → text", "POST /ocr/base64": "Send base64 PDF or image → text", "GET /health": "Health check", "GET /docs": "Swagger UI", }, } @app.get("/health", tags=["Health"]) def health(): return {"status": "ok", "model_loaded": MODEL is not None} # ── OCR endpoints ───────────────────────────────────────────────────────────── @app.post("/ocr/file", tags=["OCR"]) async def ocr_from_file(file: UploadFile = File(...)): """Upload a PDF or image file. Returns recognised handwritten text.""" data = await file.read() if len(data) > MAX_FILE_BYTES: raise HTTPException(413, f"File too large. Max {MAX_FILE_BYTES//1024//1024} MB.") file_type, images = detect_and_decode(data) results = ocr_images(images) return JSONResponse(build_response(file_type, results, file.filename or "upload")) @app.post("/ocr/base64", tags=["OCR"]) async def ocr_from_base64(payload: dict): """ JSON body: `{ "file": "", "filename": "optional" }` Legacy key `"image"` also accepted. """ raw = payload.get("file") or payload.get("image") if not raw: raise HTTPException(422, "Missing key 'file' in request body.") try: data = base64.b64decode(raw) except Exception: raise HTTPException(400, "Invalid base64 string.") if len(data) > MAX_FILE_BYTES: raise HTTPException(413, f"File too large. Max {MAX_FILE_BYTES//1024//1024} MB.") file_type, images = detect_and_decode(data) results = ocr_images(images) return JSONResponse(build_response(file_type, results, payload.get("filename", "base64_input"))) # ── Test UI ─────────────────────────────────────────────────────────────────── TEST_UI = """ Handwritten OCR — Test

✍️ Handwritten OCR

Microsoft TrOCR  ·  English handwriting  ·  PDF & Image

Upload File

Click to browse or drag & drop

JPG · PNG · WEBP · BMP · TIFF · PDF

Result

""" @app.get("/test", response_class=HTMLResponse, tags=["UI"]) def test_ui(): """Browser-based test page for the OCR API.""" return TEST_UI