프로젝트 1 — 부품 검사 수업 데모 만들기

import os
import hashlib
import cv2
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

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

folder = "ch20_demo"
names = sorted(n for n in os.listdir(folder) if n.endswith(".png"))

def md5_of(path):
    """파일의 바이트로 MD5 지문을 만든다(16차시). 바이트가 같으면 지문도 같다."""
    with open(path, "rb") as f:
        return hashlib.md5(f.read()).hexdigest()

kept, removed, seen = {}, [], {}
for name in names:
    path = os.path.join(folder, name)
    img = cv2.imread(path)
    if img is None:
        removed.append({"파일": name, "이유": "읽기 실패"})
        continue
    digest = md5_of(path)
    if digest in seen:                         # 앞에서 본 지문이면 같은 사진
        removed.append({"파일": name, "이유": f"{seen[digest]} 와 같은 사진"})
        continue
    seen[digest] = name
    kept[name] = img

print("남은 그림:", len(kept), "장")
print(pd.DataFrame(removed).to_string(index=False))

def dust_ratio(img):
    """흑백 그림과 Median 3 을 건 그림의 차이가 60 을 넘는 칸의 비율. 소금·후추 잡티가 많을수록 크다."""
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return float((cv2.absdiff(gray, cv2.medianBlur(gray, 3)) > 60).mean())

fixed, log = {}, []
for name, img in kept.items():
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    mean_before, dust = float(gray.mean()), dust_ratio(img)
    done = []
    if mean_before < 45:                                          # 어두우면 0~255 로 늘린다(5차시)
        img = cv2.normalize(img, None, 0, 255, cv2.NORM_MINMAX)
        done.append("정규화")
    if dust > 0.002:                                              # 잡티가 많으면 Median(7차시)
        img = cv2.medianBlur(img, 3)
        done.append("Median 3")
    fixed[name] = img
    log.append({"파일": name, "밝기 전": round(mean_before, 1),
                "밝기 후": round(float(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY).mean()), 1),
                "잡티 비율": round(dust, 4), "손질": " + ".join(done) or "그대로"})
print(pd.DataFrame(log).to_string(index=False))

os.environ["YOLO_VERBOSE"] = "False"         # 불러오기 전에 정해야 긴 진행 글이 줄어든다
from ultralytics import YOLO

model_path = "out_ch18_runs/parts/weights/best.pt"
if not os.path.exists(model_path):
    raise SystemExit("모델이 없습니다. 18차시 실습 2를 먼저 실행하세요")
model = YOLO("out_ch18_runs/parts/weights/best.pt")

rows = []
for name, img in fixed.items():
    # 고친 그림(BGR 배열)을 그대로 넣는다. imgsz 는 학습 때와 같은 320
    r = model.predict(img, imgsz=320, conf=0.25, device="cpu", verbose=False)[0]
    boxes = r.boxes.xyxy.cpu().numpy()
    confs = r.boxes.conf.cpu().numpy()
    classes = r.boxes.cls.cpu().numpy().astype(int)
    for (x1, y1, x2, y2), c, k in zip(boxes, confs, classes):
        rows.append({"파일": name, "클래스": r.names[int(k)], "신뢰도": round(float(c), 3),
                     "x1": int(round(float(x1))), "y1": int(round(float(y1))),
                     "x2": int(round(float(x2))), "y2": int(round(float(y2)))})
det = pd.DataFrame(rows)
print("탐지한 상자:", len(det), "개")
print(det["클래스"].value_counts().to_dict())
print(det.head(6).to_string(index=False))

def measure(img, x1, y1, x2, y2, pad=4):
    """상자 안(테두리 pad 칸 여유)을 Otsu 로 이진화해 가장 큰 바깥 윤곽선의 넓이와 원형도를 돌려준다."""
    h, w = img.shape[:2]
    x1, y1 = max(x1 - pad, 0), max(y1 - pad, 0)
    x2, y2 = min(x2 + pad, w), min(y2 + pad, h)
    roi = cv2.cvtColor(img[y1:y2, x1:x2], cv2.COLOR_BGR2GRAY)       # ROI: 상자 안만 잘라 낸다(13차시)
    _, binary = cv2.threshold(roi, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return 0, 0.0
    c = max(contours, key=cv2.contourArea)                          # 가장 큰 덩어리가 부품
    area, perimeter = cv2.contourArea(c), cv2.arcLength(c, True)
    circularity = 4 * np.pi * area / perimeter ** 2 if perimeter else 0.0
    return int(round(area)), round(float(circularity), 3)

measured = [measure(fixed[r.파일], r.x1, r.y1, r.x2, r.y2) for r in det.itertuples()]
det["넓이"] = [m[0] for m in measured]
det["원형도"] = [m[1] for m in measured]
# 클래스마다 원형도 평균: 와셔·너트는 크고 볼트는 작을 것이다
print(det.groupby("클래스")[["넓이", "원형도"]].mean().round(3))

answer = pd.read_csv("ch20_demo/answer.csv")

def iou(a, b):
    """두 상자 (x1, y1, x2, y2) 의 IoU: 겹친 넓이 ÷ 합친 넓이."""
    ix = max(0, min(a[2], b[2]) - max(a[0], b[0]))
    iy = max(0, min(a[3], b[3]) - max(a[1], b[1]))
    inter = ix * iy
    union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / union if union else 0.0

det["판정"], det["IoU"], det["정답 넓이"] = "헛것", 0.0, np.nan
missed = 0
for t in answer.itertuples():
    gt = (t.x1, t.y1, t.x2, t.y2)
    # 같은 그림, 같은 클래스, 아직 짝이 없는 탐지 상자 가운데 IoU 가 가장 큰 것
    cand = det[(det["파일"] == t.filename) & (det["클래스"] == t._2) & (det["판정"] == "헛것")]
    best_i, best_v = None, 0.0
    for i, r in cand.iterrows():
        v = iou(gt, (r.x1, r.y1, r.x2, r.y2))
        if v > best_v:
            best_i, best_v = i, v
    if best_i is not None and best_v >= 0.5:
        det.loc[best_i, ["판정", "IoU", "정답 넓이"]] = ["맞힘", round(best_v, 3), t.area]
    else:
        missed += 1

tp = int((det["판정"] == "맞힘").sum())
precision, recall = tp / len(det), tp / len(answer)
hit = det[det["판정"] == "맞힘"]
err = ((hit["넓이"] - hit["정답 넓이"]) / hit["정답 넓이"] * 100).round(1)
det.loc[hit.index, "넓이 오차(%)"] = err
print(f"정답 {len(answer)}개, 탐지 {len(det)}개, 맞힘 {tp}, 헛것 {len(det) - tp}, 놓침 {missed}")
print(f"정밀도 {precision:.3f}, 재현율 {recall:.3f}")
print("클래스별 넓이 오차 평균(%):", det.groupby("클래스")["넓이 오차(%)"].mean().round(1).to_dict())
det.to_csv("out_ch20_final.csv", index=False, encoding="utf-8-sig")

fig = plt.figure(figsize=(13, 8))
grid = fig.add_gridspec(3, 4, height_ratios=[1, 1, 1.1])
for k, (name, img) in enumerate(fixed.items()):
    view = img.copy()
    for r in det[det["파일"] == name].itertuples():
        color = (0, 200, 0) if r.판정 == "맞힘" else (0, 0, 255)    # BGR: 맞힘 초록, 헛것 빨강
        cv2.rectangle(view, (r.x1, r.y1), (r.x2, r.y2), color, 2)
        # 이름은 상자 위에 쓰고, 상자가 그림 위쪽 끝에 붙어 있으면 상자 아래에 쓴다
        text_y = r.y1 - 4 if r.y1 >= 14 else r.y2 + 13
        cv2.putText(view, r.클래스, (r.x1, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, color, 1)
    ax = fig.add_subplot(grid[k // 4, k % 4])
    ax.imshow(cv2.cvtColor(view, cv2.COLOR_BGR2RGB))
    ax.set_title(name)
    ax.axis("off")

# 아래 줄: 클래스별 정답 개수와 탐지 개수
bar = fig.add_subplot(grid[2, :])
classes = ["bolt", "nut", "washer"]
gt_n = [int((answer["class"] == c).sum()) for c in classes]
det_n = [int((det["클래스"] == c).sum()) for c in classes]
hit_n = [int(((det["클래스"] == c) & (det["판정"] == "맞힘")).sum()) for c in classes]
x = np.arange(len(classes))
bar.bar(x - 0.27, gt_n, width=0.27, label="정답")
bar.bar(x, det_n, width=0.27, label="탐지")
bar.bar(x + 0.27, hit_n, width=0.27, label="맞힘")
bar.set_xticks(x, labels=classes)
bar.set_ylim(0, max(gt_n + det_n) + 2)
bar.legend(loc="upper center", ncol=3)
fig.suptitle(f"부품 검사 데모: 정밀도 {precision:.3f}, 재현율 {recall:.3f}", fontsize=15)
fig.tight_layout()
fig.savefig("out_ch20_final.png", dpi=100)
plt.show()

sheet_name = "part_08.png"
problem = cv2.imread(os.path.join(folder, sheet_name))
key = answer[answer["filename"] == sheet_name].reset_index(drop=True)

view = problem.copy()
for k, t in key.iterrows():
    # 부품 위가 아니라 상자의 왼쪽 위 바깥에 번호를 둔다(부품을 가리지 않게)
    x, y = max(t.x1 - 14, 2), max(t.y1 - 2, 14)
    cv2.putText(view, str(k + 1), (x, y), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)

fig, (left, right) = plt.subplots(1, 2, figsize=(11, 4.2), width_ratios=[1.2, 1])
left.imshow(cv2.cvtColor(view, cv2.COLOR_BGR2RGB))
left.set_title("문제: 번호마다 부품 이름을 적고 넓이(칸)를 코드로 재어 보세요")
left.axis("off")
right.axis("off")
cells = [[str(k + 1), "", ""] for k in range(len(key))]
table = right.table(cellText=cells, colLabels=["번호", "부품 이름", "넓이(칸)"], loc="center", cellLoc="center")
table.scale(1, 2.4)                       # 칸 높이를 키워 손으로 쓸 자리를 만든다
table.set_fontsize(13)
fig.suptitle("부품 검사 학습지 (가상 그림)", fontsize=15)
fig.tight_layout()
fig.savefig("out_ch20_worksheet.png", dpi=100)

print("[강사용 정답]")
print(key.assign(번호=range(1, len(key) + 1))[["번호", "class", "area"]].to_string(index=False))
plt.show()