2.1. Custom Dataset으로 TFRecord 파일 만들기

 

 

 

 Tensorflow Object Detection API를 사용해서 모델 재학습을 하려면 TFRecord 형태의 input data를 만들어줘야 한다. TFRecord는 "Tensorflow 전용 학습 데이터 저장 형태" 정도로 알아두면 될 것 같다. 제대로 알아보지는 않았지만 Tensorflow로 설계한 모델 학습에 TFRecord를 사용하면 학습 속도가 개선되는 장점이 있다고 한다. 자세하게 알고 싶다면 여기1와 여기2에 들어가서 관련 설명을 읽어보면 될 것 같다.

 

 그리고 이미지를 byte 단위로 읽고 쓰는 법에 대해 모르는 경우라면, [Python] - 이미지 읽는 방법 / cv.imdecode( ), io.BytesIO( ) 을 읽으면 도움이 될 거다.

 

 이제 본격적으로 Custom Dataset을 TFRecord 파일로 만드는 방법에 대해 알아보자. 코드는 여기1, 여기2 그리고 './TF_Object_Detection_API/object_detection/dataset_tools/create_coco_tf_record.py 를 참고해서 작성했다.

 

 

 핵심부터 설명하자면,

다음 사진과 같이 feature_dict에서 요구하는 데이터들에 맞춰 Custom Dataset에서 뽑아내 넣어주면 된다.

 

 

 

1. Dataset 다운로드

 여기1에서 다운받아 정리하거나 여기2에서 다운받아 사용하면 된다. Dataset 폴더는 ./TF_Object_Detection_API 밑에 위치시켜주도록 하자. 이 Dataset은 내시경 영상에서 수술 도구(Surgical Instrument)를 Segmentation하기 위해 만들어진 Dataset이다.

 

 

 

 

2. image_to_tf_data( ) 함수 설계

import cv2, os, io
import numpy as np
import tensorflow as tf
from PIL import Image
from object_detection.utils import dataset_util

def image_to_tf_data(img_path, mask_path, class_name):
    filename = img_path.split('/')[-1]

    # image
    with tf.gfile.GFile(img_path, 'rb') as fid: # 이미지를 binary mode로 읽음
        encoded_img = fid.read()
    encoded_img_np = np.fromstring(encoded_img, dtype = np.uint8)
    image = cv2.imdecode(encoded_img_np, cv2.IMREAD_COLOR)
    height, width, _ = image.shape

    # mask
    with tf.gfile.GFile(mask_path, 'rb') as fid:
        encoded_mask = fid.read()
    encoded_mask_np = np.fromstring(encoded_mask, dtype = np.uint8)
    mask = cv2.imdecode(encoded_mask_np, cv2.IMREAD_GRAYSCALE) # mask는 1차원으로 사용함
    mask_pixel_vals = []
    for val in range(1, 256):
        if val in mask:
            mask_pixel_vals.append(val)

    # 필요 데이터들 추출 뒤 저장할 변수들
    classes, classes_text = [], []
    xmins, ymins, xmaxs, ymaxs = [], [], [], []
    encoded_mask_png_list = []
    for pixel_val in mask_pixel_vals:
        class_idx = 1

        # xmin, ymin, xmax, ymax 찾는 과정
        check_x_coordi = np.any(mask == pixel_val, axis = 0)
        check_y_coordi = np.any(mask == pixel_val, axis = 1)
        object_x_coordi = np.where(check_x_coordi)[0]
        object_y_coordi = np.where(check_y_coordi)[0]

        xmin = min(object_x_coordi)
        xmax = max(object_x_coordi)
        ymin = min(object_y_coordi)
        ymax = max(object_y_coordi)

        object_mask = np.uint8(mask == pixel_val) # mask의 최대값은 1 이다
        encoded_mask_png = cv2.imencode('.PNG', object_mask)[1].tobytes()  # mask는 PNG 확장자로

        classes.append(class_idx)
        classes_text.append(class_name.encode('utf8'))
        xmins.append(xmin / width)
        ymins.append(ymin / height)
        xmaxs.append(xmax / width)
        ymaxs.append(ymax / height)
        encoded_mask_png_list.append(encoded_mask_png)

    # TFRecord protocol 같은 부분임
    # 여기에 맞춰서 데이터 만들어서 넣어주면 됨
    feature_dict = {    
        'image/height':             dataset_util.int64_feature(height),
        'image/width':              dataset_util.int64_feature(width),
        'image/filename':           dataset_util.bytes_feature(filename.encode('utf8')),
        'image/source_id':          dataset_util.bytes_feature(filename.encode('utf8')),
        'image/key/sha256':         dataset_util.bytes_feature(class_name.encode('utf8')),
        'image/encoded':            dataset_util.bytes_feature(encoded_img),
        'image/format':             dataset_util.bytes_feature('png'.encode('utf8')),
        'image/object/bbox/xmin':   dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax':   dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin':   dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax':   dataset_util.float_list_feature(ymaxs),
        'image/object/class/text':  dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
        'image/object/mask':        dataset_util.bytes_list_feature(encoded_mask_png_list)}
    tf_data = tf.train.Example(features=tf.train.Features(feature=feature_dict))

    return tf_data​

• Line 10 ~ 21
  : 잘 모르겠다면 [Python] - 이미지 읽는 방법 / cv.imdecode( ), io.BytesIO( )을 읽길 바란다.
• Line 22 ~ 25
  : mask image 내에 객체가 최대 3개까지 있는데 각 객체는 다른 pixel 값으로 표현되어 있다.
   객체 별로 데이터를 만들어주기 위해서 작성한 부분이다.
   mask_pixel_vals는 Line 31에서 사용된다.
• Line 35 ~ 36
  : np.any(mask == pixel_val, axis = 0)은 axis = 0 으로 설정됐으므로
   
  

  1 2 3 4 5 6

  res = [] for i in range(width):     if pixel_val in mask[:, i]:         res.append(True)     else:         res.append(False)

  cs

  
  의 흐름으로  mask == pixel_val  여부를 판단해 결과를 반환해준다.
  ※ axis = 0은 [ height, width ]로 구성된 mask의 height 차원을 1로 만들어 반환하라는 뜻이기 때문에
     1열 벡터로 판단 후 결과 반환, 2열 벡터로 판단 후 결과 반환 … 의 흐름을 갖게 된다.
• Line 37 ~ 38
  : np.where( )에 bool array가 들어가면 True를 갖는 위치의 index들을 반환해준다.
  따라서  np.where(check_x_coordi)  는   check_x_coordi  에 포함된 True의 index를 찾아 반환해준다.
  
• Line 46
  : 이 부분 또한 잘 모르겠다면 [Python] - 이미지 읽는 방법 / cv.imdecode( ), io.BytesIO( )을 읽길 바란다.
• Line 48 ~ 75
  : feature_dict에서 요구하는 형태에 맞춰서 우리가 Custom Dataset에서 뽑아낸 정보들을 넣어주면 된다.
   (1) Line 50 ~ 53의 경우, 좌표를 이미지 크기에 상대적으로 넘겨줘야해서 width와 height으로 나눠줬다.
   (2) Line 65의 경우, Custom Dataset의 image 형식이 png라서 변경해줬다.

 

 

3. TFRecord 생성

 image_to_tf_data( ) 함수를 사용해서 뽑은 정보들을 TFRecrodWriter( )를 사용해서 TFRecord 파일로 저장해주면 된다.

base_dir = './Dataset/train/images'
img_names = os.listdir(base_dir)
img_paths = ['{}/{}' .format(base_dir, name) for name in img_names]

tf_writer = tf.python_io.TFRecordWriter('./Dataset/train.record')
for img_path in img_paths:
    print(img_path)
    mask_path = img_path.replace('/images', '/masks')
    mask_path = mask_path.replace('_image', '_mask')

    tf_example = image_to_tf_data(img_path, mask_path, 'instrument')
    
    tf_writer.write(tf_example.SerializeToString())
tf_writer.close()​

 

 

 

---

 

 

다음은 전체 코드이다.

[ create_tf_record.py ]

import cv2, os, io
import numpy as np
import tensorflow as tf
from PIL import Image
from object_detection.utils import dataset_util

def image_to_tf_data(img_path, mask_path, class_name):
    filename = img_path.split('/')[-1]

    # image
    with tf.gfile.GFile(img_path, 'rb') as fid: # 이미지를 binary mode로 읽음
        encoded_img = fid.read()
    encoded_img_np = np.fromstring(encoded_img, dtype = np.uint8)
    image = cv2.imdecode(encoded_img_np, cv2.IMREAD_COLOR)
    height, width, _ = image.shape

    # mask
    with tf.gfile.GFile(mask_path, 'rb') as fid:
        encoded_mask = fid.read()
    encoded_mask_np = np.fromstring(encoded_mask, dtype = np.uint8)
    mask = cv2.imdecode(encoded_mask_np, cv2.IMREAD_GRAYSCALE) # mask는 1차원으로 사용함
    mask_pixel_vals = []
    for val in range(1, 256):
        if val in mask:
            mask_pixel_vals.append(val)

    # 필요 데이터들 추출 뒤 저장할 변수들
    classes, classes_text = [], []
    xmins, ymins, xmaxs, ymaxs = [], [], [], []
    encoded_mask_png_list = []
    for pixel_val in mask_pixel_vals:
        class_idx = 1

        # xmin, ymin, xmax, ymax 찾는 과정
        check_x_coordi = np.any(mask == pixel_val, axis = 0)
        check_y_coordi = np.any(mask == pixel_val, axis = 1)
        object_x_coordi = np.where(check_x_coordi)[0]
        object_y_coordi = np.where(check_y_coordi)[0]

        xmin = min(object_x_coordi)
        xmax = max(object_x_coordi)
        ymin = min(object_y_coordi)
        ymax = max(object_y_coordi)

        object_mask = np.uint8(mask == pixel_val) # mask의 최대값은 1 이다
        encoded_mask_png = cv2.imencode('.PNG', object_mask)[1].tobytes()  # mask는 PNG 확장자로 저장해줘야하는 듯

        classes.append(class_idx)
        classes_text.append(class_name.encode('utf8'))
        xmins.append(xmin / width)
        ymins.append(ymin / height)
        xmaxs.append(xmax / width)
        ymaxs.append(ymax / height)
        encoded_mask_png_list.append(encoded_mask_png)

    # TFRecord protocol 같은 부분임
    # 여기에 맞춰서 데이터 만들어서 넣어주면 됨
    feature_dict = {    
        'image/height':             dataset_util.int64_feature(height),
        'image/width':              dataset_util.int64_feature(width),
        'image/filename':           dataset_util.bytes_feature(filename.encode('utf8')),
        'image/source_id':          dataset_util.bytes_feature(filename.encode('utf8')),
        'image/key/sha256':         dataset_util.bytes_feature(class_name.encode('utf8')),
        'image/encoded':            dataset_util.bytes_feature(encoded_img),
        'image/format':             dataset_util.bytes_feature('png'.encode('utf8')),
        'image/object/bbox/xmin':   dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax':   dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin':   dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax':   dataset_util.float_list_feature(ymaxs),
        'image/object/class/text':  dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
        'image/object/mask':        dataset_util.bytes_list_feature(encoded_mask_png_list)}
    tf_data = tf.train.Example(features=tf.train.Features(feature=feature_dict))

    return tf_data

base_dir = './Dataset/train/images'
img_names = os.listdir(base_dir)
img_paths = ['{}/{}' .format(base_dir, name) for name in img_names]

tf_writer = tf.python_io.TFRecordWriter('./Dataset/train.record')
for img_path in img_paths:
    print(img_path)
    mask_path = img_path.replace('/images', '/masks')
    mask_path = mask_path.replace('_image', '_mask')

    tf_example = image_to_tf_data(img_path, mask_path, 'instrument')
    
    tf_writer.write(tf_example.SerializeToString())
tf_writer.close()