main1.py

import gradio as gr
from face_engine.engine import FaceEngine
import numpy as np
# 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):
    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ý
        new_id = face_engine.save_to_db(pil_image)
        return f"Khuôn mặt đăng ký thành công với ID là {new_id}. Thời gian trích xuất: {check_result['extraction_time']:.2f} giây"
    
    # 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']}. Thời gian trích xuất: {check_result['extraction_time']:.2f} giây"


enroll = gr.Interface(enroll_new_person, gr.Image(), "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}"

verification = gr.Interface(verify, [gr.Image(),gr.Image()], "text")

def identify(pil_image):
    check_result = face_engine.get_id(pil_image)
    
    if check_result['status'] == "not_found":
        return f"Khuôn mặt không tồn tại trong hệ thống. Thời gian trích xuất: {check_result['extraction_time']:.2f} giây"
    
    return f"Khuôn mặt được đăng ký với ID là {check_result['id']}. Thời gian trích xuất: {check_result['extraction_time']:.2f} giây"


identification = gr.Interface(identify, gr.Image(), "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__":
     # Comment this for deploying
    demo.launch()