import gradio as gr
import cv2
import numpy as np
import torch
import kornia as K
from kornia.core import Tensor
from kornia.contrib import FaceDetector, FaceDetectorResult, FaceKeypoint


def draw_keypoint(img: np.ndarray, det: FaceDetectorResult, kpt_type: FaceKeypoint) -> np.ndarray:
    kpt = det.get_keypoint(kpt_type).int().tolist()
    return cv2.circle(img, kpt, 2, (255, 0, 0), 2)


def detect(img_raw):

    # preprocess
    if img_raw is not None and len(img_raw.shape) == 3:
        img = K.image_to_tensor(img_raw, keepdim=False)
        img = K.color.bgr_to_rgb(img.float())


        # create the detector and find the faces !
        face_detection = FaceDetector()

        with torch.no_grad():
            dets = face_detection(img)
        dets = [FaceDetectorResult(o) for o in dets[0]]

        img_vis = img_raw.copy()

        vis_threshold = 0.8

        face_count = 0  # Initialize face counter

        for b in dets:
            if b.score < vis_threshold:
                continue

            # Draw face bounding box
            img_vis = cv2.rectangle(img_vis, b.top_left.int().tolist(), b.bottom_right.int().tolist(), (0, 255, 0), 4)
            face_count += 1  # Increment face counter

            # Draw Keypoints
            img_vis = draw_keypoint(img_vis, b, FaceKeypoint.EYE_LEFT)
            img_vis = draw_keypoint(img_vis, b, FaceKeypoint.EYE_RIGHT)
            img_vis = draw_keypoint(img_vis, b, FaceKeypoint.NOSE)
            img_vis = draw_keypoint(img_vis, b, FaceKeypoint.MOUTH_LEFT)
            img_vis = draw_keypoint(img_vis, b, FaceKeypoint.MOUTH_RIGHT)

        return img_vis, face_count


title = "Crowd Counter"
description = "<p style='text-align: center'>This is a Gradio demo for crowd counting using Kornia's Face Detection model.</p><p style='text-align: center'>To use it, simply upload your image, or click one of the examples to load them</p>"
article = """<p style='text-align: center'><a href='https://kornia.readthedocs.io/en/latest/' target='_blank'>Kornia Docs</a> | <a href='https://github.com/kornia/kornia' target='_blank'>Kornia Github Repo</a> | <a 
href='https://kornia.readthedocs.io/en/latest/applications/face_detection.html' target='_blank'>Kornia Face Detection Tutorial</a></p>"""
outputs = [gr.Image(type="pil", label="Output Image"), gr.Textbox(label="Total (Head) Count")]

examples = ['sample.jpeg']

face = gr.Interface(
    detect,
    gr.Image(type="numpy"),
    outputs=outputs,
    examples=examples,
    title=title,
    description=description,
    article=article,
    live=True,
    allow_flagging="never"
)

face.launch()