본문 바로가기
개인 공부/인공지능, 딥러닝

1. 내 데이터로 객체 인식 학습시키기 Object Detection with Custom Dataset :: tensorflow

by seowit 2021. 2. 5.

Jetson Nano에서 Yolo를 이용해서 object detection을 수행했더니 너무 느리더라고요,,, FPS가 10도 안 나오는 것 같아요,,,

그래서 찾아보니까 SSD Mobilenet 이 젯슨 나노에서 빠르게 잘 돌아가는 예제를 보고 바로 학습을 시도했습니다!!

 


 

첫 번째 시도는, 아래 깃헙에 나와있는 Re-training SSD-Mobilenet 을 이용해서 해봤는데 open images를 안 쓰고 custom dataset 을 사용하니까 자꾸 에러가 뜨더라고요. detectnet을 뜯어보니 open image와 voc 두 가지 형태의 데이터셋을 취급하는 것 같은데, 에러가 뜨길래 다른 방법을 시도했습니다!

github.com/dusty-nv/jetson-inference

 

dusty-nv/jetson-inference

Hello AI World guide to deploying deep-learning inference networks and deep vision primitives with TensorRT and NVIDIA Jetson. - dusty-nv/jetson-inference

github.com

 


 제가 참고한 사이트 입니다 ↓

towardsdatascience.com/custom-object-detection-using-tensorflow-from-scratch-e61da2e10087

 

Custom Object Detection using TensorFlow from Scratch

Custom Dataset Training for Object Detection using TensorFlow | Dog Detection in Real time Videos | Perfect Guide for Object Detection

towardsdatascience.com

 

FYI, 제 환경은 Ubuntu 20.04 / cuda 10.2 / cudnn 7.6.5 / gpu 1080ti 이고 Anaconda 아나콘다로 가상환경 세팅 했습니다.

 

0. 가상 환경 세팅 & 설치

 

▷  tf2라는 이름의 가상환경을 생성하고 활성화합니다. tensorflow 2.*라는 뜻입니다.

$ conda create -n tf2
$ conda activate tf2

▷  아래와 같은 명령어를 입력하시면, 2021년 02월 04일 기준 텐서플로우 2.4 버전이 설치됩니다. 이후에 학습할 때, 버전 충돌 때문에 1.15.0 버전의 가상환경 tf1을 새롭게 생성합니다. 처음부터 1.15.0 버전을 설치해도 될지 모르겠습니다. 

$ pip install tensorflow
$ pip install tensorflow-gpu

▷  다른 의존성 패키지를 설치합니다.

$ pip install pillow Cython lxml jupyter matplotlib

 

▶  tensorflow 디렉토리로 이동하고, tensorflow model 텐서플로우 모델을 다운로드합니다.

# tensorflow 디렉토리 :  ~/anaconda/envs/<가상환경명>/lib/python3.7/site-packages/tensorflow

$ cd anaconda/envs/<가상환경명>/lib/python3.7/site-packages/tensorflow # 이 부분은 조금씩 다를 수 있습니다.
$ git clone https://github.com/tensorflow/models.git

▶ protobuf 컴파일하고, PYTHONPATH를 변경해줍니다. 

$ cd <path_to_your_tensorflow_installation>/models/research/
$ protoc object_detection/protos/*.proto --python_out=.
$ export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
가끔 아래와 같은 에러가 뜰 때가 있는데, 위의 박스 명령어 3줄을 입력해주시면 해결됩니다
ModuleNotFoundError: No module named 'object_detection'

 

▷ 설치가 잘 됐는지 확인합니다. 아무런 창도 뜨지 않고 뭔가 됐다는 메시지가 터미널에 출력됩니다.

$ python object_detection/builders/model_builder_test.py

 

1. 폴더 구조화 & 이미지 데이터 정리

 

▷ 추천하는 models 폴더의 structure 입니다.  annotations, images 등 폴더를 생성합니다 ( $ mkdir annotations )

annotations/xmls에 데이터 라벨링한 xml 파일을 저장하고, images 아래 train 과 test에는 jpg 파일을 적당한 비율(7:3)로 나누어 저장합니다. 

저는 Yolo 학습시킬 때 사용한 데이터를 VOC dataset format으로 변환해서 사용해서 따로 이미지 라벨링은 진행하지 않았으나, 위에 참고한 사이트를 가보시면 데이터 라벨링하는 방법도 자세하게 나와있습니다. 

models
    - annotations
          -- xmls
    - images
          -- train
          -- test
    - checkpoints
    - tf_record
    - research
    - ...

 

▶  label_map.txt 파일을 생성합니다. 1부터 시작해야합니다! 예를 들어, 사람과 강아지라는 클래스가 있다면 아래와 같이 작성합니다.

item {
  id: 1
  name: 'person'
}
item {
  id: 2
  name: 'dog'
}

 

 

▶  trainval.txt 파일을 생성합니다. 아래와 같이 이미지 이름을 확장자 없이 나열합니다.

000001
000002
...
018000
018001
저는 모든 데이터의 파일명이 숫자로 되어 있어서 다음의 코드로 trainval.txt 파일을 생성했습니다

github.com/swkim-sm/shortNsimple/blob/main/automation/make_trainval_txt.py

 

▶ XML 파일을 CSV 파일로 전환해야 합니다. 아래의 코드로 실행하시면 됩니다.

gist.github.com/iKhushPatel/ed1f837656b155d9b94d45b42e00f5e4

 

xml_to_csv.py

GitHub Gist: instantly share code, notes, and snippets.

gist.github.com

더보기

저는 float str 어쩌구 에러가 떠서 아래와 같이 바꾸어서 사용했더니 해결되었습니다. 

저는 float str 어쩌구 에러가 떠서 아래와 같이 바꾸어서 사용했더니 해결되었습니다. 

import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET


def xml_to_csv(path):
    xml_list = []
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        flag = False
        for member in root.findall('object'):
            value = (root.find('filename').text,
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member[0].text,
                     int(member.find('bndbox')[0].text),
                     int(member.find('bndbox')[1].text),
                     int(member.find('bndbox')[2].text),
                     int(member.find('bndbox')[3].text)
                     #int(member[4][0].text),
                     #int(member[4][1].text),
                     #int(member[4][2].text),
                     #int(member[4][3].text)
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
    xml_df = pd.DataFrame(xml_list, columns=column_name)
    return xml_df


def main():
	'''
    for directory in ['train','testing']:
        image_path = os.path.join(os.getcwd(), 'images/{}'.format(directory).format(directory))
        xml_df = xml_to_csv(image_path)
        xml_df.to_csv('data/{}_labels.csv'.format(directory), index=None)
        print('Successfully converted xml to csv.')
	''' 

	image_path = os.path.join(os.getcwd(), 'images/train')
	xml_df = xml_to_csv(image_path)
	xml_df.to_csv('data/train_labels.csv', index=None)

	image_path = os.path.join(os.getcwd(), 'images/test')
	xml_df = xml_to_csv(image_path)
	xml_df.to_csv('data/test_labels.csv',index=None)

main()

이를 수행하고 나면 data 폴더 안에 train_labels.csv와 test_labels.csv가 저장됩니다.

 

2. TFRecord 생성 및 config 파일 수정

 

 TFRecord 파일을 생성합니다.

아래의 코드(generate_tfrecord.py)를 자기 클래스에 맞춰서 수정하고, models/research/object_detection/dataset_tools/ 에 저장합니다. 

github.com/datitran/raccoon_dataset/blob/master/generate_tfrecord.py#L25

 

datitran/raccoon_dataset

The dataset is used to train my own raccoon detector and I blogged about it on Medium - datitran/raccoon_dataset

github.com

이 부분을

# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'raccoon':
        return 1
    else:
        None

아래와 같이 변경합니다. label_map.txt 의 id가 여기서는 return 값이 되어야 합니다.

# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'person':
        return 1
    if row_label == 'dog':
        return 2
    else:
    	None

 

models에서 코드를 수행합니다. train과 test 두 번을 수행합니다

# train

$ python research/object_detection/dataset_tools/generate_tfrecord.py --csv_input=data/train_labels.csv --output_path=train.record --image_dir=images/train
# test

$ python research/object_detection/dataset_tools/generate_tfrecord.py --csv_input=data/test_labels.csv --output_path=test.record --image_dir=images/test
AttributeError: module 'tensorflow' has no attribute 'app'
이런 에러가 뜨시는 분들은 generate_tfrecord.py 에서 
import tensorflow as tf 를
import tensorflow.compat.v1 as tf 로 바꾸어 주시면 해결 됩니다.

모두 수행하고 나면 train.record와 test.record 파일이 생성됩니다. 

 

▶  아래에서 pre-trained model을 다운로드 받습니다.

download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz

위 파일의 압축을 풀고, model.ckpt.meta, model.ckpt.index, model.ckpt.data-00000-of-00001 3개의 파일을 models/checkpoints 에 저장합니다.

 

▶  .config 파일을 수정합니다.

models/research/object_detection/samples/configs 아래에 있는 ssd_mobilenet_v2_ coco 파일을 models/ 로 복사합니다.

ssd_mobilenet_v2_ coco 파일에서 3가지 사항을 수정합니다. 

1. num_classes 변경

ssd {
    num_classes: 21 # 자기 클래스 개수로 변경
    box_coder {
        faster_rcnn_box_coder {
            y_scale: 10.0 
            x_scale: 10.0 
            height_scale: 5.0 
            width_scale: 5.0 
    } 
 }

2. fine_tune_checkpoint 변경

fine_tune_checkpoint: "checkpoints/model.ckpt" # model.ckpt 경로 입력

3. train, eval input reader 변경

train_input_reader: { 
    tf_record_input_reader { 
        input_path: "train.record" # train.record 경로 입력
    } 
    label_map_path: "annotations/label_map.pbtxt"  #  label_map.pbtxt 경로 입력
  }

eval_input_reader: { 
    tf_record_input_reader { 
        input_path: "test.record" # test.record 경로 입력
    } 
    label_map_path: "annotations/label_map.pbtxt"  # label_map.pbtxt 경로 입력
    shuffle : false
    num_readers : 1
  }

 

3. Train

 

▶  train.py 파일 실행

# models 디렉토리에서 실행

$ mkdir train
$ mkdir eval

$ python research/object_detection/legacy/train.py --logtostderr --train_dir=train --pipepline_config_path=ssd_mobilenet_v2_coco.config
Train.py not found

research/object_detection 아래에 train.py 파일이 있다고 했는데, 저는 legacy 폴더 안에 들어있었습니다. 혹시 train.py 파일을 못 찾는 분들은 legacy 폴더에 들어가보세요!
ValueError: ssd_mobilenet_v2 is not supported. See `model_builder.py` for features extractors compatible with different versions of Tensorflow

위의 에러는 tensorflow 1.15.0 버전을 설치하시면 해결됩니다. 저는 tf1이라는 가상환경을 새로 만들어서 train은 tf1에서 수행하였습니다. 

 


train.py를 실행하고 20시간 정도 지났는데, 아직 끝나지 않았습니다.

완료되면 마저 작성하도록 하겠습니다 ㅎㅎ

댓글