Ultralytics Yolo 실습(Oxford pet dataset)

Ultralytics Yolo v3패키지를 먼저 설치해 주어야 된다.

https://github.com/ultralytics/yolov3 주소에 가면 필요한 패키지를 다운받을 수 있다.

!git clone https://github.com/ultralytics/yolov3
!cd yolov3;pip install -qr requirements.txt

import torch
from IPython.display import Image, clear_output

필요한 모듈을 불러온다.

https://www.robots.ox.ac.uk/~vgg/data/pets/ 에서 필요한 oxford pet dataset을 다운로드 받는다.

이후 data 폴더를 생성후 data 폴더에 다운로드 받은 압축 파일을 푼다.

 

다운받은 oxford pet dataset은 train과 val dataset으로 분리하여 학습을 진행한다.

먼저 다운받은 oxford pet dataset의 annotations파일을 pandas를 이용하여 살펴본다.

import pandas as pd 

pd.read_csv('/content/data/annotations/trainval.txt', sep=' ', header=None, names=['img_name', 'class_id', 'etc1', 'etc2'])

 

 

 

 

 

 

 

 

이제 전체 image/annotation 파일명을 가지는 리스트 파일명을 입력 받아 메타 파일용 Dataframe 및 train/val용 Dataframe을 생성한다.

import os
import pandas as pd
from sklearn.model_selection import train_test_split

def make_train_valid_df(list_filepath, img_dir, anno_dir, test_size=0.1):
  pet_df = pd.read_csv(list_filepath, sep=' ', header=None, names=['img_name', 'class_id', 'etc1', 'etc2'])
  #class_name은 lambda 함수를 이용하여 image 파일명에서 맨 마지막 '_' 문자열 앞까지에 해당하는 문자까지 
  pet_df['class_name'] = pet_df['img_name'].apply(lambda x:x[:x.rfind('_')])
  
  # image 파일명과 annotation 파일명의 절대경로 컬럼 추가
  pet_df['img_filepath'] = img_dir + pet_df['img_name']+'.jpg'
  pet_df['anno_filepath'] = anno_dir + pet_df['img_name']+'.xml'
  # annotation xml 파일이 없는데, trainval.txt에는 리스트가 있는 경우가 있다. 
  #이들의 경우 pet_df에서 해당 rows를 삭제한다.
  pet_df = remove_no_annos(pet_df)

  # 전체 데이터의 0.1을 검증 데이터로, 나머지는 학습 데이터로 분리. 
  train_df, val_df = train_test_split(pet_df, test_size=test_size, stratify=pet_df['class_id'], random_state=2021)
  return pet_df, train_df, val_df

# annotation xml 파일이 없는데, trainval.txt에는 리스트가 있는 경우에 이들을 dataframe에서 삭제하기 위한 함수.
def remove_no_annos(df):
  remove_rows = []
  for index, row in df.iterrows():
    anno_filepath = row['anno_filepath']
    if not os.path.exists(anno_filepath):
      #해당 DataFrame index를 remove_rows list에 담음. 
      remove_rows.append(index)
  # DataFrame의 index가 담긴 list를 drop()인자로 입력하여 해당 rows를 삭제
  df = df.drop(remove_rows, axis=0, inplace=False)
  return df


pet_df, train_df, val_df = make_train_valid_df('/content/data/annotations/trainval.txt', 
                                               '/content/data/images/', '/content/data/annotations/xmls/', test_size=0.1)

전체적인 구조는 annotations txt파일에 있는 정보를 이용하여 train과 validation으로 나눠서 각각 img_filepath와 annotation_filepath를 지정해주어 저장해주는 코드이다.

pet_df.head()

실제로 경로가 알맞게 설정된 것을 볼 수 있다.

 

 

 

이후 oxford pet 데이터 셋의 annotation을 Ultralytics yolo format으로 변경하여 생성해주어야 한다.

 

oxford pet 데이터셋은 pascal voc 형태의 xml 파일구조로 되어있기 때문에 yolo 포맷용 txt 파일로 변경한다.

import glob
import xml.etree.ElementTree as ET

# 1개의 voc xml 파일을 Yolo 포맷용 txt 파일로 변경하는 함수이다.
def xml_to_txt(input_xml_file, output_txt_file, object_name):
  # ElementTree로 쉽게 xml파일을 파싱할 수 있다.
  tree = ET.parse(input_xml_file)
  root = tree.getroot()
  img_node = root.find('size')
  # img_node를 찾지 못하면 종료
  if img_node is None:
    return None
  # 원본 이미지의 너비와 높이 추출. 
  img_width = int(img_node.find('width').text)
  img_height = int(img_node.find('height').text)

  # xml 파일내에 있는 모든 object Element를 찾음. 
  value_str = None
  with open(output_txt_file, 'w') as output_fpointer:
    for obj in root.findall('object'):
        # bndbox를 찾아서 좌상단(xmin, ymin), 우하단(xmax, ymax) 좌표 추출. 
        xmlbox = obj.find('bndbox')
        x1 = int(xmlbox.find('xmin').text)
        y1 = int(xmlbox.find('ymin').text)
        x2 = int(xmlbox.find('xmax').text)
        y2 = int(xmlbox.find('ymax').text)
        # 만약 좌표중에 하나라도 0보다 작은 값이 있으면 종료. 
        if (x1 < 0) or (x2 < 0) or (y1 < 0) or (y2 < 0):
          break
        # object_name과 원본 좌표를 입력하여 Yolo 포맷으로 변환하는 convert_yolo_coord()함수 호출. 
        class_id, cx_norm, cy_norm, w_norm, h_norm = convert_yolo_coord(object_name, img_width, img_height, x1, y1, x2, y2)
        # 변환된 yolo 좌표를 object 별로 출력 text 파일에 write
        value_str = ('{0} {1} {2} {3} {4}').format(class_id, cx_norm, cy_norm, w_norm, h_norm)
        output_fpointer.write(value_str+'\n')

# object_name과 원본 좌표를 입력하여 Yolo 포맷으로 변환
def convert_yolo_coord(object_name, img_width, img_height, x1, y1, x2, y2):
  # class_id는 CLASS_NAMES 리스트에서 index 번호로 추출. 
  class_id = CLASS_NAMES.index(object_name)
  # 중심 좌표와 너비, 높이 계산. 
  center_x = (x1 + x2)/2
  center_y = (y1 + y2)/2
  width = x2 - x1
  height = y2 - y1
  # 원본 이미지 기준으로 중심 좌표와 너비 높이를 0-1 사이 값으로 scaling
  center_x_norm = center_x / img_width
  center_y_norm = center_y / img_height
  width_norm = width / img_width
  height_norm = height / img_height

  return class_id, round(center_x_norm, 7), round(center_y_norm, 7), round(width_norm, 7), round(height_norm, 7)

이후 voc format의 여러개의 xml파일을 yolo format으로 변환을 거친 후 ultralytics directory 구조로 변경한다.

import shutil

def make_yolo_anno_file(df, tgt_images_dir, tgt_labels_dir):
  for index, row in df.iterrows():
    src_image_path = row['img_filepath']
    src_label_path = row['anno_filepath']
    # 이미지 1개당 단 1개의 오브젝트만 존재하므로 class_name을 object_name으로 설정.  
    object_name = row['class_name']
    # yolo format으로 annotation할 txt 파일의 절대 경로명을 지정.
    target_label_path = tgt_labels_dir + row['img_name']+'.txt'
    # image의 경우 target images 디렉토리로 단순 copy
    shutil.copy(src_image_path, tgt_images_dir)
    # annotation의 경우 xml 파일을 target labels 디렉토리에 Ultralytics Yolo format으로 변환하여  만듬
    xml_to_txt(src_label_path, target_label_path, object_name)

# train용 images와 labels annotation 생성. 
make_yolo_anno_file(train_df, '/content/ox_pet/images/train/', '/content/ox_pet/labels/train/')
# val용 images와 labels annotation 생성. 
make_yolo_anno_file(val_df, '/content/ox_pet/images/val/', '/content/ox_pet/labels/val/')

이제 ultralytics yolo 학습을 위한 train,val img경로 설정과 annotation파일을 ultralytics yolo가 학습할 수 있는 형태로 변환을 마치고 학습 경로까지 지정을 마친 상태이다.

 

이후 dataset용 yaml파일을 생성하고

학습을 진행해보면 대락 1시간정도의 학습을 하게 된다.

!cd /content/yolov3; python train.py --img 640 --batch 16 --epochs 20 --data /content/ox_pet/ox_pet.yaml --weights yolov3.pt --project=/mydrive/ultra_workdir \
                                     --name pet --exist-ok

이미지 size 640*640, 16 batch size 20 epoch 을 설정 후 학습을 진행하였다.

학습 결과 mAP 0.5 기준일때의 mAP가 0.925가 나왔다. validation set부분을 train data set으로 설정하여 과적합이 일어났다. 성능보다는 지금은 Ultralytics Yolo v3패키지를 실습해보았다는 것에 의미를 두면 될 것 같다.

detection 수행 이미지