File size: 3,157 Bytes
a57d7b0 |
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 |
# tab 1 : 사용자가 test 데이터 셋에서 원하는 번호를 선택하여 결과 확인
# tabl 2 : 사용자가 이미지 업로드
import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import yaml
import os
import zipfile
# 모델 로드
model = YOLO("best.pt")
# 클래스 이름 로드
with open("data.yaml", "r") as f:
class_names = yaml.safe_load(f)["names"]
# test.zip 파일 압축 해제
with zipfile.ZipFile("test.zip", 'r') as zip_ref:
zip_ref.extractall("test") # 폴더명은 원하는 경로로 지정
# 객체 탐지 함수
def detect_objects(image):
results = model(image)[0]
detected_classes = []
for box in results.boxes:
cls_id = int(box.cls[0].item())
class_name = class_names[cls_id]
detected_classes.append(class_name)
result_image = Image.fromarray(results.plot())
summary = ", ".join(set(detected_classes)) if detected_classes else "탐지된 객체 없음"
return result_image, f"탐지된 클래스: {summary}"
# 이미지 파일 리스트 가져오기
def get_image_file_list():
image_dir = "test/images"
file_list = sorted([
f for f in os.listdir(image_dir) if f.endswith((".jpg", ".png"))
])
return file_list
# 선택된 인덱스로 이미지 로드
def load_image_by_index(index_display: int):
index = index_display - 1 # 사용자에게는 1부터 보여줬기 때문에
file_list = get_image_file_list()
if 0 <= index < len(file_list):
image_path = os.path.join("test/images", file_list[index])
return Image.open(image_path).convert("RGB")
else:
return Image.new("RGB", (300, 300), color="gray")
# Gradio 앱 시작
with gr.Blocks() as demo:
gr.Markdown("## ✌ YOLO 객체 탐지 템플릿 2 - 탭 기반 멀티기능 데모")
gr.Markdown("Tab 1: 테스트 이미지 선택 / Tab 2: 사용자 이미지 업로드")
with gr.Tab("📁 테스트 데이터셋에서 실행"):
file_list = get_image_file_list()
index_choices = [i + 1 for i in range(len(file_list))] # 1부터 시작
dropdown = gr.Dropdown(choices=index_choices, value=1, label="테스트 이미지 번호 선택")
load_btn = gr.Button("이미지 불러오기")
selected_image = gr.Image(label="선택된 이미지", type="pil")
detect_btn = gr.Button("탐지 실행")
result_img = gr.Image(label="탐지 결과 이미지")
result_txt = gr.Textbox(label="탐지된 클래스 요약")
load_btn.click(fn=load_image_by_index, inputs=dropdown, outputs=selected_image)
detect_btn.click(fn=detect_objects, inputs=selected_image, outputs=[result_img, result_txt])
with gr.Tab("🖼️ 사용자 이미지 업로드"):
user_img = gr.Image(type="pil", label="이미지 업로드")
user_btn = gr.Button("탐지 실행")
user_result_img = gr.Image(label="탐지 결과 이미지")
user_result_txt = gr.Textbox(label="탐지된 클래스 요약")
user_btn.click(fn=detect_objects, inputs=user_img, outputs=[user_result_img, user_result_txt])
demo.launch()
|