[mmrotate] 나만의 데이터 셋으로 rotated RetinaNet 학습하기

이번 포스팅에서는 mmrotate에서 제공하는 데이터 셋이 아닌 임의의 데이터 셋을 이용하여 모델을 학습하려 합니다.

새로운 데이터 셋으로는 앞선 포스팅에서 공부했던 COCO 데이터를 이용합니다.

 

 

---

1. DOTA 데이터 형태로 변환

먼저, COCO 데이터 셋을 DOTA 데이터 형태로 변환합니다.

COCO는 bbox가 x, y, w, h 형태이므로 이를 x1, y1, x2, y2, x3, y3, x4, y4 형태로 바꾸는 함수를 작성합니다.

여기서 COCO의 x, y는 왼쪽 위(left-top)의 꼭지점 위치이며 이미지의 좌상단을 원점으로 합니다.

DOTA도 x1, y1를 왼쪽 위로 하며 나머지들은 시계 방향으로 꼭지점들을 의미합니다.

추가로 category정보와 difficult정보(1: 어려움, 0: 쉬움)를 기록하며, 텍스트 파일로 저장하기 때문에 string 형태로 리턴하도록 정의했습니다.

def cocoToDota(bbox):
    x = round(bbox[0])
    y = round(bbox[1])
    w = round(bbox[2])
    h = round(bbox[3])
    x1, y1 = x, y
    x2, y2 = x+w, y
    x3, y3 = x+w, y+h
    x4, y4 = x, y+h
    return f'{x1} {y1} {x2} {y2} {x3} {y3} {x4} {y4} person 0\n'

 

 

영상과 라벨정보 저장하는 코드는 아래와 같습니다.

1,000장만 저장하였고 category는 사람(person)만 사용합니다.

myLoader는 앞선 포스팅(https://fictitious.tistory.com/3) 것을 그대로 사용합니다.

for i, (img, bboxes) in enumerate(myLoader):
    if i > 999:
        break
        
    imPil = to_pil_image(img[0])
    imPil.save('COCO/train/images/%04d.png' % i)

    f = open('COCO/train/labelTxt/%04d.txt' % i, 'w')
    for j in range(len(bboxes[0])):
        npBbox = bboxes[0][j].numpy()
        dotaBbox = cocoToDota(npBbox)
        f.write(dotaBbox)

 

 

 

DOTA 형태로 만든 학습데이터를 img_split.py를 이용하여 전처리를 합니다.

해당 과정은 이전 글에서 소개드려서 생략하겠습니다.

 

 

---

2. Config 파일 생성

데이터 셋이 다 준비가 되었다면 config 파일을 생성합니다.

앞서 설명했듯이 class로 사람(person)만 사용하여 'classes' 및 'num_classes' 설정하였습니다.

추가 category가 있다면 그에 맞게 설정하시면 됩니다.

data.train, val, test의 ann_file 및 img_prefix도 데이터 셋 위치에 맞게 설정합니다.

그리고 configs/myConfig.py로 저장하였습니다.

 

_base_ = 'rotated_retinanet/rotated_retinanet_hbb_r50_fpn_1x_dota_oc.py'

dataset_type = 'DOTADataset'
classes = ('person',)
data = dict(
    samples_per_gpu=2,
    workers_per_gpu=2,
    train=dict(
        type=dataset_type,
        classes=classes,
        ann_file='./data/split_1024_dota1_0/trainval/annfiles/',
        img_prefix='./data/split_1024_dota1_0/trainval/images/'),
    val=dict(
        type=dataset_type,
        classes=classes,
        ann_file='./data/split_1024_dota1_0/trainval/annfiles/',
        img_prefix='./data/split_1024_dota1_0/trainval/images/'),
    test=dict(
        type=dataset_type,
        classes=classes,
        ann_file='./data/split_1024_dota1_0/test/images/',
        img_prefix='./data/split_1024_dota1_0/test/images/'))

model = dict(
    bbox_head=dict(
        type='RotatedRetinaHead',
        num_classes=1))

 

 

---

3. 학습 및 테스트

학습 방법은 이전 포스팅과 동일한데 config 파일만 앞서 만든 파일로 변경해주시면 됩니다.

python tools/train.py configs/myConfig.py

 

 

학습을 끝내고 잘 학습되었는지 테스트를 진행합니다.

python tools/test.py configs/myConfig.py work_dirs/myConfig/latest.pth \
  --show-dir myDir/vis

 

출력:

오른쪽 두명의 사람을 각각 인식하지는 못하지만

그래도 학습은 일단 되는 것 같습니다.

 

 

 

감사합니다.