YOLOv5

YOLOv5 뢄석(2) - Train Custom Data

dhpark 2022. 2. 3. 00:24

좜처: https://docs.ultralytics.com/tutorials/train-custom-datasets/

 

Train Custom Data πŸ“Œ - YOLOv5 Documentation

Train Custom Data πŸ“Œ πŸ“š This guide explains how to train your own custom dataset with YOLOv5 πŸš€. Before You Start Clone this repo, download tutorial dataset, and install requirements.txt dependencies, including Python>=3.8 and PyTorch>=1.7. $ git clo

docs.ultralytics.com

μœ„μ˜ μ‚¬μ΄νŠΈλ₯Ό μ°Έκ³ ν•˜μ—¬ YOLOv5λ₯Ό μ»€μŠ€ν…€ λ°μ΄ν„°λ‘œ ν•™μŠ΅ν•˜λŠ” 방법을 κ°„λž΅ν•˜κ²Œ μ•Œμ•„λ³΄κ² μŠ΅λ‹ˆλ‹€.


Train Custom Data

1. Create dataset.yaml

COCO128은 COCO train2017의 μΌλΆ€λΆ„μœΌλ‘œ Training, Validation λ°μ΄ν„°λ‘œ λ‚˜λ‰˜μ–΄μ Έ 있으며, ν•™μŠ΅ 과정이 잘 λ™μž‘ν•˜λŠ”μ§€ ν…ŒμŠ€νŠΈν•˜κΈ°μ— μ ν•©ν•©λ‹ˆλ‹€(ν•™μŠ΅ λͺ©μ μ΄ μ•„λ‹™λ‹ˆλ‹€)

 

data/coco128.yaml 파일

# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../coco128/images/train2017/
val: ../coco128/images/train2017/

# number of classes
nc: 80

# class names
names: ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 
        'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 
        'teddy bear', 'hair drier', 'toothbrush']

 

2. Create Labels

라벨을 YOLO format으둜 μ€€λΉ„ ν•©λ‹ˆλ‹€. λ§Œμ•½ 이미지에 물체(Object)κ°€ μ‘΄μž¬ν•˜μ§€ μ•ŠλŠ”λ‹€λ©΄, *.txt νŒŒμΌμ€ μƒμ„±λ˜μ§€ μ•ŠμŠ΅λ‹ˆλ‹€.

 

νŠΉμ§•

- ν•œμ€„μ— ν•˜λ‚˜μ˜ 이미지λ₯Ό ν¬ν•¨ν•©λ‹ˆλ‹€.

- 각각의 행은 (class, x_center, y_center, width, height) ν˜•μ‹μž…λ‹ˆλ‹€.

- μ’Œν‘œ 값은 0 ~ 1 사이 κ°’μœΌλ‘œ μ •κ·œν™” λ˜μ–΄ μžˆμŠ΅λ‹ˆλ‹€.

 

txt 파일 κ°’ μ˜ˆμ‹œ
yolo format 이미지 μ˜ˆμ‹œ

 

3. Organize Directories

μ•„λž˜μ™€ 같은 폴더 ꡬ쑰둜 이미지와 라벨을 λ°°μΉ˜ν•©λ‹ˆλ‹€. YOLOv5λŠ” 이미지에 λŒ€μ‘ν•˜λŠ” 라벨을 μžλ™μ μœΌλ‘œ μ°ΎμŠ΅λ‹ˆλ‹€(μ•„λ§ˆλ„ ν™•μž₯μžκ°€ λ‹€λ₯Έ 동일 μ΄λ¦„μœΌλ‘œ μ°ΎλŠ” 것 κ°™μŠ΅λ‹ˆλ‹€. e.g. 'im1.jpg <-> im1.txt' 쌍 맀칭됨)

|- yolov5

|- {데이터셋 이름}

    |- images/im0.jpg, ... imN.jpg

    |- labels/im0.txt, ...imN.txt

 

4. Select a Model

- μ‚¬μ΄νŠΈμ—μ„œ λͺ©μ μ— λ§žλŠ” λͺ¨λΈμ„ μ„ νƒν•©λ‹ˆλ‹€. 즉, μ œμ–΄κΈ°κ°€ μˆ˜μš©κ°€λŠ₯ν•œ λ©”λͺ¨λ¦¬ μ‚¬μ΄μ¦ˆλ₯Ό κ³ λ €ν•΄μ„œ λͺ¨λΈμ„ μ„ νƒν•©λ‹ˆλ‹€.

 

5. Train

- 미리 ν•™μŠ΅λœ κ°€μ€‘μΉ˜λ₯Ό λ‹€μš΄ν•˜κ³ , λ°°μΉ˜μ‚¬μ΄μ¦ˆ 등을 μ„€μ •ν•˜μ—¬ ν•™μŠ΅ν•©λ‹ˆλ‹€.

- Pretrained κ°€μ€‘μΉ˜ μ‚¬μš©ν•˜μ§€ μ•Šκ³ , 랜덀으둜 μ΄ˆκΈ°ν™”λœ κ°€μ€‘μΉ˜ + yolov5X.yaml 을 μ‚¬μš©ν•˜λŠ” 것은 μΆ”μ²œν•˜μ§€ μ•ŠμŠ΅λ‹ˆλ‹€.

(e.g. --weights '' --cfg yolov5s.yaml)

# Train YOLOv5s on COCO128 for 5 epochs
$ python train.py --img 640 --batch 16 --epochs 5 --data coco128.yaml --weights yolov5s.pt

- runs/train/expX 폴더가 μƒμ„±λ˜λ©΄μ„œ ν•™μŠ΅ 과정이 μ €μž₯λ©λ‹ˆλ‹€.

 

6. Visualize

Weights & Biases λ‘œκΉ… 도ꡬλ₯Ό μ œκ³΅ν•˜μ—¬ ν΄λΌμš°λ“œ μƒμ—μ„œ μ‹€μ‹œκ°„μœΌλ‘œ ν•™μŠ΅ ν˜„ν™©μ„ μ‚΄νŽ΄ λ³Ό 수 μžˆμŠ΅λ‹ˆλ‹€.

pip install wandb # μ„€μΉ˜

λ¬Όλ‘  runs/train/expX 폴더에 λ‘œμ»¬λ‘œλ„ μ €μž₯λ˜μ–΄ 확인이 κ°€λŠ₯ν•©λ‹ˆλ‹€. ν…μ„œλ³΄λ“œ ν˜•μ‹ (+ results.txt 파일)둜 μ €μž₯λ˜μ–΄ ν…μ„œλ³΄λ“œλ‘œ μ—΄μ–΄ λ³Ό 수 μžˆμŠ΅λ‹ˆλ‹€.

from utils.plots import plot_results 
plot_results(save_dir='runs/train/exp')  # plot results.txt as results.png

 

AWS, GCP, Docker Image 둜 ν•™μŠ΅ ν™˜κ²½μ„ μ œκ³΅ν•©λ‹ˆλ‹€.