오렌지랑 사과를 판별할 수 있을까요...?

1. 오렌지랑 사과 이미지 다운로드 받기

 

https://www.kaggle.com/datasets?search=apple+image+dataset

Find Open Datasets and Machine Learning Projects | Kaggle

 

 

 

2. 다운로드 받은 이미지로 로보플로우 인식

 

https://app.roboflow.com/kimyeonju/projects

Sign in to Roboflow

 

 

 

3. 딥러닝 학습시키기

(이전 블로그 참고)

 

 

4. 딥러닝 한 best 파일로 파이썬 코드 작성 

import cv2
import torch
from ultralytics import YOLO
from collections import deque
from PIL import ImageFont, ImageDraw, Image
import numpy as np
import os # 경로 확인을 위해 추가

# =========================
# 설정
# =========================
# [수정] dy15_fruits -> day15_fruits (경로 오타 수정)
MODEL_PATH = r"C:\sundo\AIpython\AI\day15_fruits\datasets\runs\detect\orange_yolo8\weights\best.pt"

# 사과/오렌지 best.pt
CAM_INDEX = 0
CONF_THRESHOLD = 0.5          # conf는 낮게 (참고용)
HISTORY_SIZE = 10              # 최근 프레임 수
CONFIRM_COUNT = 3              # 3프레임 이상이면 확정

# =========================
# CUDA 확인
# =========================
# [수정] is_available() 줄바꿈 오타 수정
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("[INFO] CUDA available:", torch.cuda.is_available())
if DEVICE == "cuda":
    print("[INFO] GPU:", torch.cuda.get_device_name(0))

# =========================
# YOLO 모델 로드
# =========================
# [주석] 아래 경로에 파일이 없으면 FileNotFoundError가 발생합니다. 
# 반드시 탐색기에서 해당 위치에 best.pt가 있는지 확인하세요.
if not os.path.exists(MODEL_PATH):
    print(f"[ERROR] 모델 파일을 찾을 수 없습니다: {MODEL_PATH}")
    exit()

model = YOLO(MODEL_PATH).to(DEVICE)
print("[INFO] YOLO model loaded on", DEVICE)
print("[INFO] model.names =", model.names)   # {0:'apple', 1:'orange'}

# =========================
# 카메라 열기 (Windows)
# =========================
cap = cv2.VideoCapture(CAM_INDEX, cv2.CAP_DSHOW)
if not cap.isOpened():
    print("[ERROR] 카메라를 열 수 없습니다.")
    exit()
print("[OK] 카메라 연결 성공")

# =========================
# 한글 출력 준비 (PIL)
# =========================
FONT_PATH = "C:/Windows/Fonts/malgun.ttf"
# [주석] 폰트 파일이 없을 경우를 대비해 예외 처리가 필요할 수 있습니다.
try:
    font = ImageFont.truetype(FONT_PATH, 40)
except:
    font = ImageFont.load_default()
    print("[WARN] 한글 폰트를 로드할 수 없어 기본 폰트를 사용합니다.")

def put_korean_text(img, text, pos, color=(255, 0, 0)):
    """OpenCV 이미지에 한글 출력"""
    img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    draw = ImageDraw.Draw(img_pil)
    draw.text(pos, text, font=font, fill=color)
    return cv2.cvtColor(np.array(img_pil), cv2.COLOR_RGB2BGR)

# =========================
# 프레임 누적 버퍼
# =========================
history = deque(maxlen=HISTORY_SIZE)

# =========================
# 실시간 인식 루프
# =========================
while True:
    ret, frame = cap.read()
    if not ret:
        print("[ERROR] 프레임 읽기 실패")
        break

    # YOLO 추론
    results = model.predict(
        frame,
        conf=CONF_THRESHOLD,
        device=DEVICE,
        verbose=False
    )

    r0 = results[0]

    # 이번 프레임 결과 초기화
    best_conf = 0.0
    current_label = None

    if r0.boxes is not None:
        for box in r0.boxes:
            conf = float(box.conf[0])
            cls_id = int(box.cls[0])
            label = model.names.get(cls_id, "unknown")

            # 가장 높은 conf 하나 선택
            if conf > best_conf:
                best_conf = conf
                current_label = label

            # 바운딩 박스 표시
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(
                frame,
                f"{label} {conf:.2f}",
                (x1, y1 - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.7,
                (0, 255, 0),
                2
            )

    # 프레임 누적 기록
    if current_label in ["apple", "orange"]:
        history.append(current_label)

    # =========================
    # 최종 판단 (프레임 누적)
    # =========================
    final_result = "몰라"
    if history.count("apple") >= CONFIRM_COUNT:
        final_result = "사과"
    elif history.count("orange") >= CONFIRM_COUNT:
        final_result = "오렌지"

    # =========================
    # 결과 한글 출력
    # =========================
    frame = put_korean_text(
        frame,
        f"결과: {final_result}  (conf={best_conf:.2f})",
        (20, 20),
        color=(255, 0, 0)
    )

    cv2.imshow("Fruit AI (CUDA)", frame)

    # ESC 종료
    if cv2.waitKey(1) & 0xFF == 27:
        break

# =========================
# 종료 처리
# =========================
cap.release()
cv2.destroyAllWindows()

 

 

2. 결과