import gradio as gr
import numpy as np
from PIL import Image
import tensorflow as tf
import joblib

# 1. Load Model CNN (Keras)
try:
    model = tf.keras.models.load_model('cnn_model.h5')
    print("Load model thành công")
except Exception as e:
    print(f"Lỗi load model: {e}")
    model = None

# 2. Load Label Encoder
try:
    label_encoder = joblib.load('label_encoder.joblib')
    print("Load encoder thành công")
except:
    label_encoder = None

def preprocess_image(image):
    # Chuyển sang RGB
    image = image.convert("RGB")
    # Resize đúng kích thước lúc train (64x64)
    image = image.resize((64, 64))

    # Chuyển thành mảng numpy và chuẩn hóa / 255.0
    image_array = np.array(image) / 255.0

    # QUAN TRỌNG: CNN cần input 4 chiều (Batch, Height, Width, Channel)
    # Nên ta phải thêm 1 chiều ở đầu: (64,64,3) -> (1, 64, 64, 3)
    image_array = np.expand_dims(image_array, axis=0)

    return image_array

def predict(image):
    if model is None or label_encoder is None:
        return "Lỗi: Chưa load được model hoặc encoder."

    try:
        # Xử lý ảnh
        processed_img = preprocess_image(image)

        # Dự đoán
        prediction = model.predict(processed_img)

        # Lấy vị trí có xác suất cao nhất (argmax)
        class_index = np.argmax(prediction)

        # Chuyển từ số về tên nhãn
        class_name = label_encoder.inverse_transform([class_index])[0]

        # Lấy độ tin cậy (xác suất)
        confidence = float(np.max(prediction))

        return f"Kết quả: {class_name} ({confidence*100:.2f}%)"
    except Exception as e:
        return f"Lỗi dự đoán: {str(e)}"

# Giao diện Gradio
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Nhận diện Biển báo Giao thông (CNN)",
    description="Upload ảnh biển báo để model CNN dự đoán."
)

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