import gradio as gr
from PIL import Image, ImageDraw
import torch
from transformers import YolosImageProcessor, YolosForObjectDetection

# Cargar modelo
processor = YolosImageProcessor.from_pretrained("nickmuchi/yolos-small-finetuned-license-plate-detection")
model = YolosForObjectDetection.from_pretrained("nickmuchi/yolos-small-finetuned-license-plate-detection")

def borrar_numeros(image):
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)

    width, height = image.size
    target_sizes = torch.tensor([[height, width]])
    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.7)[0]

    if not results["boxes"]:
        return image

    draw = ImageDraw.Draw(image)

    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        x1, y1, x2, y2 = map(int, box.tolist())

        # Recortar matrícula
        plate = image.crop((x1, y1, x2, y2))
        pw, ph = plate.size

        # Borrar 4 dígitos (izquierda de la placa)
        borrar_x = int(pw * 0.5)
        draw_crop = ImageDraw.Draw(plate)
        draw_crop.rectangle([0, 0, borrar_x, ph], fill="white")

        # Reinsertar en la imagen original
        image.paste(plate, (x1, y1))

    return image

gr.Interface(
    fn=borrar_numeros,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Limpieza Automática de Matrículas",
    description="Detecta la matrícula y borra automáticamente los 4 números, dejando solo las letras.",
).launch()