import tensorflow as tf
from tensorflow.keras.models import load_model
from PIL import Image
import numpy as np
import gradio as gr

# Charger le modèle sauvegardé
model = load_model('model_unet_best.keras')

# Prétraitement de l'image pour l'adapter au modèle
def preprocess_image(image):
    img = image.resize((256, 256))
    img_array = np.array(img) / 255.0 
    
    if img_array.ndim == 2 or img_array.shape[-1] == 1:
        img_array = np.stack([img_array] * 3, axis=-1)

    img_array = np.expand_dims(img_array, axis=0)
    return img_array

# Seegmentation des images de feuilles de pommiers

def segment_image(image):
    
    img_array = preprocess_image(image)
    
    segmented_img = model.predict(img_array)[0] 
    
    segmented_img = np.squeeze(segmented_img)

    if segmented_img.ndim == 2 or segmented_img.shape[-1] == 1:
        segmented_img = np.stack([segmented_img] * 3, axis=-1)  

    # Vérifier le type de données et convertir en uint8 si cela n'a pas dejà été fait
    segmented_img = (segmented_img * 255).astype(np.uint8)

    return Image.fromarray(segmented_img)

def analyze_health(segmented_image):
    
    diagnostic = "La plante semble en bonne santé."
    
    return diagnostic



# Fonction principale qui combine la segmentation et le diagnostic des feuilles de pommiers
def plant_health_analysis(image):
    segmented_img = segment_image(image)
    diagnostic = analyze_health(segmented_img)
    
    return segmented_img, diagnostic

# Interface Gradio

interface = gr.Interface(
    fn=plant_health_analysis,
    inputs=gr.inputs.Image(type="pil", label="Téléchargez une image de plante"),
    outputs=[
        gr.outputs.Image(type="pil", label="Image Segmentée"),
        gr.outputs.Textbox(label="Diagnostic de la Plante")
    ],
    title="PomSafe",
    description="Téléchargez une image pour analyser l'état de santé de votre plante et détecter d'éventuelles anomalies.",
    theme="default",
    layout="horizontal"
)
interface.launch(share=True)