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"
@dataclass
class Segment:
x: int
y: int
w: int
h: int
@property
def right(self) -> int:
return self.x + self.w
@property
def bottom(self) -> int:
return self.y + self.h
@lru_cache(maxsize=1)
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"{html.escape(c)}" for c in chars)
result = f"PALAVRA RECONSTRUÍDA{html.escape(word)}{chips}
"
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("✍️ Reconhecedor de palavras manuscritas
Desenhe uma palavra, visualize a segmentação letra por letra e acompanhe a reconstrução automática.
Dica: escreva em letras de forma, com traço escuro, e deixe um pequeno espaço entre as letras.
")
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("PALAVRA RECONSTRUÍDA—
")
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()