TensorFlow를 이용한 YOLO v1 논문 구현 #5 - utils.py

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

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

utils.py

목표: 딥러닝 메인 로직 외에 유틸리티성 기능들을 조직

---

필요한 모듈 import

import cv2
import numpy as np
import tensorflow as tf
import colorsys
from operator import itemgetter

---

상세 코드

def draw_bounding_box_and_label_info(frame, x_min, y_min, x_max, y_max, label, confidence, color):
  draw_bounding_box(frame, x_min, y_min, x_max, y_max, color)
  draw_label_info(frame, x_min, y_min, label, confidence, color)
  
def draw_bounding_box(frame, x_min, y_min, x_max, y_max, color):
  cv2.rectangle(
    frame,
    (x_min, y_min),
    (x_max, y_max),
    color, 3)


def draw_label_info(frame, x_min, y_min, label, confidence, color):
  text = label + ' ' + str('%.3f' % confidence)
  bottomLeftCornerOfText = (x_min, y_min)
  font = cv2.FONT_HERSHEY_SIMPLEX
  fontScale = 0.8
  fontColor = color
  lineType = 2

  cv2.putText(frame, text,
              bottomLeftCornerOfText,
              font,
              fontScale,
              fontColor,
              lineType)

drawing하고자 하는 이미지의 label, bounding box와 추가 정보들을 출력하는 함수

위 함수의 예시

---

def find_max_confidence_bounding_box(bounding_box_info_list):
  bounding_box_info_list_sorted = sorted(bounding_box_info_list,
                                                   key=itemgetter('confidence'),
                                                   reverse=True)
  max_confidence_bounding_box = bounding_box_info_list_sorted[0]

  return max_confidence_bounding_box

정석적인 NMS 방법이 아니라, bounding box들에 대한 confidence중 가장 큰 값 하나를 선택해 return해주는 함수

---

def yolo_format_to_bounding_box_dict(xcenter, ycenter, box_w, box_h, class_name, confidence):
  bounding_box_info = {}
  bounding_box_info['left'] = int(xcenter - (box_w / 2))
  bounding_box_info['top'] = int(ycenter - (box_h / 2))
  bounding_box_info['right'] = int(xcenter + (box_w / 2))
  bounding_box_info['bottom'] = int(ycenter + (box_h / 2))
  bounding_box_info['class_name'] = class_name
  bounding_box_info['confidence'] = confidence

  return bounding_box_info

bounding box에 대한 정보들을 저장하는 함수

---

def iou(yolo_pred_boxes, ground_truth_boxes):
  # Reference : https://github.com/nilboy/tensorflow-yolo/blob/python2.7/yolo/net/yolo_tiny_net.py#L105
  boxes1 = yolo_pred_boxes
  boxes2 = ground_truth_boxes

  boxes1 = tf.stack([boxes1[:, :, :, 0] - boxes1[:, :, :, 2] / 2, boxes1[:, :, :, 1] - boxes1[:, :, :, 3] / 2,
                     boxes1[:, :, :, 0] + boxes1[:, :, :, 2] / 2, boxes1[:, :, :, 1] + boxes1[:, :, :, 3] / 2])
  boxes1 = tf.transpose(boxes1, [1, 2, 3, 0])
  boxes2 = tf.stack([boxes2[0] - boxes2[2] / 2, boxes2[1] - boxes2[3] / 2,
                     boxes2[0] + boxes2[2] / 2, boxes2[1] + boxes2[3] / 2])
  boxes2 = tf.cast(boxes2, tf.float32)

  lu = tf.maximum(boxes1[:, :, :, 0:2], boxes2[0:2])
  rd = tf.minimum(boxes1[:, :, :, 2:], boxes2[2:])

  #intersection
  intersection = rd - lu

  inter_square = intersection[:, :, :, 0] * intersection[:, :, :, 1]

  mask = tf.cast(intersection[:, :, :, 0] > 0, tf.float32) * tf.cast(intersection[:, :, :, 1] > 0, tf.float32)

  inter_square = mask * inter_square

  #bounding box1의 넓이와 bounding box2의 넓이를 이용한 계산
  square1 = (boxes1[:, :, :, 2] - boxes1[:, :, :, 0]) * (boxes1[:, :, :, 3] - boxes1[:, :, :, 1])
  square2 = (boxes2[2] - boxes2[0]) * (boxes2[3] - boxes2[1])

  return inter_square / (square1 + square2 - inter_square + 1e-6)

IOU 값을 구하기 위한 함수 정의이며, yolo가 예측한 영역과 정답 영역을 인자로 받아 계산

IOU의 계산에 대한 것은 다음 링크 참조: https://it-the-hunter.tistory.com/29

[딥러닝]IOU에 대해서 이해해보자

---

#bounding box를 그리기 위한 color 할당
def generate_color(num_classes):
  # Reference : https://github.com/qqwweee/keras-yolo3/blob/e6598d13c703029b2686bc2eb8d5c09badf42992/yolo.py#L82
  hsv_tuples = [(x / num_classes, 1., 1.)
                for x in range(num_classes)]
  colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
  colors = list(
    map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
        colors))
  np.random.seed(10101)  
  np.random.shuffle(colors)  
  np.random.seed(None)  

  return colors

output 이미지에서 오브젝트를 검출하고, 그 bounding box를 그리기 위해서는 오브젝트마다 다른 색을 쓰는게 가시성에 좋으므로, 이를 위한 함수 정의