app.py

import gradio as gr
import os
import io
import png
import tensorflow as tf
import tensorflow_text as tf_text
import tensorflow_hub as tf_hub
import numpy as np
from PIL import Image
from huggingface_hub import snapshot_download, HfFolder
from sklearn.metrics.pairwise import cosine_similarity
import traceback
import time
import pandas as pd # Para formatear la salida en tabla

# --- Configuración ---
MODEL_REPO_ID = "google/cxr-foundation"
MODEL_DOWNLOAD_DIR = './hf_cxr_foundation_space'
SIMILARITY_DIFFERENCE_THRESHOLD = 0.1
POSITIVE_SIMILARITY_THRESHOLD = 0.1
print(f"Usando umbrales: Comp Δ={SIMILARITY_DIFFERENCE_THRESHOLD}, Simp τ={POSITIVE_SIMILARITY_THRESHOLD}")

# --- Prompts ---
criteria_list_positive = [
    "optimal centering", "optimal inspiration", "optimal penetration",
    "complete field of view", "scapulae retracted", "sharp image", "artifact free"
]
criteria_list_negative = [
    "poorly centered", "poor inspiration", "non-diagnostic exposure",
    "cropped image", "scapulae overlying lungs", "blurred image", "obscuring artifact"
]

# --- Funciones Auxiliares (MISMAS que en la versión anterior de Gradio) ---
# @tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=tf.string)])
# def preprocess_text(text):
#     return bert_preprocessor_global(text) # Asume que bert_preprocessor_global está cargado

def bert_tokenize(text, preprocessor):
    if preprocessor is None: raise ValueError("BERT preprocessor no está cargado.")
    if not isinstance(text, str): text = str(text)
    out = preprocessor(tf.constant([text.lower()]))
    ids = out['input_word_ids'].numpy().astype(np.int32)
    masks = out['input_mask'].numpy().astype(np.float32)
    paddings = 1.0 - masks
    end_token_idx = (ids == 102)
    ids[end_token_idx] = 0
    paddings[end_token_idx] = 1.0
    if ids.ndim == 2: ids = np.expand_dims(ids, axis=1)
    if paddings.ndim == 2: paddings = np.expand_dims(paddings, axis=1)
    expected_shape = (1, 1, 128)
    if ids.shape != expected_shape:
         if ids.shape == (1,128): ids = np.expand_dims(ids, axis=1)
         else: raise ValueError(f"Shape incorrecta para ids: {ids.shape}, esperado {expected_shape}")
    if paddings.shape != expected_shape:
         if paddings.shape == (1,128): paddings = np.expand_dims(paddings, axis=1)
         else: raise ValueError(f"Shape incorrecta para paddings: {paddings.shape}, esperado {expected_shape}")
    return ids, paddings

def png_to_tfexample(image_array: np.ndarray) -> tf.train.Example:
    if image_array.ndim == 3 and image_array.shape[2] == 1:
        image_array = np.squeeze(image_array, axis=2)
    elif image_array.ndim != 2:
        raise ValueError(f'Array debe ser 2-D. Dimensiones: {image_array.ndim}')
    image = image_array.astype(np.float32)
    min_val, max_val = image.min(), image.max()
    if max_val <= min_val:
        if image_array.dtype == np.uint8 or (min_val >= 0 and max_val <= 255):
             pixel_array = image.astype(np.uint8); bitdepth = 8
        else:
             pixel_array = np.zeros_like(image, dtype=np.uint16); bitdepth = 16
    else:
        image -= min_val
        current_max = max_val - min_val
        if image_array.dtype != np.uint8:
            image *= 65535 / current_max
            pixel_array = image.astype(np.uint16); bitdepth = 16
        else:
            image *= 255 / current_max
            pixel_array = image.astype(np.uint8); bitdepth = 8
    output = io.BytesIO()
    png.Writer(width=pixel_array.shape[1], height=pixel_array.shape[0], greyscale=True, bitdepth=bitdepth).write(output, pixel_array.tolist())
    example = tf.train.Example()
    features = example.features.feature
    features['image/encoded'].bytes_list.value.append(output.getvalue())
    features['image/format'].bytes_list.value.append(b'png')
    return example

def generate_image_embedding(img_np, elixrc_infer, qformer_infer):
    if elixrc_infer is None or qformer_infer is None: raise ValueError("Modelos ELIXR-C o QFormer no cargados.")
    try:
        serialized_img_tf_example = png_to_tfexample(img_np).SerializeToString()
        elixrc_output = elixrc_infer(input_example=tf.constant([serialized_img_tf_example]))
        elixrc_embedding = elixrc_output['feature_maps_0'].numpy()
        qformer_input_img = {
            'image_feature': elixrc_embedding.tolist(),
            'ids': np.zeros((1, 1, 128), dtype=np.int32).tolist(),
            'paddings': np.ones((1, 1, 128), dtype=np.float32).tolist(),
        }
        qformer_output_img = qformer_infer(**qformer_input_img)
        image_embedding = qformer_output_img['all_contrastive_img_emb'].numpy()
        if image_embedding.ndim > 2:
            image_embedding = np.mean(image_embedding, axis=tuple(range(1, image_embedding.ndim - 1)))
        if image_embedding.ndim == 1: image_embedding = np.expand_dims(image_embedding, axis=0)
        if image_embedding.ndim != 2: raise ValueError(f"Embedding final no tiene 2 dims: {image_embedding.shape}")
        return image_embedding
    except Exception as e:
        print(f"Error generando embedding imagen: {e}"); traceback.print_exc(); raise

def calculate_similarities_and_classify(image_embedding, bert_preprocessor, qformer_infer):
    if image_embedding is None: raise ValueError("Embedding imagen es None.")
    if bert_preprocessor is None: raise ValueError("Preprocesador BERT es None.")
    if qformer_infer is None: raise ValueError("QFormer es None.")
    detailed_results = {}
    print("\n--- Calculando similitudes ---")
    for i in range(len(criteria_list_positive)):
        positive_text, negative_text = criteria_list_positive[i], criteria_list_negative[i]
        criterion_name = positive_text
        print(f"Procesando: \"{criterion_name}\"")
        similarity_positive, similarity_negative, difference = None, None, None
        classification_comp, classification_simp = "ERROR", "ERROR"
        try:
            tokens_pos, paddings_pos = bert_tokenize(positive_text, bert_preprocessor)
            qformer_input_pos = {'image_feature': np.zeros([1, 8, 8, 1376], dtype=np.float32).tolist(), 'ids': tokens_pos.tolist(), 'paddings': paddings_pos.tolist()}
            text_embedding_pos = qformer_infer(**qformer_input_pos)['contrastive_txt_emb'].numpy()
            if text_embedding_pos.ndim == 1: text_embedding_pos = np.expand_dims(text_embedding_pos, axis=0)

            tokens_neg, paddings_neg = bert_tokenize(negative_text, bert_preprocessor)
            qformer_input_neg = {'image_feature': np.zeros([1, 8, 8, 1376], dtype=np.float32).tolist(), 'ids': tokens_neg.tolist(), 'paddings': paddings_neg.tolist()}
            text_embedding_neg = qformer_infer(**qformer_input_neg)['contrastive_txt_emb'].numpy()
            if text_embedding_neg.ndim == 1: text_embedding_neg = np.expand_dims(text_embedding_neg, axis=0)

            if image_embedding.shape[1] != text_embedding_pos.shape[1]: raise ValueError(f"Dim mismatch: Img ({image_embedding.shape[1]}) vs Pos ({text_embedding_pos.shape[1]})")
            if image_embedding.shape[1] != text_embedding_neg.shape[1]: raise ValueError(f"Dim mismatch: Img ({image_embedding.shape[1]}) vs Neg ({text_embedding_neg.shape[1]})")

            similarity_positive = cosine_similarity(image_embedding, text_embedding_pos)[0][0]
            similarity_negative = cosine_similarity(image_embedding, text_embedding_neg)[0][0]

            difference = similarity_positive - similarity_negative
            classification_comp = "PASS" if difference > SIMILARITY_DIFFERENCE_THRESHOLD else "FAIL"
            classification_simp = "PASS" if similarity_positive > POSITIVE_SIMILARITY_THRESHOLD else "FAIL"
            print(f"  Sim(+)={similarity_positive:.4f}, Sim(-)={similarity_negative:.4f}, Diff={difference:.4f} -> Comp:{classification_comp}, Simp:{classification_simp}")
        except Exception as e:
            print(f"  ERROR criterio '{criterion_name}': {e}"); traceback.print_exc()
        detailed_results[criterion_name] = {
            'positive_prompt': positive_text, 'negative_prompt': negative_text,
            'similarity_positive': float(similarity_positive) if similarity_positive is not None else None,
            'similarity_negative': float(similarity_negative) if similarity_negative is not None else None,
            'difference': float(difference) if difference is not None else None,
            'classification_comparative': classification_comp, 'classification_simplified': classification_simp
        }
    return detailed_results

# --- Carga Global de Modelos ---
print("--- Iniciando carga global de modelos ---")
start_time = time.time()
models_loaded = False
bert_preprocessor_global = None
elixrc_infer_global = None
qformer_infer_global = None
try:
    # Añadir token si es necesario (para repos privados o gated)
    hf_token = os.environ.get("HF_TOKEN") # Leer token desde secretos del Space
    # if hf_token:
    #     print("Usando HF_TOKEN para autenticación.")
    #     HfFolder.save_token(hf_token)

    os.makedirs(MODEL_DOWNLOAD_DIR, exist_ok=True)
    print(f"Descargando/verificando modelos en: {MODEL_DOWNLOAD_DIR}")
    snapshot_download(repo_id=MODEL_REPO_ID, local_dir=MODEL_DOWNLOAD_DIR,
                      allow_patterns=['elixr-c-v2-pooled/*', 'pax-elixr-b-text/*'],
                      local_dir_use_symlinks=False, token=hf_token) # Pasar token aquí
    print("Modelos descargados/verificados.")

    print("Cargando Preprocesador BERT...")
    bert_preprocess_handle = "https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"
    bert_preprocessor_global = tf_hub.KerasLayer(bert_preprocess_handle)
    print("Preprocesador BERT cargado.")

    print("Cargando ELIXR-C...")
    elixrc_model_path = os.path.join(MODEL_DOWNLOAD_DIR, 'elixr-c-v2-pooled')
    elixrc_model = tf.saved_model.load(elixrc_model_path)
    elixrc_infer_global = elixrc_model.signatures['serving_default']
    print("Modelo ELIXR-C cargado.")

    print("Cargando QFormer (ELIXR-B Text)...")
    qformer_model_path = os.path.join(MODEL_DOWNLOAD_DIR, 'pax-elixr-b-text')
    qformer_model = tf.saved_model.load(qformer_model_path)
    qformer_infer_global = qformer_model.signatures['serving_default']
    print("Modelo QFormer cargado.")

    models_loaded = True
    end_time = time.time()
    print(f"--- Modelos cargados globalmente con éxito en {end_time - start_time:.2f} segundos ---")
except Exception as e:
    models_loaded = False
    print(f"--- ERROR CRÍTICO DURANTE LA CARGA GLOBAL DE MODELOS ---"); print(e); traceback.print_exc()

# --- Función Principal de Procesamiento para Gradio ---
def assess_quality_and_update_ui(image_pil):
    """Procesa la imagen y devuelve actualizaciones para la UI."""
    if not models_loaded:
        raise gr.Error("Error: Los modelos no se pudieron cargar. La aplicación no puede procesar imágenes.")
    if image_pil is None:
        # Devuelve valores por defecto/vacíos y controla la visibilidad
        return (
            gr.update(visible=True),  # Muestra bienvenida
            gr.update(visible=False), # Oculta resultados
            None,                     # Borra imagen de salida
            gr.update(value="N/A"),   # Borra etiqueta
            pd.DataFrame(),           # Borra dataframe
            None                      # Borra JSON
        )

    print("\n--- Iniciando evaluación para nueva imagen ---")
    start_process_time = time.time()
    try:
        # 1. Convertir a NumPy
        img_np = np.array(image_pil.convert('L'))
        # 2. Generar Embedding
        image_embedding = generate_image_embedding(img_np, elixrc_infer_global, qformer_infer_global)
        # 3. Clasificar
        detailed_results = calculate_similarities_and_classify(image_embedding, bert_preprocessor_global, qformer_infer_global)
        # 4. Formatear Resultados
        output_data, passed_count, total_count = [], 0, 0
        for criterion, details in detailed_results.items():
            total_count += 1
            sim_pos = details['similarity_positive']
            sim_neg = details['similarity_negative']
            diff = details['difference']
            comp = details['classification_comparative']
            simp = details['classification_simplified']
            output_data.append([ criterion, f"{sim_pos:.4f}" if sim_pos else "N/A",
                f"{sim_neg:.4f}" if sim_neg else "N/A", f"{diff:.4f}" if diff else "N/A", comp, simp ])
            if comp == "PASS": passed_count += 1
        df_results = pd.DataFrame(output_data, columns=[ "Criterion", "Sim (+)", "Sim (-)", "Difference", "Assessment (Comp)", "Assessment (Simp)" ])
        overall_quality = "Error"; pass_rate = 0
        if total_count > 0:
            pass_rate = passed_count / total_count
            if pass_rate >= 0.85: overall_quality = "Excellent"
            elif pass_rate >= 0.70: overall_quality = "Good"
            elif pass_rate >= 0.50: overall_quality = "Fair"
            else: overall_quality = "Poor"
        quality_label = f"{overall_quality} ({passed_count}/{total_count} passed)"
        end_process_time = time.time()
        print(f"--- Evaluación completada en {end_process_time - start_process_time:.2f} seg ---")
        # Devolver resultados y actualizar visibilidad
        return (
            gr.update(visible=False), # Oculta bienvenida
            gr.update(visible=True),  # Muestra resultados
            image_pil,                # Muestra imagen procesada
            gr.update(value=quality_label), # Actualiza etiqueta
            df_results,               # Actualiza dataframe
            detailed_results          # Actualiza JSON
        )
    except Exception as e:
        print(f"Error durante procesamiento Gradio: {e}"); traceback.print_exc()
        raise gr.Error(f"Error procesando imagen: {str(e)}")

# --- Función para Resetear la UI ---
def reset_ui():
    print("Reseteando UI...")
    return (
        gr.update(visible=True),  # Muestra bienvenida
        gr.update(visible=False), # Oculta resultados
        None,                     # Borra imagen de entrada
        None,                     # Borra imagen de salida
        gr.update(value="N/A"),   # Borra etiqueta
        pd.DataFrame(),           # Borra dataframe
        None                      # Borra JSON
    )

# --- Definir Tema Oscuro Personalizado ---
# Inspirado en los colores del HTML original y Tailwind dark grays/blues
dark_theme = gr.themes.Default(
    primary_hue=gr.themes.colors.blue,       # Azul como color primario
    secondary_hue=gr.themes.colors.blue,     # Azul secundario
    neutral_hue=gr.themes.colors.gray,       # Gris neutro
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "Consolas", "monospace"],
).set(
    # Fondos
    body_background_fill="#111827",          # Fondo principal muy oscuro (gray-900)
    background_fill_primary="#1f2937",       # Fondo de componentes (gray-800)
    background_fill_secondary="#374151",     # Fondo secundario (gray-700)
    block_background_fill="#1f2937",         # Fondo de bloques (gray-800)

    # Texto
    body_text_color="#d1d5db",               # Texto principal claro (gray-300)
    text_color_subdued="#9ca3af",            # Texto secundario (gray-400)
    block_label_text_color="#d1d5db",        # Etiquetas de bloque (gray-300)
    block_title_text_color="#ffffff",        # Títulos de bloque (blanco)

    # Bordes
    border_color_accent="#374151",           # Borde (gray-700)
    border_color_primary="#4b5563",          # Borde primario (gray-600)

    # Botones y Elementos Interactivos
    button_primary_background_fill="*primary_600", # Usa color primario (azul)
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="*neutral_700",
    button_secondary_text_color="#ffffff",
    input_background_fill="#374151",         # Fondo de inputs (gray-700)
    input_border_color="#4b5563",            # Borde de inputs (gray-600)
    input_text_color="#ffffff",              # Texto en inputs

    # Sombras y Radios
    shadow_drop="rgba(0,0,0,0.2) 0px 2px 4px",
    block_shadow="rgba(0,0,0,0.2) 0px 2px 5px",
    radius_size="*radius_lg",                # Bordes redondeados
)


# --- Definir la Interfaz Gradio con Bloques y Tema ---
with gr.Blocks(theme=dark_theme, title="CXR Quality Assessment") as demo:
    # --- Cabecera ---
    with gr.Row():
        gr.Markdown(
            """
            # <span style="color: #e5e7eb;">CXR Quality Assessment</span>
            <p style="color: #9ca3af;">Evaluate chest X-ray technical quality using AI (ELIXR family)</p>
            """, # Usar blanco/gris claro para texto cabecera
            elem_id="app-header"
        )

    # --- Contenido Principal (Dos Columnas) ---
    with gr.Row(equal_height=False): # Permitir alturas diferentes

        # --- Columna Izquierda (Carga) ---
        with gr.Column(scale=1, min_width=350):
            gr.Markdown("### 1. Upload Image", elem_id="upload-title")
            input_image = gr.Image(type="pil", label="Upload Chest X-ray", height=300) # Altura fija para imagen entrada
            with gr.Row():
                 analyze_btn = gr.Button("Analyze Image", variant="primary", scale=2)
                 reset_btn = gr.Button("Reset", variant="secondary", scale=1)
            # Añadir ejemplos si tienes imágenes de ejemplo
            # gr.Examples(
            #     examples=[os.path.join("examples", "sample_cxr.png")],
            #     inputs=input_image, label="Example CXR"
            # )
            gr.Markdown(
                "<p style='color:#9ca3af; font-size:0.9em;'>Model loading on startup takes ~1 min. Analysis takes ~15-40 sec.</p>"
            )


        # --- Columna Derecha (Bienvenida / Resultados) ---
        with gr.Column(scale=2):

            # --- Bloque de Bienvenida (Visible Inicialmente) ---
            with gr.Column(visible=True, elem_id="welcome-section") as welcome_block:
                gr.Markdown(
                    """
                    ### Welcome!
                    Upload a chest X-ray image (PNG, JPG, etc.) on the left panel and click "Analyze Image".

                    The system will evaluate its technical quality based on 7 standard criteria using the ELIXR model family.
                    The results will appear here once the analysis is complete.
                    """, elem_id="welcome-text"
                )
                # Podrías añadir un icono o imagen aquí si quieres
                # gr.Image("path/to/welcome_icon.png", interactive=False, show_label=False, show_download_button=False)


            # --- Bloque de Resultados (Oculto Inicialmente) ---
            with gr.Column(visible=False, elem_id="results-section") as results_block:
                gr.Markdown("### 2. Quality Assessment Results", elem_id="results-title")
                with gr.Row(): # Fila para imagen de salida y resumen
                     with gr.Column(scale=1):
                          output_image = gr.Image(type="pil", label="Analyzed Image", interactive=False)
                     with gr.Column(scale=1):
                          gr.Markdown("#### Summary", elem_id="summary-title")
                          output_label = gr.Label(value="N/A", label="Overall Quality Estimate", elem_id="quality-label")
                          # Podríamos añadir más texto de resumen aquí si quisiéramos

                gr.Markdown("#### Detailed Criteria Evaluation", elem_id="detailed-title")
                output_dataframe = gr.DataFrame(
                    headers=["Criterion", "Sim (+)", "Sim (-)", "Difference", "Assessment (Comp)", "Assessment (Simp)"],
                    label=None, # Quitar etiqueta redundante
                    wrap=True,
                    # La altura ahora se maneja mejor automáticamente o con CSS
                    # row_count=(7, "dynamic") # Mostrar 7 filas, permitir scroll si hay más
                    max_rows=10, # Limitar filas visibles con scroll
                    overflow_row_behaviour="show_ends", # Muestra inicio/fin al hacer scroll
                    interactive=False, # No editable
                    elem_id="results-dataframe"
                )
                with gr.Accordion("Raw JSON Output (for debugging)", open=False):
                      output_json = gr.JSON(label=None)

                gr.Markdown(
                    f"""
                    #### Technical Notes
                    *   **Criterion:** Quality aspect evaluated.
                    *   **Sim (+/-):** Cosine similarity with positive/negative prompt.
                    *   **Difference:** Sim (+) - Sim (-).
                    *   **Assessment (Comp):** PASS if Difference > {SIMILARITY_DIFFERENCE_THRESHOLD}. (Main Result)
                    *   **Assessment (Simp):** PASS if Sim (+) > {POSITIVE_SIMILARITY_THRESHOLD}.
                    """, elem_id="notes-text"
                )

    # --- Pie de página ---
    gr.Markdown(
        """
        ----
        <p style='text-align:center; color:#9ca3af; font-size:0.8em;'>
        CXR Quality Assessment Tool | Model: google/cxr-foundation | Interface: Gradio
        </p>
        """, elem_id="app-footer"
        )


    # --- Conexiones de Eventos ---
    analyze_btn.click(
        fn=assess_quality_and_update_ui,
        inputs=[input_image],
        outputs=[
            welcome_block,      # -> actualiza visibilidad bienvenida
            results_block,      # -> actualiza visibilidad resultados
            output_image,       # -> muestra imagen analizada
            output_label,       # -> actualiza etiqueta resumen
            output_dataframe,   # -> actualiza tabla
            output_json         # -> actualiza JSON
        ]
    )

    reset_btn.click(
        fn=reset_ui,
        inputs=None, # No necesita inputs
        outputs=[
            welcome_block,
            results_block,
            input_image,        # -> limpia imagen entrada
            output_image,
            output_label,
            output_dataframe,
            output_json
        ]
    )


# --- Iniciar la Aplicación Gradio ---
if __name__ == "__main__":
     # server_name="0.0.0.0" para accesibilidad en red local
     # server_port=7860 es el puerto estándar de HF Spaces
     # auth=("user", "password") # Si quieres añadir autenticación básica localmente
     demo.launch(server_name="0.0.0.0", server_port=7860) #, share=True) # Quita share=True para despliegue normal