실습 1 — 사전학습 모델로 첫 추론

import os
# 불러오기 전에 둔다: ultralytics 의 긴 진행 글을 줄인다
os.environ["YOLO_VERBOSE"] = "False"
from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt

plt.rcParams["font.family"] = "Malgun Gothic"
plt.rcParams["axes.unicode_minus"] = False

# 사전학습 가중치를 불러온다(data 폴더의 yolo11n.pt)
model = YOLO("yolo11n.pt")
# names: 클래스 번호 → 이름 사전. COCO 의 80가지가 들어 있다
print("클래스 수:", len(model.names))
print("앞의 다섯:", [model.names[k] for k in range(5)])
# 과일 이름의 번호를 찾아 둔다. lime(라임)은 COCO 에 없다
fruits = {name: k for k, name in model.names.items() if name in ("apple", "orange", "banana", "lime")}
print("과일 번호:", fruits)

img = cv2.imread("fruit_stand.png")
print("그림 크기:", img.shape)

# predict: 추론. imgsz=320 은 긴 쪽을 320 에 맞춰 넣는 크기, conf=0.25 는 신뢰도 기준값
results = model.predict(img,
                        imgsz=320,          # 넣는 크기
                        conf=0.25,          # 이보다 자신 없는 상자는 버린다
                        device="cpu",       # CPU 로 계산
                        verbose=False)      # 진행 글을 찍지 않는다
r = results[0]                              # 그림 한 장을 넣었으니 결과도 하나

# 결과를 NumPy 배열로 꺼낸다. .cpu() 는 계산 장치에서 꺼내 오는 일
xyxy = r.boxes.xyxy.cpu().numpy()
conf = r.boxes.conf.cpu().numpy()
cls = r.boxes.cls.cpu().numpy().astype(int)
print("찾은 물체 수:", len(xyxy))
for box, p, c in zip(xyxy, conf, cls):
    x1, y1, x2, y2 = [int(v) for v in box]
    print(f"{r.names[int(c)]:12} 신뢰도 {float(p):.2f}  상자 ({x1}, {y1}, {x2}, {y2})")

drawn = img.copy()                          # 원본은 그대로 두고 사본에 그린다
for box, p, c in zip(xyxy, conf, cls):
    x1, y1, x2, y2 = [int(v) for v in box]
    cv2.rectangle(drawn, (x1, y1), (x2, y2), (255, 255, 255), 2)     # 흰 상자
    label = f"{r.names[int(c)]} {float(p):.2f}"
    # 글자 뒤에 검은 띠를 깔아 어느 바탕에서도 읽히게 한다. getTextSize: 글자가 차지할 크기
    (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.45, 1)
    ty = max(y1, th + 4)
    cv2.rectangle(drawn, (x1, ty - th - 4), (x1 + tw + 2, ty), (0, 0, 0), -1)   # -1: 속을 채운다
    cv2.putText(drawn, label, (x1 + 1, ty - 3), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 255), 1)

# 클래스별로 몇 개 찾았는지 센다
counts = {}
for c in cls:
    counts[r.names[int(c)]] = counts.get(r.names[int(c)], 0) + 1
print("클래스별 개수:", counts)

fig, axes = plt.subplots(1, 2, figsize=(11, 4.2), layout="constrained")
axes[0].imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
axes[0].set_title("fruit_stand.png 원본")
axes[1].imshow(cv2.cvtColor(drawn, cv2.COLOR_BGR2RGB))
axes[1].set_title(f"yolo11n 이 찾은 물체 {len(xyxy)}개 (conf 0.25 이상)")
for ax in axes:
    ax.axis("off")
plt.show()