실습 5 — YOLO 데이터셋 만들기

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

CLASSES = ["bolt", "nut", "washer"]                    # 클래스 번호: bolt 0, nut 1, washer 2
src = Path("ch17_raw")
ann = pd.read_csv(src / "annotations.csv")
sizes = {p.name: cv2.imread(str(p)).shape[:2] for p in sorted(src.glob("*.png"))}

ann = ann[ann["filename"].isin(sizes)]                 # 그림 없는 주석: 지운다
ann = ann.drop_duplicates()                            # 중복 줄: 하나만 남긴다
# 모르는 클래스가 있는 그림은 통째로 빼서 사람에게 돌려보낸다
recheck = set(ann.loc[~ann["class"].isin(CLASSES), "filename"])
recheck |= set(sizes) - set(ann["filename"])           # 주석 없는 그림도 뺀다

rows, count = [], {"뒤집기": 0, "자르기": 0, "크기 0 지움": 0}
for r in ann.to_dict("records"):
    if r["filename"] in recheck:
        continue
    h, w = sizes[r["filename"]]
    x1, x2 = sorted((r["x1"], r["x2"]))                 # 뒤집힌 상자: 작은 값이 앞으로
    y1, y2 = sorted((r["y1"], r["y2"]))
    count["뒤집기"] += (x1, x2) != (r["x1"], r["x2"]) or (y1, y2) != (r["y1"], r["y2"])
    cx1, cy1, cx2, cy2 = max(x1, 0), max(y1, 0), min(x2, w), min(y2, h)   # 범위 밖: 테두리에 맞춰 자른다
    count["자르기"] += (cx1, cy1, cx2, cy2) != (x1, y1, x2, y2)
    if cx2 - cx1 < 1 or cy2 - cy1 < 1:                 # 크기 0 상자: 지운다
        count["크기 0 지움"] += 1
        continue
    rows.append({"filename": r["filename"], "class": r["class"], "x1": cx1, "y1": cy1, "x2": cx2, "y2": cy2})

clean = pd.DataFrame(rows)
print("고친 것:", count)
print("다시 주석할 그림:", sorted(recheck))
print("남은 그림:", clean["filename"].nunique(), "장 · 남은 물체:", len(clean), "개")

names = sorted(clean["filename"].unique())
rng = np.random.default_rng(1717)                      # 씨앗을 고정해 다시 돌려도 같게 나눈다
order = rng.permutation(len(names))                    # permutation: 0..n-1 을 무작위 차례로
n_train = round(len(names) * 0.8)
split = {names[i]: ("train" if k < n_train else "val") for k, i in enumerate(order)}
clean["split"] = clean["filename"].map(split)

print("그림 수:", pd.Series(split).value_counts().to_dict())
print("물체 수:", clean["split"].value_counts().to_dict())
# crosstab: 두 열을 가로세로로 놓고 개수를 센 표
print(pd.crosstab(clean["class"], clean["split"]))

import shutil

root = Path("out_ch17_yolo")
if root.exists():
    shutil.rmtree(root)                                # 다시 돌릴 때 지난 결과가 섞이지 않게 지우고 새로 만든다
for kind in ("images", "labels"):
    for part in ("train", "val"):
        (root / kind / part).mkdir(parents=True, exist_ok=True)

for name in names:
    part = split[name]
    # shutil.copy2: 파일을 그대로 복사한다. 원본은 ch17_raw 에 그대로 남는다
    shutil.copy2(src / name, root / "images" / part / name)
    h, w = sizes[name]
    lines = []
    for r in clean[clean["filename"] == name].itertuples():
        xc, yc = (r.x1 + r.x2) / 2 / w, (r.y1 + r.y2) / 2 / h
        bw, bh = (r.x2 - r.x1) / w, (r.y2 - r.y1) / h
        lines.append(f"{CLASSES.index(r[2])} {xc:.6f} {yc:.6f} {bw:.6f} {bh:.6f}")   # r[2] = class 열
    # 그림과 같은 이름(stem)의 .txt 에 한 줄에 물체 하나씩 쓴다
    (root / "labels" / part / (Path(name).stem + ".txt")).write_text("\n".join(lines) + "\n", encoding="utf-8")

# data.yaml: path 는 data 폴더에서 실행할 때의 상대 경로, train/val 은 path 아래의 그림 폴더
yaml = ["path: out_ch17_yolo", "train: images/train", "val: images/val", "names:"]
yaml += [f"  {k}: {c}" for k, c in enumerate(CLASSES)]
(root / "data.yaml").write_text("\n".join(yaml) + "\n", encoding="utf-8")

for kind in ("images", "labels"):
    for part in ("train", "val"):
        print(f"{kind}/{part}: {len(list((root / kind / part).iterdir()))}개")
print("--- data.yaml")
print((root / "data.yaml").read_text(encoding="utf-8"), end="")
print("--- labels/train 의 첫 파일:", sorted((root / "labels" / "train").iterdir())[0].name)
print(sorted((root / "labels" / "train").iterdir())[0].read_text(encoding="utf-8"), end="")

def draw_labels(img_path, label_path):
    """YOLO 라벨을 읽어 그림에 상자와 이름을 그린 RGB 그림을 돌려준다."""
    img = cv2.imread(str(img_path))
    h, w = img.shape[:2]
    colors = [(0, 0, 255), (0, 200, 255), (255, 200, 0)]   # BGR: bolt 빨강, nut 노랑, washer 하늘
    for line in label_path.read_text(encoding="utf-8").split("\n"):
        if not line.strip():
            continue
        cls, xc, yc, bw, bh = line.split()
        cls, xc, yc, bw, bh = int(cls), float(xc), float(yc), float(bw), float(bh)
        # 비율 → 픽셀로 되돌린다
        x1, y1 = round((xc - bw / 2) * w), round((yc - bh / 2) * h)
        x2, y2 = round((xc + bw / 2) * w), round((yc + bh / 2) * h)
        cv2.rectangle(img, (x1, y1), (x2, y2), colors[cls], 2)
        cv2.putText(img, CLASSES[cls], (x1, max(y1 - 4, 12)), cv2.FONT_HERSHEY_SIMPLEX, 0.45, colors[cls], 1)
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

val_images = sorted((root / "images" / "val").glob("*.png"))[:6]
fig, axes = plt.subplots(2, 3, figsize=(11, 5.6), layout="constrained")
for ax, p in zip(axes.ravel(), val_images):
    ax.imshow(draw_labels(p, root / "labels" / "val" / (p.stem + ".txt")))
    ax.set_title(p.name, fontsize=11)
    ax.axis("off")
fig.suptitle("검증용 그림 여섯 장에 YOLO 라벨을 되돌려 그린 결과")
plt.show()