app.py

"""
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": "<base64>", "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 = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>Handwritten OCR — Test</title>
<style>
  *, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }

  :root {
    --bg:       #0f1117;
    --surface:  #1a1d27;
    --border:   #2e3348;
    --accent:   #6c7fff;
    --accent2:  #a78bfa;
    --text:     #e2e4f0;
    --muted:    #6b7099;
    --success:  #34d399;
    --error:    #f87171;
    --radius:   12px;
  }

  body {
    background: var(--bg);
    color: var(--text);
    font-family: 'Segoe UI', system-ui, sans-serif;
    min-height: 100vh;
    display: flex;
    flex-direction: column;
    align-items: center;
    padding: 40px 16px 60px;
  }

  header {
    text-align: center;
    margin-bottom: 36px;
  }
  header h1 {
    font-size: 2rem;
    font-weight: 700;
    background: linear-gradient(135deg, var(--accent), var(--accent2));
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
  }
  header p {
    color: var(--muted);
    margin-top: 6px;
    font-size: 0.95rem;
  }

  .card {
    background: var(--surface);
    border: 1px solid var(--border);
    border-radius: var(--radius);
    padding: 28px;
    width: 100%;
    max-width: 760px;
    margin-bottom: 20px;
  }
  .card h2 {
    font-size: 1rem;
    font-weight: 600;
    color: var(--muted);
    text-transform: uppercase;
    letter-spacing: .08em;
    margin-bottom: 18px;
  }

  /* Drop zone */
  .dropzone {
    border: 2px dashed var(--border);
    border-radius: var(--radius);
    padding: 40px 20px;
    text-align: center;
    cursor: pointer;
    transition: border-color .2s, background .2s;
    position: relative;
  }
  .dropzone.over, .dropzone:hover {
    border-color: var(--accent);
    background: rgba(108,127,255,.05);
  }
  .dropzone input[type=file] {
    position: absolute; inset: 0; opacity: 0; cursor: pointer; width: 100%; height: 100%;
  }
  .dropzone svg { margin-bottom: 10px; opacity: .5; }
  .dropzone p   { color: var(--muted); font-size: .9rem; }
  .dropzone p strong { color: var(--text); }

  /* Preview */
  #preview-wrap {
    display: none;
    margin-top: 18px;
    border-radius: var(--radius);
    overflow: hidden;
    border: 1px solid var(--border);
    text-align: center;
    background: #111;
    max-height: 340px;
  }
  #preview-wrap img {
    max-height: 340px;
    max-width: 100%;
    object-fit: contain;
  }
  #preview-wrap .pdf-badge {
    padding: 60px 0;
    font-size: 3rem;
    letter-spacing: .05em;
    color: var(--muted);
  }

  /* File info bar */
  #file-info {
    display: none;
    margin-top: 12px;
    padding: 10px 14px;
    background: rgba(108,127,255,.08);
    border: 1px solid rgba(108,127,255,.2);
    border-radius: 8px;
    font-size: .85rem;
    color: var(--muted);
    display: flex;
    align-items: center;
    gap: 10px;
  }
  #file-info span { color: var(--text); font-weight: 500; }

  /* Button */
  button#run-btn {
    display: block;
    width: 100%;
    margin-top: 20px;
    padding: 13px;
    border: none;
    border-radius: var(--radius);
    background: linear-gradient(135deg, var(--accent), var(--accent2));
    color: #fff;
    font-size: 1rem;
    font-weight: 600;
    cursor: pointer;
    transition: opacity .15s, transform .1s;
  }
  button#run-btn:hover:not(:disabled) { opacity: .88; }
  button#run-btn:active:not(:disabled) { transform: scale(.98); }
  button#run-btn:disabled { opacity: .4; cursor: not-allowed; }

  /* Spinner */
  .spinner {
    display: none;
    width: 22px; height: 22px;
    border: 3px solid rgba(255,255,255,.2);
    border-top-color: #fff;
    border-radius: 50%;
    animation: spin .7s linear infinite;
    margin: 0 auto;
  }
  @keyframes spin { to { transform: rotate(360deg); } }

  /* Results */
  #result-card { display: none; }

  .meta-row {
    display: flex;
    flex-wrap: wrap;
    gap: 10px;
    margin-bottom: 16px;
  }
  .badge {
    padding: 4px 12px;
    border-radius: 20px;
    font-size: .8rem;
    font-weight: 600;
  }
  .badge.type  { background: rgba(108,127,255,.15); color: var(--accent); }
  .badge.pages { background: rgba(167,139,250,.15); color: var(--accent2); }
  .badge.time  { background: rgba(52,211,153,.12);  color: var(--success); }

  /* Page tabs */
  .tabs {
    display: flex;
    flex-wrap: wrap;
    gap: 6px;
    margin-bottom: 14px;
  }
  .tab-btn {
    padding: 5px 14px;
    border-radius: 20px;
    border: 1px solid var(--border);
    background: transparent;
    color: var(--muted);
    font-size: .83rem;
    cursor: pointer;
    transition: all .15s;
  }
  .tab-btn.active, .tab-btn:hover {
    background: var(--accent);
    color: #fff;
    border-color: var(--accent);
  }

  .result-box {
    background: #0b0d14;
    border: 1px solid var(--border);
    border-radius: var(--radius);
    padding: 18px;
    font-family: 'Cascadia Code', 'Fira Code', monospace;
    font-size: .92rem;
    line-height: 1.65;
    white-space: pre-wrap;
    word-break: break-word;
    min-height: 80px;
    color: var(--text);
  }
  .result-box.error { color: var(--error); border-color: rgba(248,113,113,.3); }

  .copy-btn {
    margin-top: 12px;
    padding: 7px 18px;
    border-radius: 8px;
    border: 1px solid var(--border);
    background: transparent;
    color: var(--muted);
    font-size: .85rem;
    cursor: pointer;
    transition: all .15s;
  }
  .copy-btn:hover { border-color: var(--accent); color: var(--accent); }
  .copy-btn.copied { color: var(--success); border-color: var(--success); }

  footer {
    margin-top: 40px;
    color: var(--muted);
    font-size: .8rem;
    text-align: center;
  }
  footer a { color: var(--accent); text-decoration: none; }
</style>
</head>
<body>

<header>
  <h1>✍️ Handwritten OCR</h1>
  <p>Microsoft TrOCR &nbsp;·&nbsp; English handwriting &nbsp;·&nbsp; PDF &amp; Image</p>
</header>

<div class="card">
  <h2>Upload File</h2>

  <div class="dropzone" id="dropzone">
    <input type="file" id="file-input" accept="image/*,.pdf">
    <svg width="40" height="40" fill="none" stroke="currentColor" stroke-width="1.5"
         viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg">
      <path stroke-linecap="round" stroke-linejoin="round"
            d="M3 16.5v2.25A2.25 2.25 0 005.25 21h13.5A2.25 2.25 0 0021 18.75V16.5m-13.5-9L12 3m0 0l4.5 4.5M12 3v13.5"/>
    </svg>
    <p><strong>Click to browse</strong> or drag &amp; drop</p>
    <p style="margin-top:4px">JPG · PNG · WEBP · BMP · TIFF · PDF</p>
  </div>

  <div id="file-info" style="display:none">
    <svg width="16" height="16" fill="none" stroke="currentColor" stroke-width="2"
         viewBox="0 0 24 24"><path stroke-linecap="round" stroke-linejoin="round"
         d="M19.5 14.25v-2.625a3.375 3.375 0 00-3.375-3.375h-1.5A1.125 1.125 0 0113.5 7.125v-1.5a3.375 3.375 0 00-3.375-3.375H8.25m0 12.75h7.5m-7.5 3H12M10.5 2.25H5.625c-.621 0-1.125.504-1.125 1.125v17.25c0 .621.504 1.125 1.125 1.125h12.75c.621 0 1.125-.504 1.125-1.125V11.25a9 9 0 00-9-9z"/></svg>
    <span id="file-name">—</span>
    <span id="file-size" style="margin-left:auto"></span>
  </div>

  <div id="preview-wrap"></div>

  <button id="run-btn" disabled>Run OCR</button>
  <div class="spinner" id="spinner"></div>
</div>

<div class="card" id="result-card">
  <h2>Result</h2>
  <div class="meta-row" id="meta-row"></div>
  <div class="tabs" id="tabs"></div>
  <div class="result-box" id="result-box"></div>
  <button class="copy-btn" id="copy-btn">Copy text</button>
</div>

<footer>
  <a href="/docs" target="_blank">Swagger UI</a> &nbsp;·&nbsp;
  <a href="/health" target="_blank">Health</a> &nbsp;·&nbsp;
  Model: microsoft/trocr-large-handwritten
</footer>

<script>
const fileInput  = document.getElementById('file-input');
const dropzone   = document.getElementById('dropzone');
const runBtn     = document.getElementById('run-btn');
const spinner    = document.getElementById('spinner');
const resultCard = document.getElementById('result-card');
const resultBox  = document.getElementById('result-box');
const metaRow    = document.getElementById('meta-row');
const tabs       = document.getElementById('tabs');
const copyBtn    = document.getElementById('copy-btn');
const previewWrap= document.getElementById('preview-wrap');
const fileInfo   = document.getElementById('file-info');
const fileName   = document.getElementById('file-name');
const fileSize   = document.getElementById('file-size');

let selectedFile  = null;
let pageResults   = [];   // [{page, text, inference_seconds}]
let activePage    = 0;    // index into pageResults

/* ── File selection ─────────────────────────────────────────────────────── */

function humanSize(bytes) {
  if (bytes < 1024) return bytes + ' B';
  if (bytes < 1048576) return (bytes/1024).toFixed(1) + ' KB';
  return (bytes/1048576).toFixed(1) + ' MB';
}

function showPreview(file) {
  previewWrap.style.display = 'block';
  previewWrap.innerHTML = '';
  if (file.type === 'application/pdf' || file.name.toLowerCase().endsWith('.pdf')) {
    previewWrap.innerHTML = '<div class="pdf-badge">PDF 📄</div>';
  } else {
    const img = document.createElement('img');
    img.src = URL.createObjectURL(file);
    previewWrap.appendChild(img);
  }
}

function handleFile(file) {
  if (!file) return;
  selectedFile = file;
  fileName.textContent = file.name;
  fileSize.textContent = humanSize(file.size);
  fileInfo.style.display = 'flex';
  showPreview(file);
  runBtn.disabled = false;
  resultCard.style.display = 'none';
}

fileInput.addEventListener('change', () => handleFile(fileInput.files[0]));

dropzone.addEventListener('dragover', e => { e.preventDefault(); dropzone.classList.add('over'); });
dropzone.addEventListener('dragleave', () => dropzone.classList.remove('over'));
dropzone.addEventListener('drop', e => {
  e.preventDefault();
  dropzone.classList.remove('over');
  handleFile(e.dataTransfer.files[0]);
});

/* ── OCR call ───────────────────────────────────────────────────────────── */

runBtn.addEventListener('click', async () => {
  if (!selectedFile) return;

  runBtn.disabled = true;
  runBtn.textContent = '';
  spinner.style.display = 'block';
  resultCard.style.display = 'none';

  const fd = new FormData();
  fd.append('file', selectedFile);

  try {
    const res  = await fetch('/ocr/file', { method: 'POST', body: fd });
    const data = await res.json();

    spinner.style.display = 'none';
    runBtn.textContent = 'Run OCR';
    runBtn.disabled = false;

    if (!res.ok || !data.success) {
      showError(data.detail || 'OCR failed. Please try another file.');
      return;
    }

    renderResults(data);

  } catch (err) {
    spinner.style.display = 'none';
    runBtn.textContent = 'Run OCR';
    runBtn.disabled = false;
    showError('Network error: ' + err.message);
  }
});

/* ── Render ─────────────────────────────────────────────────────────────── */

function badge(cls, text) {
  const b = document.createElement('span');
  b.className = 'badge ' + cls;
  b.textContent = text;
  return b;
}

function renderResults(data) {
  resultCard.style.display = 'block';
  metaRow.innerHTML = '';
  tabs.innerHTML = '';
  resultBox.className = 'result-box';

  metaRow.appendChild(badge('type', data.file_type.toUpperCase()));

  if (data.file_type === 'pdf') {
    metaRow.appendChild(badge('pages', data.page_count + ' page' + (data.page_count > 1 ? 's' : '')));
    metaRow.appendChild(badge('time',  data.total_seconds.toFixed(2) + 's total'));

    pageResults = [{ page: 0, text: data.full_text, label: 'All pages' }, ...data.pages];
    activePage  = 0;

    pageResults.forEach((p, i) => {
      const btn = document.createElement('button');
      btn.className = 'tab-btn' + (i === 0 ? ' active' : '');
      btn.textContent = i === 0 ? 'All pages' : 'Page ' + p.page;
      btn.dataset.idx = i;
      btn.addEventListener('click', () => switchTab(i));
      tabs.appendChild(btn);
    });

    resultBox.textContent = data.full_text || '(no text recognised)';
  } else {
    metaRow.appendChild(badge('time', data.inference_seconds.toFixed(2) + 's'));
    pageResults = [{ text: data.text }];
    resultBox.textContent = data.text || '(no text recognised)';
  }

  resultCard.scrollIntoView({ behavior: 'smooth', block: 'nearest' });
}

function switchTab(idx) {
  activePage = idx;
  document.querySelectorAll('.tab-btn').forEach((b, i) => {
    b.classList.toggle('active', i === idx);
  });
  resultBox.textContent = pageResults[idx].text || '(no text recognised)';
  copyBtn.classList.remove('copied');
  copyBtn.textContent = 'Copy text';
}

function showError(msg) {
  resultCard.style.display = 'block';
  metaRow.innerHTML = '';
  tabs.innerHTML = '';
  metaRow.appendChild(badge('type', 'ERROR'));
  resultBox.className = 'result-box error';
  resultBox.textContent = msg;
}

/* ── Copy ───────────────────────────────────────────────────────────────── */

copyBtn.addEventListener('click', () => {
  const text = pageResults[activePage]?.text || resultBox.textContent;
  navigator.clipboard.writeText(text).then(() => {
    copyBtn.textContent = 'Copied ✓';
    copyBtn.classList.add('copied');
    setTimeout(() => {
      copyBtn.textContent = 'Copy text';
      copyBtn.classList.remove('copied');
    }, 2000);
  });
});
</script>
</body>
</html>
"""


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