251003 - Project (9)

시작하며

오늘부터 연휴가 시작됐지만, 연휴 동안 계속 프로젝트 작업을 할 예정이다. 오늘은 YOLO를 학습시키기 위한 폴더 구조 변환을 하였고, yaml 파일을 만들어서 기본적인 정보들을 저장시키는 작업을 하였다.

YOLO 폴더 변환

import os
import shutil
from pathlib import Path

# 클래스 매핑 (알파벳 순서)
class_mapping = {
    'can': 0,
    'glass': 1,
    'paper': 2,
    'plastic': 3,
    'styrofoam': 4,
    'vinyl': 5
}

# YOLO형식으로 폴더 구조 변환
def prepare_yolo_dataset(source_dir="datasets", output_dir="yolo_dataset"):
    # 긱 폴더를 순회
    for folder in ["train", "valid", "test"]:
        # 이미지 / 라벨 폴더 생성
        img_path = os.path.join(source_dir, output_dir, "images", folder)
        label_path = os.path.join(source_dir, output_dir, "labels", folder)
        os.makedirs(img_path, exist_ok=True)
        os.makedirs(label_path, exist_ok=True)

        # 클래스별 이미지 처리
        for class_name in class_mapping.keys():
            # 각 클래스 이미지 폴더에 접근
            source_class_dir = os.path.join(source_dir, "resnet", folder, class_name)
            # 이미지 파일 찾기
            for img_file in Path(source_class_dir).glob("*.jpg"):
                # 새 파일명 생성
                new_name = f"{class_name}_{img_file.name}"
                # 복사할 전체 경로 지정
                dst_img = os.path.join(img_path, new_name)
                # 이미지 복사
                shutil.copy(img_file, dst_img)
                # 라벨 파일명 생성
                label_name = new_name.replace(".jpg", ".txt")
                # 라벨 파일 전체 경로 지정
                dst_label = os.path.join(label_path, label_name)
                # 라벨 내용 작성
                class_id = class_mapping[class_name]
                with open(dst_label, "w") as f:
                    f.write(f"{class_id} 0.5 0.5 1.0 1.0\n")

yaml

# 재활용품 6종 분류 데이터셋 설정

# 데이터셋 경로
path: datasets/yolo_dataset
train: images/train
val: images/valid
test: images/test

# 클래스 정의
nc: 6  # number of classes
names: ["can", "glass", "paper", "plastic", "styrofoam", "vinyl"]

마치며

내일은 부족한 데이터를 추가적으로 수집하고 YOLO 학습까지 진행해보는 것을 목표로 하고 있다.