import gradio as gr
from ultralytics import YOLO
import os
import cv2
from PIL import Image

# 모델 로드
model_path = "falldetect-11x.pt"
model = YOLO(model_path)

def detect_fall(image):
    # 이미지에서 낙상 감지 수행
    results = model(image)
    result_img = results[0].plot()
    
    result_img_rgb = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
        
    # numpy 배열을 PIL 이미지로 변환하여 반환
    return Image.fromarray(result_img_rgb)

# Gradio 인터페이스 생성
interface = gr.Interface(
    fn=detect_fall,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(),
    title="Fall detect",
    description="Fall detect",
    # examples=[["example1.jpg"], ["example2.jpg"]]  # 예시 이미지가 있으면 추가
)

interface.launch()