Spaces:
Sleeping
Sleeping
| """ | |
| BilinguaScript HTR API | |
| FastAPI application factory with /health and /transcribe routes. | |
| """ | |
| from fastapi import FastAPI, UploadFile, File, Form, HTTPException | |
| from fastapi.responses import JSONResponse | |
| from predict import run_inference, normalize_lang, EN_MODEL_ID, FR_MODEL_ID | |
| import preprocess, logging | |
| logging.basicConfig(level=logging.INFO) | |
| logger = logging.getLogger(__name__) | |
| app = FastAPI(title="BilinguaScript HTR API", version="1.1.0") | |
| def root(): | |
| """Landing response for the Space page - the real work happens on /api/v1/transcribe.""" | |
| return { | |
| "service": "BiLinguaScript HTR API", | |
| "description": "Bilingual English/French handwriting recognition for Cameroonian documents (fine-tuned TrOCR).", | |
| "models": {"en": EN_MODEL_ID, "fr": FR_MODEL_ID}, | |
| "endpoints": { | |
| "health": "GET /health", | |
| "transcribe": "POST /api/v1/transcribe (multipart: file, language=en|fr|auto, segment=true|false)", | |
| }, | |
| "docs": "/docs", | |
| } | |
| def health(): | |
| return { | |
| "status": "ok", | |
| "models": {"en": EN_MODEL_ID, "fr": FR_MODEL_ID}, | |
| } | |
| async def transcribe( | |
| file: UploadFile = File(...), | |
| language: str = Form(default="auto"), | |
| segment: bool = Form(default=True), | |
| ): | |
| """Transcribe an uploaded image. | |
| - language: 'en'/'english', 'fr'/'french', or 'auto' (defaults to English). | |
| - segment: true -> split the image into lines first (full-page input); | |
| false -> treat the image as a single pre-cropped line (this is | |
| what the Django backend sends, since it segments with its own | |
| pipeline before calling us — re-segmenting a line crop here | |
| would waste time and risk splitting tall ascenders/descenders). | |
| """ | |
| data = await file.read() | |
| if not data: | |
| raise HTTPException(status_code=400, detail="Empty file.") | |
| img = preprocess.decode_image(data) | |
| if img is None: | |
| raise HTTPException(status_code=422, detail="Could not decode image.") | |
| lang = normalize_lang(language) | |
| if segment: | |
| lines = preprocess.segment_lines(img) | |
| else: | |
| lines = [(0, img)] | |
| results = [] | |
| for idx, line_img in lines: | |
| encoded = preprocess.encode_png(line_img) | |
| result = run_inference(encoded, lang) | |
| results.append({"line_index": idx, **result}) | |
| return JSONResponse({"language": lang, "lines": results}) | |
| # ── Kraken blla baseline segmentation (primary engine) ─────────────────────── | |
| # kraken's blla neural segmenter is the open-source counterpart of | |
| # Transkribus' ARU-Net baseline detection: it emits reading-order baselines | |
| # with bounding polygons that already include ascenders and descenders, so no | |
| # vertical-padding heuristics are needed. kraken is installed --no-deps on the | |
| # Space (its torch~=2.1 metadata pin conflicts with our torch==2.2.2, which it | |
| # runs on fine), so treat any import/model/segment failure as non-fatal and | |
| # fall back to the docTR word-box grouping below. | |
| _kraken_model = None | |
| _kraken_disabled = False | |
| def _get_kraken_model(): | |
| """Load and cache kraken's default blla model (shipped as package data).""" | |
| global _kraken_model | |
| if _kraken_model is None: | |
| import os | |
| import kraken | |
| from kraken.lib import vgsl | |
| logger.info("Loading kraken blla segmentation model...") | |
| path = os.path.join(os.path.dirname(kraken.__file__), 'blla.mlmodel') | |
| _kraken_model = vgsl.TorchVGSLModel.load_model(path) | |
| return _kraken_model | |
| def _kraken_lines_to_bboxes(seg, page_w, page_h): | |
| """Convert a kraken Segmentation (or legacy dict) to reading-order bboxes.""" | |
| lines = getattr(seg, 'lines', None) | |
| if lines is None and isinstance(seg, dict): | |
| lines = seg.get('lines', []) | |
| out = [] | |
| for ln in lines or []: | |
| poly = getattr(ln, 'boundary', None) | |
| if poly is None and isinstance(ln, dict): | |
| poly = ln.get('boundary') | |
| if not poly: | |
| continue | |
| # kraken's boundary polygon is conservative where strokes are faint - | |
| # it stops short and the first or last words of a line get cropped away | |
| # before the recogniser ever sees them. The baseline it returns for the | |
| # same line spans the full written extent, so take the union of both | |
| # for the horizontal reach while keeping the polygon's vertical reach | |
| # (which is what correctly captures ascenders and descenders). | |
| base = getattr(ln, 'baseline', None) | |
| if base is None and isinstance(ln, dict): | |
| base = ln.get('baseline') | |
| xs = [p[0] for p in poly] + [p[0] for p in (base or [])] | |
| ys = [p[1] for p in poly] | |
| x0 = max(0, int(min(xs))) | |
| y0 = max(0, int(min(ys))) | |
| x1 = min(page_w, int(max(xs)) + 1) | |
| y1 = min(page_h, int(max(ys)) + 1) | |
| if (x1 - x0) >= 4 and (y1 - y0) >= 4: | |
| out.append({'x': x0, 'y': y0, 'w': x1 - x0, 'h': y1 - y0}) | |
| return out | |
| def _merge_overlapping_lines(lines, min_overlap=0.5): | |
| """Merge line boxes that are really fragments of one visual line. | |
| kraken emits a separate baseline wherever a handwritten line has a wide | |
| word gap, so a single sentence can arrive as four or five boxes. They sit | |
| side by side and overlap vertically, whereas genuinely different lines do | |
| not - so chain by vertical overlap with the nearest neighbour. Measured on | |
| a ruled-notebook page: 27 kraken boxes -> 11 lines, matching the 11 lines | |
| actually written. | |
| """ | |
| out = [] | |
| for b in sorted(lines, key=lambda z: z['x']): | |
| y0, y1 = b['y'], b['y'] + b['h'] | |
| bh = max(1, b['h']) | |
| best, best_score = None, 0.0 | |
| for L in out: | |
| ov = min(L['y1'], y1) - max(L['y0'], y0) | |
| if ov <= 0: | |
| continue | |
| score = ov / min(bh, max(1, L['y1'] - L['y0'])) | |
| if score > best_score: | |
| best, best_score = L, score | |
| if best is not None and best_score >= min_overlap: | |
| best['y0'] = min(best['y0'], y0); best['y1'] = max(best['y1'], y1) | |
| best['x0'] = min(best['x0'], b['x']); best['x1'] = max(best['x1'], b['x'] + b['w']) | |
| else: | |
| out.append({'y0': y0, 'y1': y1, 'x0': b['x'], 'x1': b['x'] + b['w']}) | |
| return [ | |
| {'x': L['x0'], 'y': L['y0'], 'w': L['x1'] - L['x0'], 'h': L['y1'] - L['y0']} | |
| for L in sorted(out, key=lambda L: (L['y0'] + L['y1']) / 2) | |
| ] | |
| def _kraken_segment(bgr_img, page_w, page_h): | |
| """Run kraken blla on a BGR ndarray; returns bbox lines or None on failure.""" | |
| global _kraken_disabled | |
| if _kraken_disabled: | |
| return None | |
| try: | |
| import cv2 | |
| from PIL import Image | |
| from kraken import blla | |
| pil = Image.fromarray(cv2.cvtColor(bgr_img, cv2.COLOR_BGR2RGB)) | |
| seg = blla.segment(pil, model=_get_kraken_model()) | |
| lines = _kraken_lines_to_bboxes(seg, page_w, page_h) | |
| if lines: | |
| return _merge_overlapping_lines(lines) | |
| logger.warning("kraken found no lines; falling back to docTR for this page") | |
| return None | |
| except ImportError as e: | |
| _kraken_disabled = True | |
| logger.warning("kraken unavailable (%s); docTR fallback active for this process", e) | |
| return None | |
| except Exception as e: | |
| logger.warning("kraken segmentation failed (%s); docTR fallback for this page", e) | |
| return None | |
| # ── Layout-aware line segmentation (docTR text detection, fallback) ────────── | |
| # The naive projection segmenter fails on phone photos, ruled paper, tables | |
| # and diagrams. docTR's pretrained DB-ResNet detector finds word boxes | |
| # robustly; the grouping below merges them into reading-order lines using the | |
| # same tolerance rule validated in the dissertation corpus pipeline: | |
| # tau = min(alpha * median(word_height), beta * page_height) | |
| _det_model = None | |
| def _get_detector(): | |
| global _det_model | |
| if _det_model is None: | |
| from doctr.models import detection_predictor | |
| logger.info("Loading docTR detection model (db_resnet50)...") | |
| _det_model = detection_predictor('db_resnet50', pretrained=True) | |
| return _det_model | |
| def _group_words_into_lines(boxes, page_w, page_h, min_overlap=0.35): | |
| """boxes: list of (x0, y0, x1, y1) absolute pixels -> list of line bboxes. | |
| Words are chained into lines by VERTICAL OVERLAP with the nearest word | |
| already placed, not by distance from a line's mean centre-y. The old | |
| centre-distance rule fragmented real lines badly: handwriting drifts | |
| downward across a page and ascenders/descenders shift a word's centre, | |
| so words on one visual line easily exceeded a fixed tolerance and each | |
| became its own "line". Comparing a word only against its horizontal | |
| neighbour is the same locality that makes baseline-following systems | |
| (Transkribus' ARU-Net, kraken's blla) robust to slanted, uneven lines. | |
| """ | |
| import numpy as np | |
| if not boxes: | |
| return [] | |
| heights = np.array([b[3] - b[1] for b in boxes], dtype=float) | |
| med_h = float(np.median(heights[heights <= np.percentile(heights, 90)])) if len(heights) > 2 else float(np.median(heights)) | |
| # Left-to-right, so each word chains onto the line its neighbour extended. | |
| lines = [] | |
| for b in sorted(boxes, key=lambda z: z[0]): | |
| bh = max(1.0, b[3] - b[1]) | |
| best, best_score, best_dist = None, 0.0, None | |
| for line in lines: | |
| near = line['last'] # rightmost word placed on this line | |
| ov = min(near[3], b[3]) - max(near[1], b[1]) | |
| if ov <= 0: | |
| continue | |
| score = ov / min(bh, max(1.0, near[3] - near[1])) | |
| dist = max(0.0, b[0] - near[2]) | |
| if score > best_score or (abs(score - best_score) < 0.1 | |
| and best_dist is not None and dist < best_dist): | |
| best, best_score, best_dist = line, max(score, best_score), dist | |
| if best is not None and best_score >= min_overlap: | |
| best['boxes'].append(b) | |
| if b[2] > best['last'][2]: | |
| best['last'] = b | |
| else: | |
| lines.append({'boxes': [b], 'last': b}) | |
| # Reading order: top to bottom by the line's vertical centre. | |
| def line_cy(line): | |
| return sum((x[1] + x[3]) / 2 for x in line['boxes']) / len(line['boxes']) | |
| # Expand each line box the way baseline-based systems (Transkribus' | |
| # ARU-Net, kraken's blla) build their bounding polygons: detectors fire | |
| # on the x-height band, so without generous vertical margins the crops | |
| # decapitate ascenders and cut descenders - exactly the strokes cursive | |
| # readers need. ~30% of the line height above and ~35% below recovers | |
| # them; horizontal padding stays small. | |
| out = [] | |
| pad_x = max(2, int(med_h * 0.15)) | |
| for line in sorted(lines, key=line_cy): | |
| xs0 = min(b[0] for b in line['boxes']); ys0 = min(b[1] for b in line['boxes']) | |
| xs1 = max(b[2] for b in line['boxes']); ys1 = max(b[3] for b in line['boxes']) | |
| line_h = ys1 - ys0 | |
| pad_top = int(line_h * 0.30) | |
| pad_bot = int(line_h * 0.35) | |
| y = max(0, int(ys0 - pad_top)) | |
| out.append({ | |
| 'x': max(0, int(xs0 - pad_x)), 'y': y, | |
| 'w': min(page_w, int(xs1 + pad_x)) - max(0, int(xs0 - pad_x)), | |
| 'h': min(page_h, int(ys1 + pad_bot)) - y, | |
| }) | |
| return out | |
| async def segment_endpoint(file: UploadFile = File(...)): | |
| """Detect handwriting lines on a full page; returns reading-order bboxes. | |
| Primary engine: kraken blla baseline segmentation. Fallback: docTR word | |
| boxes grouped into lines. Response: | |
| {"width": W, "height": H, "engine": "kraken"|"doctr", "lines": [{"x","y","w","h"}, ...]} | |
| """ | |
| data = await file.read() | |
| if not data: | |
| raise HTTPException(status_code=400, detail="Empty file.") | |
| img = preprocess.decode_image(data) | |
| if img is None: | |
| raise HTTPException(status_code=422, detail="Could not decode image.") | |
| h, w = img.shape[:2] | |
| kraken_lines = _kraken_segment(img, w, h) | |
| if kraken_lines is not None: | |
| return JSONResponse({"width": w, "height": h, "engine": "kraken", | |
| "lines": kraken_lines}) | |
| import cv2 | |
| rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
| det = _get_detector() | |
| result = det([rgb]) | |
| # doctr returns normalized (x0, y0, x1, y1[, score]) per word | |
| raw = result[0].get('words', result[0]) if isinstance(result[0], dict) else result[0] | |
| abs_boxes = [] | |
| for b in raw: | |
| x0, y0, x1, y1 = float(b[0]) * w, float(b[1]) * h, float(b[2]) * w, float(b[3]) * h | |
| if (x1 - x0) >= 4 and (y1 - y0) >= 4: | |
| abs_boxes.append((x0, y0, x1, y1)) | |
| lines = _group_words_into_lines(abs_boxes, w, h) | |
| return JSONResponse({"width": w, "height": h, "engine": "doctr", "lines": lines}) | |