app.py

"""
Sistema de Control de Acceso Vehicular (ANPR + Barrera)
=========================================================
Aplicación standalone. Ejecutar con:

    python app.py

Requiere que 'best.pt' (el modelo entrenado) esté en la misma carpeta que este script.
Al correr, abre automáticamente una interfaz web en el navegador donde se sube un video
y se obtiene: video anotado, registro de accesos, y CSV descargable.
"""

import os
import cv2
import torch
import tempfile
import pandas as pd
from difflib import SequenceMatcher
from ultralytics import YOLO
import easyocr
import gradio as gr

# ---------------------------------------------------------------------------
# Carga de modelo y OCR
# ---------------------------------------------------------------------------
MODEL_PATH = 'best.pt'
if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(
        "No se encontró 'best.pt' en esta carpeta. "
        "Colocá el archivo del modelo entrenado junto a app.py antes de ejecutar."
    )

GPU_DISPONIBLE = torch.cuda.is_available()
print(f'GPU disponible: {GPU_DISPONIBLE}')
if not GPU_DISPONIBLE:
    print('Corriendo en CPU. El procesamiento será más lento, pero funciona igual.')

model = YOLO(MODEL_PATH)  # Ultralytics usa GPU automáticamente si está disponible
reader = easyocr.Reader(['en'], gpu=GPU_DISPONIBLE)
print('Modelo y OCR cargados correctamente.')


# ---------------------------------------------------------------------------
# OCR sobre el recorte de la placa
# ---------------------------------------------------------------------------
def leer_placa(frame, box):
    x1, y1, x2, y2 = map(int, box)
    recorte = frame[max(0, y1):y2, max(0, x1):x2]
    if recorte.size == 0:
        return '', 0.0
    resultado = reader.readtext(recorte)
    if not resultado:
        return '', 0.0
    texto = ''.join([r[1] for r in resultado]).strip().upper()
    confianza = sum(r[2] for r in resultado) / len(resultado)
    return texto, confianza


# ---------------------------------------------------------------------------
# Lógica de control de acceso
# ---------------------------------------------------------------------------
def es_similar(a, b, umbral=0.75):
    if not a or not b:
        return False
    return SequenceMatcher(None, a, b).ratio() >= umbral


def verificar_acceso(texto_leido, placas_autorizadas):
    for autorizada in placas_autorizadas:
        if es_similar(texto_leido, autorizada.strip().upper().replace(' ', '')):
            return True, autorizada
    return False, None


def cajas_se_superponen(box1, box2, iou_min=0.3):
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])
    inter = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area1 + area2 - inter
    iou = inter / union if union > 0 else 0.0
    return iou >= iou_min


def dibujar_barrera(frame, permitido, texto_placa):
    h, w = frame.shape[:2]
    color = (0, 200, 0) if permitido else (0, 0, 220)
    mensaje = (
        f'ACCESO PERMITIDO - BARRERA ABIERTA: {texto_placa}'
        if permitido else
        f'ACCESO DENEGADO - BARRERA CERRADA: {texto_placa}'
    )
    overlay = frame.copy()
    cv2.rectangle(overlay, (0, 0), (w, 50), color, -1)
    frame = cv2.addWeighted(overlay, 0.6, frame, 0.4, 0)
    cv2.putText(frame, mensaje, (10, 33), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    return frame


# ---------------------------------------------------------------------------
# Procesamiento del video (detección + OCR + decisión estabilizada + registro)
# ---------------------------------------------------------------------------
def procesar_video(video_path, placas_autorizadas_str, confianza=0.25,
                    ocr_cada_n_frames=20, frames_gracia_desaparicion=15, progress=gr.Progress()):
    placas_autorizadas = [p for p in placas_autorizadas_str.split(',') if p.strip()]

    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 1

    out_path = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False).name
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    writer = cv2.VideoWriter(out_path, fourcc, fps, (w, h))

    registro = []  # (placa, timestamp, resultado)

    def estado_vacio():
        return {'texto': '', 'permitido': False, 'activo': False, 'confianza': 0.0,
                 'box': None, 'decidido': False, 'frames_sin_deteccion': 0}

    estado_actual = estado_vacio()
    frame_idx = 0
    progress(0, desc=f'Iniciando procesamiento ({total_frames} frames totales)...')

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        results = model.predict(frame, conf=confianza, verbose=False)[0]
        timestamp = round(frame_idx / fps, 1)
        vehiculo_presente_este_frame = False

        for box in results.boxes.xyxy.cpu().numpy():
            x1, y1, x2, y2 = map(int, box)

            es_el_que_seguimos = (
                estado_actual['box'] is None or
                cajas_se_superponen(box, estado_actual['box'])
            )

            if not es_el_que_seguimos:
                # Un vehículo distinto mientras ya se sigue a otro: se ignora hasta que el actual desaparezca.
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 255), 2)
                continue

            estado_actual['box'] = box
            estado_actual['frames_sin_deteccion'] = 0
            vehiculo_presente_este_frame = True

            if not estado_actual['decidido'] and frame_idx % ocr_cada_n_frames == 0:
                texto, conf_lectura = leer_placa(frame, box)
                if texto and conf_lectura > estado_actual['confianza']:
                    permitido, _ = verificar_acceso(texto, placas_autorizadas)
                    es_primer_registro = not estado_actual['activo']
                    estado_actual.update({'texto': texto, 'permitido': permitido, 'activo': True, 'confianza': conf_lectura})

                    if permitido:
                        estado_actual['decidido'] = True  # ya está confirmado: se deja de leer, solo se sigue el tracker

                    if es_primer_registro:
                        registro.append((texto, timestamp, 'PERMITIDO' if permitido else 'DENEGADO'))
                    elif registro:
                        registro[-1] = (texto, registro[-1][1], 'PERMITIDO' if permitido else 'DENEGADO')

            color_caja = (0, 200, 0) if estado_actual['permitido'] else (0, 0, 220)
            cv2.rectangle(frame, (x1, y1), (x2, y2), color_caja, 2)

        if estado_actual['box'] is not None and not vehiculo_presente_este_frame:
            estado_actual['frames_sin_deteccion'] += 1
            if estado_actual['frames_sin_deteccion'] >= frames_gracia_desaparicion:
                # El vehículo ya no está en cuadro: se libera el seguimiento para poder leer el próximo.
                estado_actual = estado_vacio()

        if estado_actual['activo']:
            frame = dibujar_barrera(frame, estado_actual['permitido'], estado_actual['texto'])

        writer.write(frame)
        frame_idx += 1
        if frame_idx % 50 == 0 or frame_idx == total_frames - 1:
            progress(min(frame_idx / total_frames, 1.0), desc=f'Procesando frame {frame_idx}/{total_frames}')

    cap.release()
    writer.release()

    df = pd.DataFrame(registro, columns=['Placa detectada', 'Timestamp (s)', 'Resultado'])
    csv_path = tempfile.NamedTemporaryFile(suffix='.csv', delete=False).name
    df.to_csv(csv_path, index=False)

    n_permitidos = sum(1 for r in registro if r[2] == 'PERMITIDO')
    n_denegados = len(registro) - n_permitidos
    resumen = f'Vehículos detectados: {len(registro)} | Acceso permitido: {n_permitidos} | Acceso denegado: {n_denegados}'

    return out_path, df, csv_path, resumen


# ---------------------------------------------------------------------------
# Interfaz web (Gradio)
# ---------------------------------------------------------------------------
demo = gr.Interface(
    fn=procesar_video,
    inputs=[
        gr.Video(label='Subí un video de vehículos'),
        gr.Textbox(label='Placas autorizadas (separadas por coma)', placeholder='ABC123, XYZ789', value='TS07JS9670'),
        gr.Slider(0.1, 0.9, value=0.25, step=0.05, label='Umbral de confianza de detección')
    ],
    outputs=[
        gr.Video(label='Video con resultado de acceso'),
        gr.Dataframe(label='Registro de accesos'),
        gr.File(label='Descargar registro (CSV)'),
        gr.Textbox(label='Resumen')
    ],
    title='Sistema de Control de Acceso Vehicular (ANPR)',
    description=(
        'Detección de placas (YOLOv8n) + lectura OCR + verificación contra lista de '
        'autorizados, simulando la apertura/cierre de una barrera de estacionamiento.'
    )
)

if __name__ == '__main__':
    demo.launch()