import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from tensorflow.keras.applications.efficientnet import preprocess_input

# Đường dẫn model SavedModel
MODEL_PATH = "exported_model"
IMG_SIZE = (224, 224)
CLASS_NAMES = ['bad', 'good', 'very_good']

# Load model
model = tf.saved_model.load(MODEL_PATH)

infer = model.signatures["serving_default"]

def predict_guava_quality(img_input):
    if img_input is None:
        return "❌ Vui lòng tải ảnh", 0.0

    # Convert image
    img = Image.fromarray(img_input).convert("RGB")
    img = img.resize(IMG_SIZE)

    arr = np.array(img).astype("float32")
    arr = preprocess_input(arr)
    arr = np.expand_dims(arr, axis=0)

    # TensorFlow serving
    outputs = infer(tf.constant(arr))
    preds = list(outputs.values())[0].numpy()[0]

    idx = np.argmax(preds)
    confidence = preds[idx]
    label = CLASS_NAMES[idx]

    return f"✅ Kết quả: {label}", float(confidence)

demo = gr.Interface(
    fn=predict_guava_quality,
    inputs=gr.Image(type="numpy", label="Tải ảnh Quả Ổi"),
    outputs=[
        gr.Textbox(label="Dự đoán"),
        gr.Number(label="Độ tin cậy (%)", precision=4)
    ],
    title="Phân loại chất lượng Ổi",
    description="Model EfficientNetB0 | very_good / good / bad"
)

if __name__ == "__main__":
    demo.launch()