bbox_adder.py

import gradio as gr
from PIL import Image
from typing import List, Dict, Any, Optional

# ----------------------------------------------------------------------
# FUNCIONES DE ADICIÓN DE BBOX MANUAL
# ----------------------------------------------------------------------

def add_new_bbox_mode(
    tokens_data: List[Dict[str, Any]], 
    image_orig: Optional[Image.Image],
    select_data: gr.SelectData
):
    """
    Captura las coordenadas de un Bounding Box dibujado en la interfaz 
    mediante el evento .select(type="select") en gr.Image.

    Args:
        tokens_data: Estado actual de los tokens.
        image_orig: Imagen original (no usada directamente, pero necesaria para la firma).
        select_data: Objeto gr.SelectData devuelto por el evento .select(type="select").

    Returns:
        Tupla de actualizaciones de Gradio:
        (tb_token_editor (Ocultar), dd_tag_selector (Mostrar), 
         tb_new_token_text (Mostrar), btn_add_new_token (Mostrar), 
         STATE_NEW_BBOX (Coordenadas [x1, y1, x2, y2]))
    """
    print("-> Evento de selección (type='select') detectado.")
    
    # Valores de retorno por defecto en caso de fallo (UI ocultos, BBox limpio)
    FAIL_RETURN = (
        gr.update(visible=False), 
        gr.update(visible=False), 
        gr.update(visible=False, value=""), 
        gr.update(visible=False), 
        None
    )
    
    # 1. Verificar y extraer coordenadas del BBox arrastrado
    new_bbox_data = None
    
    # Gradio 4+ con type="select" devuelve un diccionario de coords
    if select_data and hasattr(select_data, 'coords') and select_data.coords:
        coords = select_data.coords
        try:
            x1 = int(coords['x_min'])
            y1 = int(coords['y_min'])
            x2 = int(coords['x_max'])
            y2 = int(coords['y_max'])
            new_bbox_data = [x1, y1, x2, y2]
        except (ValueError, TypeError, KeyError) as e:
            print(f"Error al parsear coordenadas de select_data.coords: {e}")
            return FAIL_RETURN
            
    # Alternativa (si la versión de Gradio devuelve atributos directamente)
    elif select_data and hasattr(select_data, 'x_min') and select_data.x_min is not None:
        try:
            x1 = int(select_data.x_min)
            y1 = int(select_data.y_min)
            x2 = int(select_data.x_max)
            y2 = int(select_data.y_max)
            new_bbox_data = [x1, y1, x2, y2]
        except (ValueError, TypeError) as e:
            print(f"Error al parsear coordenadas de select_data: {e}")
            return FAIL_RETURN
    
    if new_bbox_data is None:
        print("Error: Los datos de selección no contienen las coordenadas (bbox).")
        return FAIL_RETURN

    print(f"-> BBox extraído: {new_bbox_data}")
    
    # 2. Retornar actualizaciones para mostrar la UI de entrada
    # Nota: tb_token_editor debe ocultarse porque solo se usa para editar tokens existentes.
    return (
        gr.update(visible=False), # Ocultar tb_token_editor
        gr.update(visible=True),  # Mostrar dd_tag_selector
        gr.update(visible=True, value=""), # Mostrar tb_new_token_text (limpio)
        gr.update(visible=True),  # Mostrar btn_add_new_token
        new_bbox_data             # Guardar STATE_NEW_BBOX
    )


def append_new_token(
    tokens_data: List[Dict[str, Any]],
    image_orig: Optional[Image.Image],
    bbox: List[int],
    text: str,
    tag: str
):
    """
    Agrega el nuevo token con el BBox dibujado al estado de tokens_data.
    
    Args:
        tokens_data: Lista de tokens existente.
        image_orig: Imagen original (no usada directamente).
        bbox: Coordenadas del BBox [x1, y1, x2, y2].
        text: Texto del nuevo token.
        tag: Etiqueta NER del nuevo token.

    Returns:
        Tupla de actualizaciones de Gradio:
        (tokens_data actualizada, df_label_input (para forzar refresh), 
         tb_token_editor, tb_new_token_text, btn_add_new_token, STATE_NEW_BBOX)
    """
    if not bbox or not text:
        print("Error: BBox o Texto están vacíos. No se puede añadir el token.")
        # Retornar el estado actual sin cambios, pero limpia los campos de entrada
        return (
            tokens_data, 
            gr.update(), # No forzamos actualización de DF
            gr.update(visible=False), 
            gr.update(visible=False, value=""), 
            gr.update(visible=False), 
            None
        )

    # Crear el nuevo token
    new_token = {
        'token': text.strip(),
        'ner_tag': tag,
        'bbox': bbox,
        'id': len(tokens_data), # Asignar un ID simple
        'page': 0 # Asumimos página 0 si no hay lógica de multipágina
    }
    
    # Añadir el token al final de la lista
    tokens_data.append(new_token)
    print(f"-> Nuevo token añadido: {new_token}")

    # Preparar la nueva tabla para Gradio (actualizar df_label_input)
    # Nota: El dataframe de Gradio espera una lista de listas [[token, tag], [token, tag], ...]
    df_rows = [[t['token'], t['ner_tag']] for t in tokens_data]
    
    # Retornar los estados actualizados y limpiar la UI de dibujo
    return (
        tokens_data, 
        gr.update(value=df_rows), # Actualiza el DataFrame
        gr.update(visible=False), # Ocultar tb_token_editor
        gr.update(visible=False, value=""), # Limpiar y Ocultar tb_new_token_text
        gr.update(visible=False), # Ocultar btn_add_new_token
        None                      # Limpiar STATE_NEW_BBOX
    )