TensorFlow를 이용한 YOLO v1 논문 구현 #4 - datasets.py

이전 글: https://it-the-hunter.tistory.com/28

TensorFlow를 이용한 YOLO v1 논문 구현 #3 - loss.py

datasets.py

목표: 데이터 전처리 및 batch 단위로 묶는 로직

---

필요한 모듈 import

import tensorflow as tf
import numpy as np

---

상세 코드

#train.py에서 처리하는 데이터를 변경
def process_each_ground_truth(original_image,
                              bbox,
                              class_labels,
                              input_width,
                              input_height
                              ):

train.py에서 처리하는, 위와 같은 데이터들을 원하는 형태로 변경

• original_image: 이미지 resize
• bbox: 각 값들을 normalize해 적용, 절대 좌표로 변경
• class_labels: 정답 class label intger로 받음
• input_ : yolo 모델에 입력하는 폭, 높이 값

---

def bounds_per_dimension(ndarray):
  return map(
    lambda e: range(e.min(), e.max() + 1),
    np.where(ndarray != 0)
  )


def zero_trim_ndarray(ndarray):
  return ndarray[np.ix_(*bounds_per_dimension(ndarray))]

  image = original_image.numpy()
  image = zero_trim_ndarray(image)

padding된, 검은 색으로 표시된 부분을 제거해주며 나머지 부분을 반환

---

  #원본 폭, 높이 설정
  original_h = image.shape[0]
  original_w = image.shape[1]

  width_rate = input_width * 1.0 / original_w
  height_rate = input_height * 1.0 / original_h

  image = tf.image.resize(image, [input_height, input_width])

논문에 따르면 input 이미지를 resize해주기 때문에, 비율을 맞추기 위한 전처리

---

  #해당 이미지에 실제 오브젝트 수를 parsing
  object_num = np.count_nonzero(bbox, axis=0)[0]
  labels = [[0, 0, 0, 0, 0]] * object_num
  #절대 좌표로 변경
  for i in range(object_num):
    xmin = bbox[i][1] * original_w
    ymin = bbox[i][0] * original_h
    xmax = bbox[i][3] * original_w
    ymax = bbox[i][2] * original_h

    class_num = class_labels[i]
  
    xcenter = (xmin + xmax) * 1.0 / 2 * width_rate
    ycenter = (ymin + ymax) * 1.0 / 2 * height_rate

    box_w = (xmax - xmin) * width_rate
    box_h = (ymax - ymin) * height_rate

    labels[i] = [xcenter, ycenter, box_w, box_h, class_num]
  
  return [image.numpy(), labels, object_num]  #resize된 이미지, 절대 좌표의 label, 실제 오브젝트의 수

계산에 필요한 각 값들을 batch단위로 묶기 위한 계산 부분

마지막에 resize된 input 이미지, 절대 좌표로 변경된 label, 실제 오브젝트의 수를 반환하여 train.py에 반영