1. 개발 환경 구축

YOLOv5 기반 객체 탐지 모델을 직접 학습시키기 위한 환경을 마련해보자. 먼저 Ultralytics에서 제공하는 소스 코드를 로컬 머신으로 가져온다. 터미에서 다음 명령을 실행하면 된다.

git clone https://github.com/ultralytics/yolov5.git

Git 사용이 어렵다면 GitHub 페이지에서 ZIP 파일을 직접 내려받아 원하는 경로에 압축 해제하면 된다. 프로젝트 폴더명은 yolov5로 유지하는 것이 이후 단계에서 혼란을 줄인다.

가상 환경 설정 후 PyTorch를 설치한다. NVIDIA GPU와 CUDA가 구축된 환경이라면 GPU 가속 버전을, 그렇지 않다면 CPU 전용 버전을 선택한다.

# CUDA 11.8 환경
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# CPU 전용
pip install torch torchvision torchaudio

핵심 프레임워크 설치가 끝나면 프로젝트 의존성 패키지를 일괄 설치한다.

cd yolov5
pip install -r requirements.txt

2. 데이터셋 아키텍처 설계

YOLOv5는 특정 디렉터리 규칙을 따르는 데이터셋 구조를 요구한다. 다음과 같이 폴더를 구성한다.

dataset/
├── images/
│   ├── train/
│   ├── val/
│   └── test/
└── labels/
    ├── train/
    ├── val/
    └── test/

images 폴더에 원본 이미지를, labels 폴더에 어노테이션 정보를 배치한다. 두 폴더 내의 파일명은 서로 대응되어야 하며, 이미지 확장자가 .jpg라면 라벨 파일은 동일한 이름의 .txt로 생성한다.

어노테이션 파일은 클래스 인덱스와 중심 좌표, 너비, 높이를 공백으로 구분하여 기록한다. 모든 수치는 이미지 크기 대비 정규화된 0~1 범위의 실수여야 한다.

0 0.5 0.5 0.3 0.4
2 0.2 0.3 0.1 0.15

위 예시는 한 이미지 내에 클래스 0번 객체 하나와 클래스 2번 객체 하나가 존재함을 의미한다.

3. YAML 구성 파일 작성

데이터셋 경로와 클래스 정보를 담은 YAML 파일을 프로젝트 내에 생성한다. 파일명은 임의로 지정 가능하며, 여기서는 custom_data.yaml로 한다.

path: ../dataset
train: images/train
val: images/val
test: images/test

nc: 3
names: ['defect_a', 'defect_b', 'normal']

nc는 탐지 대상 클래스의 개수, names는 각 클래스의 식별자 목록이다. 순서는 어노테이션 파일의 클래스 인덱스와 정확히 일치해야 한다.

4. 모델 학습 실행

준비가 완료되면 train.py 스크립트를 호출하여 학습을 개시한다. 사전 학습된 가중치를 활용하면 수렴 속도와 최종 정확도 측면에서 유리하다.

python train.py --img 640 --batch 16 --epochs 100 --data custom_data.yaml --weights yolov5s.pt

주요 인자의 역할은 다음과 같다.

• --img: 입력 이미지의 해상도. 640은 640×640 픽셀을 의미하며, 값이 클수록 메모리 소모 증가
• --batch: 한 번의 반복(iteration)에서 처리하는 이미지 묶음 크기. GPU 메모리에 따라 조정
• --epochs: 전체 데이터셋을 반복 학습할 횟수
• --data: 앞서 작성한 YAML 설정 파일 경로
• --weights: 전이 학습에 사용할 기반 가중치 파일

학습 과정에서 runs/train/exp 디렉터리에 로그와 체크포인트가 자동 저장된다. results.png를 통해 mAP, 손실값 등의 추이를 시각적으로 확인할 수 있다.

5. 추론 파이프라인 구현

학습이 완료된 모델로 실제 이지에서 객체를 탐지해보자. 가장 최근 학습 결과의 가중치 파일 경로를 지정한다.

python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.25 --source inference_images/

--conf는 신뢰도 임계값으로, 이 값 이상의 확률로 탐지된 결과만 최종 출력에 포함된다. 필요에 따라 --source에 단일 이미지, 폴더, 또는 영상 파일 경로를 전달할 수 있다.

탐지 결과는 runs/detect/exp 폴더에 바운딩 박스가 그려진 이미지 형태로 저장되며, 터미널에는 각 객체의 클래스, 신뢰도, 좌표 정보가 출력된다.