import cv2
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from supervision import Detections
from PIL import Image
import torch
import numpy as np
import gradio as gr

# 함수 정의: 사각형 그리고 정확도 표시
def draw_rect_with_conf(image, detections):
    for detection in detections:
        # 탐지된 객체의 경계 상자 좌표 (NumPy 배열)와 신뢰도 점수
        box, _, conf, _, _ = detection
        x1, y1, x2, y2 = box.astype(int)

        # 사각형 그리기
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)

        # 신뢰도 표시
        cv2.putText(image, f'Accuracy: {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    return image

def detect_faces(input_img):
    # download model
    model_path = hf_hub_download(
    repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt"
    )

    # load model
    model = YOLO(model_path)
    
    # 이미지를 YOLO 모델 입력 형식으로 변환
    image_cv = cv2.cvtColor(np.array(input_img), cv2.COLOR_RGB2BGR)

    # 얼굴 탐지
    output = model(input_img)
    results = Detections.from_ultralytics(output[0])

    # 탐지 결과 그리기
    drawn_image = draw_rect_with_conf(image_cv, results)

    # OpenCV 이미지를 PIL 이미지로 변환
    drawn_image_pil = Image.fromarray(cv2.cvtColor(drawn_image, cv2.COLOR_BGR2RGB))
    return drawn_image_pil

def gradio_interface(input_img):
    # 얼굴 탐지 함수 호출
    detected_img = detect_faces(input_img)
    return detected_img

# Gradio 인터페이스 설정
demo = gr.Interface(fn=gradio_interface, inputs=gr.Image(type="pil"), outputs="image")

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