app.py

import os
import numpy as np
import tensorflow as tf
from PIL import Image
import gradio as gr
import pickle

# Thông số
IMG_HEIGHT = 64
IMG_WIDTH = 64

# Load model và label encoder
model = tf.keras.models.load_model('traffic_sign_model.keras')
with open('label_encoder.pkl', 'rb') as f:
    le = pickle.load(f)

def predict_traffic_sign(image):
    """
    Hàm dự đoán biển báo giao thông từ ảnh đầu vào
    """
    try:
        # Xử lý ảnh đầu vào
        img = Image.fromarray(image.astype('uint8')).convert('RGB')
        img = img.resize((IMG_HEIGHT, IMG_WIDTH))
        img_array = np.array(img).astype('float32') / 255.0
        img_array = np.expand_dims(img_array, axis=0)
        
        # Dự đoán
        predictions = model.predict(img_array, verbose=0)
        predicted_class_idx = np.argmax(predictions[0])
        confidence = predictions[0][predicted_class_idx]
        
        # Lấy tên lớp
        predicted_class_name = le.inverse_transform([predicted_class_idx])[0]
        
        # Tạo dictionary kết quả cho tất cả các lớp
        results = {}
        for idx, class_name in enumerate(le.classes_):
            results[class_name] = float(predictions[0][idx])
        
        return results
    
    except Exception as e:
        return {f"Error: {str(e)}": 0.0}

# Tạo Gradio Interface
demo = gr.Interface(
    fn=predict_traffic_sign,
    inputs=gr.Image(label="Tải ảnh biển báo giao thông"),
    outputs=gr.Label(num_top_classes=5, label="Kết quả dự đoán"),
    title="🚦 Nhận diện Biển báo Giao thông",
    description="""
    **Upload một ảnh biển báo giao thông để nhận diện.**
    
    Model CNN được huấn luyện để phân loại các loại biển báo giao thông Việt Nam.
    
    📊 Kết quả hiển thị top 5 dự đoán có xác suất cao nhất.
    """,
    examples=[
        # Thêm đường dẫn đến ảnh mẫu trong thư mục examples/
        # ["examples/stop_sign.jpg"],
        # ["examples/speed_limit.jpg"],
    ],
    theme=gr.themes.Soft(),
    allow_flagging="never",
    analytics_enabled=False
)