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

# --- CẤU HÌNH QUAN TRỌNG ---
# 1. Tên file model (Phải khớp chính xác tên file bạn đã upload)
MODEL_FILENAME = "model.h5"

# 2. Kích thước ảnh model yêu cầu
# BẮT BUỘC PHẢI CÓ. Bạn cần nhớ lại lúc train bạn chọn size bao nhiêu.
# Nếu là 150 thì giữ nguyên, nếu khác thì sửa số này.
IMG_SIZE = 150

# --- LOAD MODEL ---
try:
    model = tf.keras.models.load_model(MODEL_FILENAME)
    print(f"Đã load thành công model: {MODEL_FILENAME}")
except Exception as e:
    print(f"LỖI NGHIÊM TRỌNG: Không tìm thấy file {MODEL_FILENAME}. Hãy kiểm tra lại tab Files.")
    model = None

# --- HÀM XỬ LÝ VÀ DỰ ĐOÁN ---
def du_doan_nguoi(image):
    if model is None:
        return "Lỗi: Chưa load được model."
    
    if image is None:
        return "Vui lòng tải ảnh lên."
        
    # --- BƯỚC CHUẨN HÓA ẢNH (BẮT BUỘC) ---
    # Dù người dùng up ảnh size nào, ta cũng phải đưa về đúng chuẩn của model.
    # Nếu bỏ bước này, code sẽ báo lỗi ngay lập tức.
    image_resized = image.resize((IMG_SIZE, IMG_SIZE))
    
    # Chuyển sang dạng số và chuẩn hóa về khoảng 0-1
    img_array = np.array(image_resized)
    img_array = img_array / 255.0
    
    # Thêm chiều batch (ví dụ từ (150,150,3) thành (1,150,150,3))
    img_array = np.expand_dims(img_array, axis=0)
    
    # --- DỰ ĐOÁN ---
    prediction = model.predict(img_array)[0]
    
    # --- XỬ LÝ KẾT QUẢ ---
    # Cần xác định quy ước lúc train: 0 là gì, 1 là gì?
    # GIẢ ĐỊNH PHỔ BIẾN: 0 = Không phải người, 1 = Người.
    # Nếu kết quả bị ngược, hãy đảo vị trí hai nhãn này cho nhau.
    labels = ['Con người','Không phải người'] 
    
    if len(prediction) == 1:
        # Dạng 1 đầu ra (Sigmoid)
        score_nguoi = float(prediction[0])
        return {labels[1]: score_nguoi, labels[0]: 1 - score_nguoi}
    else:
        # Dạng 2 đầu ra (Softmax)
        return {labels[i]: float(prediction[i]) for i in range(len(labels))}

# --- GIAO DIỆN WEB ---
demo = gr.Interface(
    fn=du_doan_nguoi,
    inputs=gr.Image(type="pil", label="Tải ảnh của bạn lên đây (Kích thước bất kỳ)"),
    outputs=gr.Label(num_top_classes=2, label="Kết quả phân tích AI"),
    title="🤖 Web Nhận Diện: Con Người hay Không?",
    description=f"Tải lên một bức ảnh bất kỳ, AI sẽ kiểm tra xem trong ảnh có phải là con người hay không. (Model đang chạy ở kích thước {IMG_SIZE}x{IMG_SIZE})",
    theme="default"
)

# Chạy ứng dụng
if __name__ == "__main__":
    # Thêm debug=True để dễ xem lỗi nếu có
    demo.launch()