image_track commit

app.py

@@ -1,7 +1,42 @@
+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 greet(name):
-    return "Hello " + name + "!"
+def detect_faces(input_img):
+    # download model
+    model_path = hf_hub_download(
+    repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt"
+    )
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()   
+    # 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()