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