main3.py

import gradio as gr
from face_engine.engine import FaceEngine
from PIL import Image

# Danh sách mô hình cho người dùng lựa chọn
model_list = [
    "ir18_webface4m", 
    "ir50_webface4m", 
    "ir101_webface4m", 
    "ir101_webface12m", 
    "vit_base_kprpe_webface4m", 
    "vit_base_kprpe_webface12m", 
    "vit_base_webface4m"
]

# Khởi tạo FaceEngine
face_engine = FaceEngine()

# Hàm cập nhật mô hình dựa trên lựa chọn của người dùng
def update_model(model_name):
    face_engine.set_model(model_name)
    return f"Đã chuyển đổi sang mô hình: {model_name}"

# Giao diện chọn mô hình
model_selector = gr.Dropdown(
    choices=model_list, 
    label="Chọn mô hình trích xuất đặc trưng", 
    value="vit_base_kprpe_webface4m", 
    type="value"
)

# Hàm để đăng ký người dùng mới
def enroll_new_person(pil_image, user_id):
    # Kiểm tra xem khuôn mặt đã tồn tại trong hệ thống chưa
    check_result = face_engine.get_id(pil_image)
    
    if check_result['status'] == "not_found":
        # Nếu khuôn mặt chưa có trong hệ thống, tiến hành đăng ký
        face_engine.save_to_db(pil_image, user_id)
        return f"Khuôn mặt đã đăng ký thành công với ID là {user_id}."
    
    # Nếu khuôn mặt đã có trong hệ thống
    return f"Khuôn mặt đã tồn tại với ID là {check_result['id']}."

# Giao diện đăng ký (enroll)
enroll_interface = gr.Interface(
    fn=enroll_new_person, 
    inputs=[gr.Image(), gr.Textbox(label="Nhập ID (ví dụ: CCCD)")], 
    outputs="text",
    title="Đăng ký khuôn mặt mới",
    description="Tải lên ảnh và nhập ID của bạn để đăng ký vào hệ thống."
)

# Hàm để xác minh hai khuôn mặt có giống nhau không
def verify(pil_image_1, pil_image_2):
    cosine_similarity = face_engine.compute_cosine_similarity(pil_image_1, pil_image_2)
    if cosine_similarity > face_engine.threshold:
        return f"Hai ảnh là cùng một người, điểm tương đồng: {cosine_similarity:.2f}"
    return f"Hai ảnh là hai người khác nhau, điểm tương đồng: {cosine_similarity:.2f}"

# Giao diện xác minh (verify)
verification_interface = gr.Interface(
    fn=verify, 
    inputs=[gr.Image(), gr.Image()], 
    outputs="text",
    title="Xác minh khuôn mặt",
    description="Tải lên hai ảnh để xác minh xem chúng có phải cùng một người hay không."
)

# Hàm để nhận diện khuôn mặt từ ảnh đầu vào
def identify(pil_image):
    check_result = face_engine.get_id(pil_image)
    
    if check_result['status'] == "not_found":
        return "Khuôn mặt không tồn tại trong hệ thống."
    
    # Nếu tìm thấy khuôn mặt
    return f"Khuôn mặt được nhận diện với ID là {check_result['id']}."

# Giao diện nhận diện (identify)
identification_interface = gr.Interface(
    fn=identify, 
    inputs=gr.Image(), 
    outputs="text",
    title="Nhận diện khuôn mặt",
    description="Tải lên một ảnh để nhận diện khuôn mặt từ hệ thống."
)

# Giao diện chính với các tab
demo = gr.TabbedInterface(
    interface_list=[
        gr.Interface(fn=update_model, inputs=model_selector, outputs="text", title="Chọn Mô Hình"),
        enroll_interface,
        verification_interface,
        identification_interface
    ], 
    tab_names=["Chọn Mô Hình", "Enroll", "Verification", "Identification"]
)

# Khởi chạy ứng dụng Gradio
if __name__ == "__main__":
    demo.launch()