import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np

# 모델 로드
model = YOLO("best.pt")

def predict_image(image):
    if image is None:
        return None
    
    # 원본 RGB 이미지 저장
    original_image = np.array(image)
    original_h, original_w = original_image.shape[:2]
    
    # 원본 이미지가 grayscale이면 RGB로 변환
    if len(original_image.shape) == 2:
        original_image = cv2.cvtColor(original_image, cv2.COLOR_GRAY2RGB)
    elif len(original_image.shape) == 3 and original_image.shape[2] == 1:
        original_image = cv2.cvtColor(original_image, cv2.COLOR_GRAY2RGB)
    
    # 전처리용 이미지 복사
    input_image = original_image.copy()
    
    # 전처리: grayscale 변환, 리사이즈 
    # RGB 이미지를 grayscale로 변환
    if len(input_image.shape) == 3 and input_image.shape[2] == 3:
        input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2GRAY)
    # grayscale 이미지를 3채널로 복제 (YOLO의 입력 : 3채널)
    if len(input_image.shape) == 2:
        input_image = cv2.cvtColor(input_image, cv2.COLOR_GRAY2RGB)
    
    # 256x256으로 리사이즈
    input_image = cv2.resize(input_image, (256, 256))
    
    # 모델 추론
    results = model.predict(
        source=input_image,
        conf=0.25,
        iou=0.7,
        imgsz=256,
        save=False  # 임시 파일 저장하지 않음
    )
    
    # bbox 좌표를 원본 크기로 스케일링
    result = results[0]
    boxes = result.boxes
    
    # 스케일 비율 계산 (256x256 -> 원본 크기)
    scale_x = original_w / 256.0
    scale_y = original_h / 256.0
    
    # 원본 RGB 이미지에 bbox 그리기
    result_image = original_image.copy()
    
    for box in boxes:
        # bbox 좌표 가져오기 (256x256 기준)
        x1, y1, x2, y2 = box.xyxy[0].cpu().numpy()
        
        # 원본 크기로 스케일링
        x1 = int(x1 * scale_x)
        y1 = int(y1 * scale_y)
        x2 = int(x2 * scale_x)
        y2 = int(y2 * scale_y)
        
        # 클래스와 confidence 가져오기
        cls = int(box.cls[0].cpu().numpy())
        conf = float(box.conf[0].cpu().numpy())
        
        # bbox 그리기
        cv2.rectangle(result_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        
        # 라벨 텍스트
        label = f"{result.names[cls]} {conf:.2f}"
        
        # 텍스트 배경 그리기
        (text_width, text_height), baseline = cv2.getTextSize(
            label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
        )
        cv2.rectangle(
            result_image,
            (x1, y1 - text_height - baseline - 5),
            (x1 + text_width, y1),
            (0, 255, 0),
            -1
        )
        
        # 텍스트 그리기
        cv2.putText(
            result_image,
            label,
            (x1, y1 - baseline - 2),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 0, 0),
            1
        )
    
    return result_image

# Gradio Interface 생성
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil", label="입력 이미지"),
    outputs=gr.Image(type="numpy", label="탐지 결과"),
    title="YOLO 객체 탐지 데모",
    description="이미지를 업로드하면 YOLO 모델이 객체를 탐지합니다.",
    examples=None
)

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