Untitled

import gradio as gr
from face_engine.engine import FaceEngine

# 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"]

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")

def enroll_new_person(pil_image, user_id):
    check_result = face_engine.get_id(pil_image)
    
    if check_result['status'] == "not_found":
        # Khuôn mặt không tồn tại, 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}."
    
    # Khuôn mặt đã tồn tại
    return f"Khuôn mặt đã tồn tại với ID là {check_result['id']}."

# Giao diện đăng ký (enroll)
enroll = gr.Interface(enroll_new_person, inputs=[gr.Image(), gr.Textbox(label="Nhập ID (ví dụ: CCCD)")], outputs="text")

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 khuôn mặt của cùng một người, điểm tương đồng: {cosine_similarity}"
    return f"Hai ảnh khuôn mặt của hai người khác nhau, điểm tương đồng: {cosine_similarity}"

# Giao diện xác minh (verify)
verification = gr.Interface(verify, inputs=[gr.Image(), gr.Image()], outputs="text")

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."
    
    return f"Khuôn mặt được đăng ký với ID là {check_result['id']}."

# Giao diện nhận dạng (identify)
identification = gr.Interface(identify, inputs=gr.Image(), outputs="text")

# Tạo giao diện với tab chọn mô hình và các chức năng khác
demo = gr.TabbedInterface([
    gr.Interface(update_model, model_selector, "text"),
    enroll,
    verification,
    identification
], ["Chọn Mô Hình", "Enroll", "Verification", "Identification"])

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