teste

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"<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()