YOLO를 활용한 쓰레기 분류 모델
prem.py 파일들 라벨링에 따라 클래스 넘버->문자 or 문자->클래스 넘버
이미지 속성 번역
이미지 속성을 영어로 번역하는 딕셔너리를 설정하여, 데이터의 일관성을 유지합니다.
필요 이미지 속성 리스트
모델 학습에 필요한 이미지 속성 리스트를 정의합니다:
해상도, ISO, 낮/밤, 장소, 물체 클래스, 물체 색상, 경계 박스 좌표 등.
파일 및 폴더 관리
확장자 설정: 처리할 이미지와 JSON 파일의 확장자를 설정합니다.
백업 디렉터리: 존재하지 않을 경우 백업 디렉터리를 생성합니다.
데이터 정리
필요 없는 데이터 및 오류 이미지를 제거하여 데이터를 정리합니다.
라벨링 데이터에서 필요 없는 클래스의 자료를 삭제합니다.
이미지 속성 값 추출
정리된 데이터에서 이미지 속성 값을 추출하여 CSV 파일로 저장합니다. 이 파일은 모델 학습에 필수적인 정보를 포함합니다.
데이터 전처리
이미지 크기를 640x640으로 조정하고 클래스 정보를 기반으로 데이터를 전처리합니다.
최종적으로 필요한 데이터 파일을 지정된 경로로 분배
data = {
'train': os.path.join(data_path, 'train/images'),
'val': os.path.join(data_path, 'val/images'),
'test': os.path.join(data_path, 'test/images'),
'names': ['metal', 'vinyl', 'styrofoam', 'glass', 'paper', 'plastic'],
'nc': 6 # 클래스 수
}
yaml_path = os.path.join(data_path, 'data.yaml')
with open(yaml_path, 'w') as f:
yaml.dump(data, f)
데이터 경로와 클래스 이름을 설정하고, 이를 YAML 파일로 저장.
results = model.train(
model='yolo11n.pt',
data='data/data.yaml',
epochs=10,
patience=5,
batch=16,
seed=42,
optimizer='AdamW',
pretrained=True,
device='cuda',
save=True,
save_period=1,
lr0=0.001
)
지정된 하이퍼파라미터를 사용하여 모델 훈련.
테스트 이미지 기준 95%의 정확도