import spaces
import gradio as gr
import cv2
import numpy as np
from rmn import RMN

# Khởi tạo model RMN.
# Quá trình này sẽ mất một chút thời gian ở lần chạy đầu tiên để tải weights.
m = RMN()

@spaces.GPU
def detect_emotion(image):
    """
    Hàm xử lý ảnh đầu vào, gọi model RMN và trả về ảnh đã vẽ bounding box.
    """
    if image is None:
        return None, "Please upload an image first."
        
    # Xử lý kênh màu: ảnh từ webcam thường có 4 kênh (RGBA), cần chuyển về 3 kênh (RGB)
    if len(image.shape) == 3 and image.shape[2] == 4:
        image = image[:, :, :3]
        
    # Chuyển đổi từ RGB sang OpenCV BGR để tương thích với RMN
    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    
    # Dự đoán cảm xúc
    results = m.detect_emotion_for_single_frame(image_bgr)
    
    # RMN cung cấp sẵn hàm draw() để vẽ bounding box và label lên ảnh
    image_with_boxes = m.draw(image_bgr, results)
    
    # Chuyển ngược lại từ BGR sang RGB để Gradio hiển thị đúng màu
    output_image = cv2.cvtColor(image_with_boxes, cv2.COLOR_BGR2RGB)
    
    # Có thể format lại text kết quả để hiển thị ra màn hình
    results_text = "No face detected"
    if results:
        results_text = "\n".join([f"Face {i+1}: {res['emo_label']} (Prob: {res['emo_proba']:.2f})" for i, res in enumerate(results)])
        
    return output_image, results_text

# Khởi tạo giao diện Gradio
with gr.Blocks(title="Facial Expression Recognition") as demo:
    gr.Markdown("# 🎭 Facial Expression Recognition (RMN)")
    gr.Markdown(
        "Upload an image or use your webcam to detect facial expressions using the **Residual Masking Network (RMN)**. "
        "This model achieves state-of-the-art results on the FER2013 dataset.\n\n"
        "**MẸO SỬ DỤNG WEBCAM:** Hãy bấm vào **biểu tượng chiếc máy ảnh 📷 lơ lửng giữa khung hình** để chụp. "
        "Ngay sau khi bức ảnh tĩnh được chụp xong, hệ thống sẽ **tự động** phân tích khuôn mặt mà không cần bạn bấm thêm nút nào cả!"
    )
    
    with gr.Row():
        with gr.Column():
            # Đầu vào ảnh
            input_image = gr.Image(label="Input Image", type="numpy")
            btn = gr.Button("Detect Emotion", variant="primary")
            
        with gr.Column():
            # Đầu ra ảnh và text
            output_image = gr.Image(label="Detected Result", type="numpy")
            output_text = gr.Textbox(label="Detection Details", lines=3)
            
    # Tự động nhận diện ngay khi có ảnh được upload hoặc chụp xong từ webcam
    input_image.change(fn=detect_emotion, inputs=input_image, outputs=[output_image, output_text])
    
    # Kết nối nút bấm thủ công (phòng hờ trường hợp auto không chạy hoặc user thích tự bấm)
    btn.click(fn=detect_emotion, inputs=input_image, outputs=[output_image, output_text])
    
    # Cho phép demo các ảnh có sẵn
    gr.Examples(
        examples=[
            # Bạn có thể thêm đường dẫn các ảnh mẫu vào đây nếu có
            # ["examples/happy.jpg"],
            # ["examples/sad.jpg"]
        ],
        inputs=input_image
    )

# Chạy ứng dụng
if __name__ == "__main__":
    # share=True tạo ra một đường link public tạm thời (rất tiện để khoe nhanh)
    demo.launch(server_name="0.0.0.0", server_port=7860, share=True)