import streamlit as st
import requests
from PIL import Image
from io import BytesIO
from transformers import AutoTokenizer, AutoModelForCausalLM

# Configuración de la API
API_URL = "https://api-inference.huggingface.co/models/riabayonaor/modelo_prediccion_enfermedades_pepinos"

# Cargar el modelo y el tokenizador de BLOOM
tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom")
model = AutoModelForCausalLM.from_pretrained("bigscience/bloom")

def query(image_bytes):
    response = requests.post(API_URL, data=image_bytes)
    if response.status_code != 200:
        st.error(f"Error al clasificar la imagen: {response.status_code}")
        return None
    return response.json()

def bloom_query(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=500, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def main():
    st.set_page_config(page_title="Predicción de Enfermedades en Pepinos")
    st.title("Predicción de Enfermedades en Pepinos")
    st.write("Sube una foto de una planta de pepino o un pepino para clasificar posibles enfermedades y obtener soluciones.")

    uploaded_file = st.file_uploader("Sube una foto de una planta de pepino o un pepino", type=["jpg", "jpeg", "png"])

    if uploaded_file is not None:
        image = Image.open(uploaded_file)
        st.image(image, caption='Imagen subida.', use_column_width=True)
        st.write("Clasificando...")

        # Convertir la imagen a bytes
        img_byte_arr = BytesIO()
        image.save(img_byte_arr, format='PNG')
        img_byte_arr = img_byte_arr.getvalue()

        # Enviar la imagen al modelo de Hugging Face
        predictions = query(img_byte_arr)

        if predictions is not None:
            if "error" not in predictions:
                # Suponiendo que las predicciones están en el formato [{label: "Downy Mildew", score: 0.95}, {label: "Fresh Cucumber", score: 0.05}]
                top_prediction = max(predictions, key=lambda x: x["score"])
                class_label = top_prediction['label']
                st.write(f"Predicción principal: {class_label} con confianza {top_prediction['score']:.2f}")

                # Usar la etiqueta principal para el modelo BLOOM
                prompt = f"Esta enfermedad es {class_label}. Explica qué es y sugiere posibles insecticidas o soluciones en español."

                # Llamar al modelo BLOOM
                explanation = bloom_query(prompt)
                
                st.write(f"Esta enfermedad es {class_label}:")
                st.write(explanation)
            else:
                st.write("No se pudo clasificar la imagen.")
        else:
            st.write("No se pudo clasificar la imagen.")

if __name__ == "__main__":
    main()