import cv2 as cv
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download
from yunet import YuNet
from ediffiqa import eDifFIQA

# Download face detection model (YuNet)
model_path_yunet = hf_hub_download(
    repo_id="opencv/face_detection_yunet",
    filename="face_detection_yunet_2023mar.onnx"
)

# Download face quality assessment model (eDifFIQA Tiny)
model_path_quality = hf_hub_download(
    repo_id="opencv/face_image_quality_assessment_ediffiqa",
    filename="ediffiqa_tiny_jun2024.onnx"
)

# Backend and target
backend_id = cv.dnn.DNN_BACKEND_OPENCV
target_id = cv.dnn.DNN_TARGET_CPU

# Initialize YuNet for face detection
face_detector = YuNet(
    modelPath=model_path_yunet,
    inputSize=[320, 320],
    confThreshold=0.9,
    nmsThreshold=0.3,
    topK=5000,
    backendId=backend_id,
    targetId=target_id
)

# Initialize eDifFIQA for quality assessment
quality_model = eDifFIQA(
    modelPath=model_path_quality,
    inputSize=[112, 112]
)
quality_model.setBackendAndTarget(
    backendId=backend_id,
    targetId=target_id
)

REFERENCE_FACIAL_POINTS = np.array([
    [38.2946  , 51.6963  ],
    [73.5318  , 51.5014  ],
    [56.0252  , 71.7366  ],
    [41.5493  , 92.3655  ],
    [70.729904, 92.2041  ]
], dtype=np.float32)

def align_image(image, detection_data):
    src_pts = np.float32(detection_data[0][4:-1]).reshape(5, 2)
    tfm, _ = cv.estimateAffinePartial2D(src_pts, REFERENCE_FACIAL_POINTS, method=cv.LMEDS)
    face_img = cv.warpAffine(image, tfm, (112, 112))
    return face_img

def assess_face_quality(input_image):
    bgr_image = cv.cvtColor(input_image, cv.COLOR_RGB2BGR)
    h, w, _ = bgr_image.shape

    face_detector.setInputSize([w, h])
    detections = face_detector.infer(bgr_image)

    if detections is None or len(detections) == 0:
        return "No face detected.", input_image

    aligned_face = align_image(bgr_image, detections)
    score = np.squeeze(quality_model.infer(aligned_face)).item()

    output_image = aligned_face.copy()
    cv.putText(output_image, f"{score:.3f}", (0, 20), cv.FONT_HERSHEY_DUPLEX, 0.8, (0, 0, 255), 2)
    output_image = cv.cvtColor(output_image, cv.COLOR_BGR2RGB)

    return f"Quality Score: {score:.3f}", output_image

# Gradio Interface
with gr.Blocks(css='''.example * {
    font-style: italic;
    font-size: 18px !important;
    color: #0ea5e9 !important;
    }''') as demo:

    gr.Markdown("### Face Image Quality Assessment (eDifFIQA + YuNet)")
    gr.Markdown("Upload a face image. The app detects and aligns the face, then evaluates image quality using the eDifFIQA model.")

    with gr.Row():
        input_image = gr.Image(type="numpy", label="Upload Face Image")
        with gr.Column():
            quality_score = gr.Text(label="Quality Score")
            aligned_face = gr.Image(type="numpy", label="Aligned Face with Score")

    # Clear output when new image is uploaded
    input_image.change(fn=lambda: ("", None), outputs=[quality_score, aligned_face])

    with gr.Row():
        submit_btn = gr.Button("Submit", variant="primary")
        clear_btn = gr.Button("Clear")

    submit_btn.click(fn=assess_face_quality, inputs=input_image, outputs=[quality_score, aligned_face])
    clear_btn.click(fn=lambda: (None, "", None), outputs=[input_image, quality_score, aligned_face])

    gr.Markdown("Click on any example to try it.", elem_classes=["example"])

    gr.Examples(
        examples=[
            ["examples/lena.jpg"],
            ["examples/gray_face.png"]
        ],
        inputs=input_image
    )

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