실습 3 — 학습한 모델로 부품 찾기

import os
from pathlib import Path
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

weights = Path("out_ch18_runs/parts/weights/best.pt")
if not weights.exists():
    raise SystemExit("best.pt 가 없습니다. 18차시 실습 2를 먼저 실행하세요")
model = YOLO("out_ch18_runs/parts/weights/best.pt")
print("이 모델의 클래스:", model.names)        # 결과: bolt, nut, washer 셋만 남는다

root = Path("out_ch17_yolo")
val_paths = sorted((root / "images" / "val").glob("*.png"))
img = cv2.imread(str(val_paths[0]))
r = model.predict(img, imgsz=320, conf=0.25, device="cpu", verbose=False)[0]
print(val_paths[0].name, "에서 찾은 물체:")
for box, p, c in zip(r.boxes.xyxy.cpu().numpy(), r.boxes.conf.cpu().numpy(), r.boxes.cls.cpu().numpy().astype(int)):
    print(f"  {r.names[int(c)]:6} 신뢰도 {float(p):.2f}  상자 {[int(v) for v in box]}")

def read_labels(label_path, w, h):
    """YOLO 라벨을 읽어 [(클래스번호, (x1, y1, x2, y2)), ...] 픽셀 상자로 돌려준다(17차시의 되돌리는 식)."""
    boxes = []
    for line in label_path.read_text(encoding="utf-8").split("\n"):
        if not line.strip():
            continue
        c, xc, yc, bw, bh = line.split()
        xc, yc, bw, bh = float(xc), float(yc), float(bw), float(bh)
        boxes.append((int(c), ((xc - bw / 2) * w, (yc - bh / 2) * h, (xc + bw / 2) * w, (yc + bh / 2) * h)))
    return boxes

def detect(model, img, conf=0.25):
    """추론해서 [(클래스번호, 신뢰도, (x1, y1, x2, y2)), ...] 로 돌려준다."""
    r = model.predict(img, imgsz=320, conf=conf, device="cpu", verbose=False)[0]
    return [(int(c), float(p), tuple(float(v) for v in b))
            for b, p, c in zip(r.boxes.xyxy.cpu().numpy(), r.boxes.conf.cpu().numpy(),
                               r.boxes.cls.cpu().numpy().astype(int))]

def draw(img, gts, preds):
    """정답은 초록 얇은 선, 찾은 상자는 빨강 굵은 선과 이름·신뢰도로 그린 사본을 돌려준다."""
    out = img.copy()
    for c, (x1, y1, x2, y2) in gts:
        cv2.rectangle(out, (round(x1), round(y1)), (round(x2), round(y2)), (0, 255, 0), 1)
    for c, p, (x1, y1, x2, y2) in preds:
        cv2.rectangle(out, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
        cv2.putText(out, f"{model.names[c]} {p:.2f}", (int(x1), max(int(y1) - 4, 12)),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1)
    return out

fig, axes = plt.subplots(2, 3, figsize=(11, 5.6), layout="constrained")
for ax, p in zip(axes.ravel(), val_paths[:6]):
    img = cv2.imread(str(p))
    h, w = img.shape[:2]
    gts = read_labels(root / "labels" / "val" / (p.stem + ".txt"), w, h)
    preds = detect(model, img)
    ax.imshow(cv2.cvtColor(draw(img, gts, preds), cv2.COLOR_BGR2RGB))
    ax.set_title(f"{p.name}  정답 {len(gts)} · 찾음 {len(preds)}", fontsize=11)
    ax.axis("off")
fig.suptitle("초록: 정답 상자, 빨강: 학습한 모델이 찾은 상자 (conf 0.25)")
plt.show()

import pandas as pd

rows = []
for p in val_paths:
    img = cv2.imread(str(p))
    h, w = img.shape[:2]
    gts = read_labels(root / "labels" / "val" / (p.stem + ".txt"), w, h)
    # 0.25 로 한 번만 추론하고, 0.5 는 그 결과에서 신뢰도로 다시 거른다(결과가 같고 빠르다)
    preds = detect(model, img, conf=0.25)
    rows.append({"그림": p.name, "정답": len(gts),
                 "conf 0.25": len(preds),
                 "conf 0.5": sum(1 for _, s, _ in preds if s >= 0.5)})
table = pd.DataFrame(rows)
print(table.to_string(index=False))
print("합계:", table[["정답", "conf 0.25", "conf 0.5"]].sum().to_dict())