[ROADs] Dataset pre-processing & Training

Modeling Process Architecture

목표: 기존 YOLOv8n에 custom dataset을 이용해서 transfer learning 을 통한 obstacle detection model 완성

 

modeling process

 

Dataset setting

AIhub 내 이미지 데이터 수집

지자체 도로부속시설물 파손 데이터 : 도로정비 대상이 되는 도로부속시설물 파손데이터 16종 (데이터 구축년도 - 2022년 / 구축량 - 422,610장)

 

Data Converting

annotation 파일에서 사용할 정보 추출

AI HUB dataset 에서 다운로드 받은 자료를 살펴보니, 이미지 한 장 당 annotation 정보를 담은 json 파일이 1개씩 존재한다.

따라서 각 json 파일에서 필요한 정보만 추출해서 사용하기로 결정하고 최종으로 필요한 json 파일 구조는 아래와 같다.

 

 

Bounding Box 형식 변경

YOLOv8 모델은 바운딩 박스를 [x_center, y_center, width, height] 형식으로 입력해야 하는데,

우리가 가지고 있는 annotation file 의 bbox format은 [좌상단 x좌표, 좌상단 y좌표, bbox_width, bbox_height]이므로

center point를 잡으면서 data normalizing 하는 과정을 먼저 반영했다.

 

Bounding Box (왼: dataset format / 오: YOLO format)

위 과정을 위해서 사용한 코드는 아래와 같다.

def convert_to_yolov8_format(bbox, image_width, image_height):

    x_center = (bbox[0] + bbox[2] / 2) / image_width
    y_center = (bbox[1] + bbox[3] / 2) / image_height
    width = bbox[2] / image_width
    height = bbox[3] / image_height

    return x_center, y_center, width, height

 

Training using YOLOv8n

YOLO v8n 전이 학습 기반 이미지 객체 바운딩 박스 검출 학습 과정을 수행했다.

모델 성능 개선을 위해 MS COCO dataset 으로 pre-trained model 에 우리 데이터를 넣어서 transfer learning 하는 방식으로 진행

Transfer Learning → COCO Dataset으로 학습된 YOLOv8n + Custom Dataset (학습환경: colab GPU, Tesla T4)

 

  -  train set, validation set 비율은 8:2로 진행

  -  학습에 포함되지 않은 test dataset 약 3,000장으로 성능평가(mAP)를 해서 가장 성능이 높은 모델 채택하기로 함

 

ISSUE (1) : Image resizing 관련

YOLOv8 모델의 경우 정사각형 이미지만 Input으로 받는다는 것을 확인 → 우리가 가진 dataset의 경우는 (1920, 1080) 사이즈의 이미지이기 때문에 정사각형으로 resizing을 고려했다.
하지만, training 과정에서 'rect=True' 로 설정해 주면 직사각형 형태도 넣을 수 있다는 것을 확인했고, 아래와 같은 조건으로 학습을 진행했다.

    !yolo task=detect 
             mode=train 
             model=yolov8n.pt 
             data=/content/drive/MyDrive/FinalProject/data.yaml 
             epochs=20 
             imgsz=[1920,1080] 
             rect=True

 

ISSUE (2) : Resuming Training with the last.pt file

20000장의 이미지로 모델 훈련하는 과정에서 1epoch 당 거의 40분 정도가 소요되었다. (환경: colab T4)

따라서 총 20epoch을 돌렸기 때문에 전체 epoch을 한 번에 훈련시킬 수 없어서 학습을 끊었다가 다시 재개하는 과정이 필요했다.

아래 코드를 이용하면 이전에 학습 완료했던 YOLOv8n 모델의 weight 들을 가지고 이어서 학습을 재개할 수 있었다.

model = YOLO("/content/drive/MyDrive/FinalProject/runs/detect/train/weights/last.pt")
model.train(resume=True)

 

Training Result

1차 학습 : 20,000개의 무작위 데이터로 진행

모든 class 에 대한 mAP는 0.873으로 비교적 높게 나왔으나(*20 epoch), 데이터별로 편차가 심함

  → 이는 전체 train dataset에서 검출된 객체를 나타내는 instances 의 숫자가 너무 차이나서 발생한 결과라고 볼 수 있겠음

 

1차학습 - validation 결과

2차 학습 : instances 개수를 비교적 균등하게 맞춤

모든 class 를 비슷하게 해서(각 class 당 1000개씩은 모으기! 1000개 찍으면 다른 걸로 넘어가는 걸로, minimum 개수를 설정)

    →  Dataset : 총 14,821 장 사용

그 외 조건은 모두 1차 Training과 동일한 조건으로 진행

 

1차학습보다 사용한 이미지 양이 적어서, 전체 카테고리에 대해서 평균을 낸 mAP는 0.01 정도 감소했지만, 1차학습에서 낮게 나온 카테고리들의 mAP 값을 높인 2차학습 모델을 채택!

 

2차학습 - validation 결과

 

 

Model Evaluation

학습에 사용되지 않은 새로운 test dataset 약 3,000장으로 test 했을 때, 1차 모델은 0.675, 2차 모델은 0.802로 약 18.81% mAP가 증가한 것을 확인할 수 있었다.
또한 개별 category를 살펴보아도, 대부분의 category mAP 값이 1차 모델 대비 2차모델에서 상승한 것을 볼 수 있다.

 

1차 2차 모델 category별 mAP 비교

 

[Training Code]

https://github.com/juooo1117/ROADs/blob/main/modeling/model_training_2.ipynb