Spaces:
Running
Running
| from __future__ import annotations | |
| import html | |
| import re | |
| from dataclasses import dataclass | |
| from functools import lru_cache | |
| import cv2 | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image, ImageOps | |
| from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
| MODEL_ID = "microsoft/trocr-small-handwritten" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| class Segment: | |
| x: int | |
| y: int | |
| w: int | |
| h: int | |
| def right(self) -> int: | |
| return self.x + self.w | |
| def bottom(self) -> int: | |
| return self.y + self.h | |
| def load_model(): | |
| processor = TrOCRProcessor.from_pretrained(MODEL_ID) | |
| model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID).to(DEVICE) | |
| model.eval() | |
| return processor, model | |
| def editor_image(value) -> np.ndarray | None: | |
| if value is None: | |
| return None | |
| if isinstance(value, dict): | |
| composite = value.get("composite") | |
| value = composite if composite is not None else value.get("background") | |
| if value is None: | |
| return None | |
| if isinstance(value, Image.Image): | |
| value = np.asarray(value) | |
| array = np.asarray(value) | |
| if array.ndim == 2: | |
| return cv2.cvtColor(array.astype(np.uint8), cv2.COLOR_GRAY2RGB) | |
| if array.shape[2] == 4: | |
| alpha = array[:, :, 3:4].astype(np.float32) / 255.0 | |
| rgb = array[:, :, :3].astype(np.float32) | |
| return (rgb * alpha + 255 * (1 - alpha)).astype(np.uint8) | |
| return array[:, :, :3].astype(np.uint8) | |
| def binarize(image: np.ndarray) -> np.ndarray: | |
| gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
| gray = cv2.GaussianBlur(gray, (3, 3), 0) | |
| _, ink = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) | |
| # The canvas is white, but this also tolerates an inverted uploaded image. | |
| if np.count_nonzero(ink) > ink.size * 0.45: | |
| ink = cv2.bitwise_not(ink) | |
| ink = cv2.morphologyEx(ink, cv2.MORPH_OPEN, np.ones((2, 2), np.uint8)) | |
| return ink | |
| def merge_dots(boxes: list[Segment], median_h: float) -> list[Segment]: | |
| """Merge detached dots/accents with the closest overlapping stem.""" | |
| boxes = sorted(boxes, key=lambda b: (b.x, b.y)) | |
| consumed: set[int] = set() | |
| merged: list[Segment] = [] | |
| for i, box in enumerate(boxes): | |
| if i in consumed: | |
| continue | |
| if box.h < median_h * 0.38 and box.w < median_h * 0.45: | |
| candidates = [] | |
| cx = box.x + box.w / 2 | |
| for j, stem in enumerate(boxes): | |
| if i == j or j in consumed or stem.h < median_h * 0.45: | |
| continue | |
| overlap = max(0, min(box.right, stem.right) - max(box.x, stem.x)) | |
| horizontal_distance = abs(cx - (stem.x + stem.w / 2)) | |
| vertical_gap = stem.y - box.bottom | |
| if (overlap > 0 or horizontal_distance < median_h * 0.35) and -2 <= vertical_gap < median_h: | |
| candidates.append((abs(vertical_gap) + horizontal_distance, j, stem)) | |
| if candidates: | |
| _, j, stem = min(candidates, key=lambda item: item[0]) | |
| x1, y1 = min(box.x, stem.x), min(box.y, stem.y) | |
| x2, y2 = max(box.right, stem.right), max(box.bottom, stem.bottom) | |
| merged.append(Segment(x1, y1, x2 - x1, y2 - y1)) | |
| consumed.update({i, j}) | |
| continue | |
| merged.append(box) | |
| consumed.add(i) | |
| return sorted(merged, key=lambda b: b.x) | |
| def split_wide_box(ink: np.ndarray, box: Segment, median_w: float) -> list[Segment]: | |
| """Split touching letters at deep valleys of the vertical ink projection.""" | |
| if box.w < max(median_w * 1.75, box.h * 0.9): | |
| return [box] | |
| roi = ink[box.y : box.bottom, box.x : box.right] | |
| projection = np.count_nonzero(roi, axis=0).astype(np.float32) | |
| if projection.max() == 0: | |
| return [box] | |
| smooth = np.convolve(projection, np.ones(7) / 7, mode="same") | |
| expected = max(2, int(round(box.w / max(median_w, box.h * 0.55)))) | |
| cuts: list[int] = [] | |
| min_spacing = max(8, int(box.h * 0.25)) | |
| for _ in range(expected - 1): | |
| candidates = np.argsort(smooth) | |
| selected = None | |
| for pos in candidates: | |
| if pos < min_spacing or pos > box.w - min_spacing: | |
| continue | |
| if all(abs(int(pos) - cut) >= min_spacing for cut in cuts): | |
| selected = int(pos) | |
| break | |
| if selected is None or smooth[selected] > projection.max() * 0.28: | |
| break | |
| cuts.append(selected) | |
| if not cuts: | |
| return [box] | |
| edges = [0] + sorted(cuts) + [box.w] | |
| return [Segment(box.x + a, box.y, b - a, box.h) for a, b in zip(edges, edges[1:]) if b - a > 5] | |
| def segment_letters(ink: np.ndarray, min_area: int) -> list[Segment]: | |
| n, _, stats, _ = cv2.connectedComponentsWithStats(ink, connectivity=8) | |
| raw = [] | |
| for x, y, w, h, area in stats[1:]: | |
| if area >= min_area and h >= 5 and w >= 2: | |
| raw.append(Segment(int(x), int(y), int(w), int(h))) | |
| if not raw: | |
| return [] | |
| median_h = float(np.median([b.h for b in raw])) | |
| boxes = merge_dots(raw, median_h) | |
| substantial = [b for b in boxes if b.h >= median_h * 0.45] | |
| median_w = float(np.median([b.w for b in substantial or boxes])) | |
| result = [] | |
| for box in boxes: | |
| result.extend(split_wide_box(ink, box, median_w)) | |
| return sorted(result, key=lambda b: b.x) | |
| def prepare_crop(ink: np.ndarray, box: Segment) -> Image.Image: | |
| pad = max(6, int(max(box.w, box.h) * 0.15)) | |
| x1, y1 = max(0, box.x - pad), max(0, box.y - pad) | |
| x2, y2 = min(ink.shape[1], box.right + pad), min(ink.shape[0], box.bottom + pad) | |
| roi = ink[y1:y2, x1:x2] | |
| ys, xs = np.where(roi > 0) | |
| if len(xs): | |
| roi = roi[ys.min() : ys.max() + 1, xs.min() : xs.max() + 1] | |
| side = max(roi.shape) + 40 | |
| canvas = np.full((side, side), 255, dtype=np.uint8) | |
| glyph = 255 - roi | |
| oy, ox = (side - glyph.shape[0]) // 2, (side - glyph.shape[1]) // 2 | |
| canvas[oy : oy + glyph.shape[0], ox : ox + glyph.shape[1]] = glyph | |
| return Image.fromarray(canvas).convert("RGB").resize((384, 384), Image.Resampling.LANCZOS) | |
| def decode_character(crop: Image.Image) -> str: | |
| processor, model = load_model() | |
| pixels = processor(images=crop, return_tensors="pt").pixel_values.to(DEVICE) | |
| with torch.inference_mode(): | |
| ids = model.generate(pixels, max_new_tokens=4, num_beams=4, early_stopping=True) | |
| text = processor.batch_decode(ids, skip_special_tokens=True)[0].strip() | |
| match = re.search(r"[A-Za-zÀ-ÖØ-öø-ÿ0-9]", text) | |
| return match.group(0).upper() if match else "?" | |
| def recognize(value, min_area: int, uppercase: bool): | |
| image = editor_image(value) | |
| if image is None: | |
| raise gr.Error("Desenhe uma palavra no quadro antes de reconhecer.") | |
| ink = binarize(image) | |
| boxes = segment_letters(ink, int(min_area)) | |
| if not boxes: | |
| raise gr.Error("Nenhuma letra foi encontrada. Use traços escuros e mais espessos.") | |
| if len(boxes) > 20: | |
| raise gr.Error("Foram detectados mais de 20 segmentos. Limpe o quadro e tente novamente.") | |
| annotated = image.copy() | |
| gallery = [] | |
| chars = [] | |
| for index, box in enumerate(boxes, 1): | |
| crop = prepare_crop(ink, box) | |
| char = decode_character(crop) | |
| if not uppercase: | |
| char = char.lower() | |
| chars.append(char) | |
| gallery.append((crop, f"Letra {index}: {char}")) | |
| cv2.rectangle(annotated, (box.x, box.y), (box.right, box.bottom), (16, 185, 129), 3) | |
| cv2.putText(annotated, str(index), (box.x, max(22, box.y - 7)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (30, 64, 175), 2) | |
| word = "".join(chars) | |
| chips = " ".join(f"<span class='chip'>{html.escape(c)}</span>" for c in chars) | |
| result = f"<div class='result'><small>PALAVRA RECONSTRUÍDA</small><strong>{html.escape(word)}</strong><div>{chips}</div></div>" | |
| binary_preview = cv2.cvtColor(255 - ink, cv2.COLOR_GRAY2RGB) | |
| return result, word, annotated, binary_preview, gallery, f"{len(boxes)} letra(s) segmentada(s) • modelo: {MODEL_ID} • dispositivo: {DEVICE}" | |
| def clear_all(): | |
| return None, "", "", None, None, [], "Aguardando desenho…" | |
| CSS = """ | |
| .gradio-container {max-width: 1180px !important; margin: auto !important;} | |
| .hero {padding: 22px 26px; border-radius: 20px; color: white; background: linear-gradient(120deg,#172554,#1d4ed8 60%,#0f766e); margin-bottom: 18px;} | |
| .hero h1 {font-size: 2rem; margin: 0 0 6px;} | |
| .hero p {margin: 0; opacity: .9;} | |
| .result {text-align:center; border:1px solid #bfdbfe; background:#eff6ff; padding:20px; border-radius:18px;} | |
| .result small {display:block; color:#475569; font-weight:700; letter-spacing:.08em;} | |
| .result strong {display:block; color:#172554; font-size:3rem; line-height:1.2; margin:6px 0 12px;} | |
| .chip {display:inline-block; min-width:32px; padding:5px 9px; margin:3px; color:white; background:#0f766e; border-radius:9px; font-weight:800;} | |
| .hint {border-left:4px solid #14b8a6; padding:10px 14px; background:#f0fdfa; border-radius:8px;} | |
| """ | |
| with gr.Blocks(css=CSS, title="Reconhecedor de Palavras Manuscritas") as demo: | |
| gr.HTML("<div class='hero'><h1>✍️ Reconhecedor de palavras manuscritas</h1><p>Desenhe uma palavra, visualize a segmentação letra por letra e acompanhe a reconstrução automática.</p></div>") | |
| with gr.Row(): | |
| with gr.Column(scale=6): | |
| canvas = gr.Sketchpad( | |
| label="1. Desenhe uma palavra (separe levemente as letras)", | |
| type="numpy", | |
| image_mode="RGBA", | |
| canvas_size=(900, 320), | |
| fixed_canvas=True, | |
| height=360, | |
| ) | |
| gr.HTML("<div class='hint'>Dica: escreva em letras de forma, com traço escuro, e deixe um pequeno espaço entre as letras.</div>") | |
| with gr.Row(): | |
| recognize_btn = gr.Button("🔎 Segmentar e reconhecer", variant="primary") | |
| clear_btn = gr.Button("🧹 Limpar") | |
| with gr.Column(scale=4): | |
| result_html = gr.HTML("<div class='result'><small>PALAVRA RECONSTRUÍDA</small><strong>—</strong></div>") | |
| word_text = gr.Textbox(label="Resultado em texto", interactive=False) | |
| status = gr.Textbox(label="Status", value="Aguardando desenho…", interactive=False) | |
| with gr.Accordion("Ajustes avançados", open=False): | |
| min_area = gr.Slider(10, 500, value=40, step=5, label="Área mínima do componente") | |
| uppercase = gr.Checkbox(value=True, label="Reconstruir em maiúsculas") | |
| gr.Markdown("## Como o sistema chegou ao resultado") | |
| with gr.Row(): | |
| annotated = gr.Image(label="2. Segmentos detectados", type="numpy") | |
| binary = gr.Image(label="Pré-processamento binário", type="numpy") | |
| gallery = gr.Gallery(label="3. Inferência individual por letra", columns=6, rows=2, height="auto", object_fit="contain") | |
| gr.Markdown( | |
| "**Pipeline:** canvas → binarização → componentes conectados/projeção vertical → " | |
| "recorte e centralização → TrOCR por letra → reconstrução da palavra.\n\n" | |
| "> Este é um demonstrador educacional. Letras cursivas muito conectadas podem exigir um modelo treinado especificamente para o seu conjunto de escrita." | |
| ) | |
| outputs = [result_html, word_text, annotated, binary, gallery, status] | |
| recognize_btn.click(recognize, [canvas, min_area, uppercase], outputs) | |
| clear_btn.click(clear_all, outputs=[canvas, result_html, word_text, annotated, binary, gallery, status]) | |
| if __name__ == "__main__": | |
| demo.queue(default_concurrency_limit=2).launch() | |