Upload app.py

app.py

@@ -0,0 +1,131 @@
+import gradio as gr
+from ultralytics import YOLO
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
+import pandas as pd
+from collections import Counter
+import os
+
+# 모델 로드
+model = YOLO("best.pt")
+
+# 클래스별 색상 정의
+class_colors = {
+    'freshripe': (50, 205, 50),
+    'freshunripe': (173, 255, 47),
+    'overripe': (255, 165, 0),
+    'ripe': (0, 128, 0),
+    'rotten': (128, 0, 0),
+    'unripe': (255, 255, 0)
+}
+
+def predict_image(image):
+    if image is None:
+        return None, "이미지를 업로드해주세요."
+    
+    # 이미지 복사본 생성
+    img = image.copy()
+    h, w = img.shape[:2]
+    
+    # 모델 예측
+    results = model.predict(
+        source=img,
+        imgsz=640,  # 기본 이미지 크기
+        conf=0.5,   # 신뢰도 임계값
+        save=False,
+        verbose=False
+    )
+    
+    result = results[0]
+    result_img = result.plot()  # 모델이 그린 예측 결과
+    
+    # 결과 분석
+    boxes = result.boxes.data.cpu().numpy()
+    cls_names = result.names
+    
+    class_counts = Counter()
+    result_text = "예측 결과:\n"
+    
+    # matplotlib 그림 생성
+    fig, ax = plt.subplots(1, figsize=(12, 8))
+    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    ax.imshow(img_rgb)
+    ax.set_title("바나나 익음 분류 결과")
+    
+    for box in boxes:
+        x1, y1, x2, y2, conf, cls_id = box
+        pred_class = cls_names[int(cls_id)]
+        class_counts[pred_class] += 1
+        
+        # 결과 텍스트 추가
+        result_text += f"- 클래스: {pred_class}, 확률: {conf:.2f}\n"
+        result_text += f"  위치: [{int(x1)}, {int(y1)}, {int(x2)}, {int(y2)}]\n"
+        
+        # 박스 그리기
+        rect = plt.Rectangle((x1, y1), x2-x1, y2-y1,
+                           fill=False,
+                           edgecolor=tuple(c/255 for c in class_colors.get(pred_class, (255, 0, 0))),
+                           linewidth=2)
+        ax.add_patch(rect)
+        
+        # 레이블 텍스트 추가
+        label = f"{pred_class} {conf:.2f}"
+        ax.text(x1, y1-5, label,
+               color='white',
+               fontsize=10,
+               bbox=dict(facecolor=tuple(c/255 for c in class_colors.get(pred_class, (255, 0, 0))),
+                        alpha=0.7,
+                        pad=2))
+    
+    ax.axis('off')
+    
+    # 클래스별 개수 합계 추가
+    result_text += "\n클래스별 개수:\n"
+    for cls_name, count in class_counts.items():
+        result_text += f"- {cls_name}: {count}개\n"
+    
+    if len(boxes) == 0:
+        result_text += "\n바나나가 감지되지 않았습니다."
+    
+    # 그림을 이미지로 변환
+    plt.tight_layout()
+    # 그림을 메모리상의 이미지로 저장
+    fig.canvas.draw()
+    vis_img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    vis_img = vis_img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    plt.close(fig)
+    
+    return vis_img, result_text
+
+# Gradio 인터페이스 생성
+with gr.Blocks(title="바나나 익음 분류기") as demo:
+    gr.Markdown("# 🍌 바나나 익음 분류기")
+    gr.Markdown("바나나 사진을 업로드하면 익음 정도를 분석해 드립니다.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="numpy")
+            submit_btn = gr.Button("분석하기")
+        
+        with gr.Column():
+            output_image = gr.Image(type="numpy")
+            output_text = gr.Textbox(label="분석 결과", lines=10)
+    
+    submit_btn.click(
+        fn=predict_image,
+        inputs=input_image,
+        outputs=[output_image, output_text]
+    )
+    
+    # 예시 이미지 추가 (예시 이미지를 업로드했다면)
+    if os.path.exists("examples"):
+        example_images = [f"examples/{f}" for f in os.listdir("examples") if f.endswith(('.jpg', '.jpeg', '.png'))]
+        if example_images:
+            gr.Examples(
+                examples=example_images,
+                inputs=input_image
+            )
+
+demo.launch()