Spaces:
Sleeping
Sleeping
File size: 3,402 Bytes
23d40fc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
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()
|