import tensorflow as tf
import keras
import gradio as gr
import numpy as np
from PIL import Image
import os
from huggingface_hub import hf_hub_download
from tensorflow.keras.applications.inception_v3 import preprocess_input
from tensorflow.keras.preprocessing.image import img_to_array
import base64

# Descargar el modelo desde Hugging Face
model_path = hf_hub_download(repo_id="Daniel00611/InceptionV3_72", filename="InceptionV3_72.keras")
# Cargar tu modelo de TensorFlow en formato .keras
model = keras.models.load_model(model_path)

# Función auxiliar para decodificar base64 a una imagen PIL
def base64_to_image(base64_str):
    image_data = base64.b64decode(base64_str)
    image = Image.open(BytesIO(image_data))
    return image

# Preprocesamiento de la imagen
def preprocess_image(img, target_size=(299, 299)):
    # Convertir a RGB si la imagen no está en ese formato
    if img.mode != "RGB":
        img = img.convert("RGB")
    img = img.resize(target_size)
    img_array = img_to_array(img)
    img_array = np.expand_dims(img_array, axis=0)
    img_array = preprocess_input(img_array)
    return img_array
    

# Función de predicción
def predict_fn(img=None,base64_img=None):
    # Verificar si se proporciona una imagen en base64
    if base64_img is not None:
        try:
            img = base64_to_image(base64_img)
        except Exception as e:
            return {"error": "Error decoding base64 image: " + str(e)}
    elif img is None:
        return {"error": "No image provided."}
    
    img = preprocess_image(img)
    
    labels = ['acanthoica', 'akashiwo', 'alexandrium', 'amoeba', 'amphidinium', 'amylax', 'apedinella', 'asterionellopsis', 'bacillaria', 'bacteriastrum', 'biddulphia', 'calciopappus', 'cerataulina', 'ceratium', 'chaetoceros', 'chrysochromulina', 'cochlodinium', 'corethron', 'corymbellus', 'coscinodiscus', 'cryptophyta', 'cylindrotheca', 'dactyliosolen', 'delphineis', 'dictyocha', 'dinobryon', 'dinophysis', 'ditylum', 'emiliania', 'ephemera', 'eucampia', 'euglena', 'gonyaulax', 'guinardia', 'gyrodinium', 'hemiaulus', 'heterocapsa', 'karenia', 'katodinium', 'kryptoperidinium', 'laboea', 'lauderia', 'leptocylindrus', 'licmophora', 'nanoneis', 'odontella', 'ophiaster', 'ostreopsis', 'oxytoxum', 'paralia', 'parvicorbicula', 'phaeocystis', 'pleuronema', 'pleurosigma', 'polykrikos', 'prorocentrum', 'proterythropsis', 'protoperidinium', 'pseudo-nitzschia', 'pseudochattonella', 'pyramimonas', 'rhabdolithes', 'rhizosolenia', 'scrippsiella', 'skeletonema', 'stephanopyxis', 'syracosphaera', 'thalassionema', 'thalassiosira', 'trichodesmium', 'vicicitus', 'warnowia']
    
    preds = model.predict(img)[0]
    top_10_indices = preds.argsort()[-10:][::-1]
    top_10_classes = [labels[i] for i in top_10_indices]
    top_10_probabilities = preds[top_10_indices]
    
    result = {top_10_classes[i]: float(top_10_probabilities[i]) for i in range(10)}
    return result


# Interfaz de Gradio para manejar tanto imágenes como base64
iface = gr.Interface(
    fn=predict_fn,
    inputs=[
        gr.Image(type="pil", label="Imagen (opcional)"),
        gr.Textbox(label="Imagen Base64 (opcional)")
    ],
    outputs=gr.Label(),
    title="Clasificación de Fitoplancton con InceptionV3",
    description="Sube una imagen o ingresa una cadena en base64 para clasificar fitoplancton en 72 géneros.",
    live=False
)

# Lanzar la interfaz
iface.launch()