[이미지 학습AI] 이미지를 학습으로 annotation을 해보자.[2]

이전 게시글에서 jpg , xml 형식의 Pascal VOC Dataset의 이미지 데이터들을 다운로드 받았다.

우리가 사용할 모델은 yolov5로 링크는 다음과 같다.

https://github.com/ultralytics/yolov5 

물론 이전 게시글에서 클론 했지만 혹시하고 다시 적어본다.

기본적으로 yolov5는 yolo dataset을 사용하고있다.

VOC도 변환할 수 있으니 너무 걱정은 안해도 된다.

data/VOC.yaml 파일에서 아래에 스크립트가 보인다.

import xml.etree.ElementTree as ET

from tqdm import tqdm
from utils.general import download, Path


def convert_label(path, lb_path, year, image_id):
  def convert_box(size, box):
      dw, dh = 1. / size[0], 1. / size[1]
      x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
      return x * dw, y * dh, w * dw, h * dh

  in_file = open(path / f'VOC{year}/Annotations/{image_id}.xml')
  out_file = open(lb_path, 'w')
  tree = ET.parse(in_file)
  root = tree.getroot()
  size = root.find('size')
  w = int(size.find('width').text)
  h = int(size.find('height').text)

  for obj in root.iter('object'):
      cls = obj.find('name').text
      if cls in yaml['names'] and not int(obj.find('difficult').text) == 1:
          xmlbox = obj.find('bndbox')
          bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
          cls_id = yaml['names'].index(cls)  # class id
          out_file.write(" ".join([str(a) for a in (cls_id, *bb)]) + '\n')

convert_label 함수를 보면, infile = open(...) 이 부분을 우리가 가지고있는 xml 파일을 넣어주면 txt형태의 yolo 데이터셋으로 바꿔준다.

https://github.com/ultralytics/yolov5/issues/3128#issuecomment-984467383

 

yolo로 변환이 완료되면, 다시 한번 data/VOC.yaml 파일을 열어서 아래의 내용을 확인해보자.

path: ../datasets/VOC
train: # train images (relative to 'path')  16551 images
  - images/train2012
  - images/train2007
  - images/val2012
  - images/val2007
val: # val images (relative to 'path')  4952 images
  - images/test2007
test: # test images (optional)
  - images/test2007

데이터 셋이 없을 경우 다운로드 받을 수 있는 VOC 데이터를 기준으로 작성되어있으므로, 주석은 무시해도 좋다.

이 yaml을 복사해서 다른 이름으로 저장하고, 우리의 데이터셋에 맞게 수정하자.

우리가 봐야할 항목은  train, val, test 인데, 위와같이 디렉터리를 구성하면 된다.

디렉터리 위치는 중요하지않고, 구성만 맞으면 된다.

예시로 든 것과는 다르지만 나는 다음과 같이 구성했다.

path: /home/<user>/image_data/
train: images/train
val: images/val

nc: <your classes>
names: ['1stclass','2ndclass', ....]

디렉터리 구성은 아래와 같다. images에 사진파일, labels에 txt파일을 동일한 이름으로 넣어주면 된다.

.
├── images
│   ├── dataset1
│   ├── dataset2
│   └── dataset3
└── labels
    ├── dataset1
    ├── dataset1.cache
    ├── dataset2
    ├── dataset2.cache
    ├── dataset3
    └── dataset3.cache

.cache는 train시에 생성되니 무시해도 좋다.

다음은 cfg 파일이다. models 디렉터리에 들어가보자.

yolov5x ~ yolov5l 이 있다. 각각의 Conv가 다르기때문에 성능과 속도적인 측면에서 퍼포먼스가 차이가 난다.

나는 yolov5s.yaml을 복사해서 사용했다.

사실 바꿔야 할 부분은 nc부분 하나이다. 우리가 학습시킬 데이터셋의 클래스 개수.

backbone, head등을 수정할 수 있을정도라면 이 포스트를 보지 않을테니 배제하도록 하자.

이렇게 수정하고나면? 준비 끝이다.

python train.py --data data/<yourname>.yaml --epoch 500 --batch 2 --cfg models/<yourname>.yaml

끝이다.

아까 복사해서 이름을 바꿨던 data/ yaml, models/ yaml 을 인자로 넣어주고 반복횟수와 배치사이즈를 조절해주고 학습시키면 된다.

기타 다른 옵션들이 있지만. 이는 직접 train의 opt를 확인해보면 된다.