[Python] Simple YoLo 실습

# 나만의 YOLO 모델 생성
# Darknet 수준의 YOLO 모델을 개인이 학습하기에는 
# 컴퓨터 리소스도 많고, 시간도 너무 많이 걸린다는 점.

# 검은색 바탕에 간단한 도형3개만 탐지하는 YOLO 모델을 구현

# YOLO 논문에서는 이미지를 가로 세로 각 7개의 셀로 나누어 총 49개의 셀을 기본으로 하지만,
# 우리는 가로 세로 3개의 셀로 나누는 방식으로 문제를 단순화.

# 논문에서는 한 셀당 2개의 박스를 그리지만,
# 우리는 한 셀당 1개의 박스를 그리는 방식으로 수정.
# 마지막 탐지할 객체의 종류인 Class도 3개로 줄여서 구현

In [1]:

# 필요한 패키지를 임포트 함
import tensorflow as tf
import numpy as np
import cv2
from google.colab.patches import cv2_imshow

# 파라미터 설정

# 이미지 크기
width_size = 256
hight_size = 256
channel_size = 3
img_size = (width_size,hight_size,channel_size)

# 이미지를 나눌 크기
cell_num = 3

# 찾고자 하는 객체 개수
class_num = 3

# 한셀에 그릴 박스 수
anchor_num = 1
label_num = anchor_num * (5 + class_num)

# 학습 수
epoch_num = 20000

# 로스 비중
loss_p_rate = 1.0
loss_cod_rate = 5.0
loss_c_rate = 1.0
loss_p_no_rate = 0.5

In [ ]:

# 제공하는 3개의 이미지 파일(0.png, 1.png, 2.png)을 사용.
# 코랩 폴더에 3개의 이미지 파일을 업로드

In [ ]:

# CV2를 이용하여 랜덤한 위치에 3개의 도형 이미지를 그린다
# 해당 이미지의 위치를 찾아서 경계박스로 나타내고, 정답 클래스 레이블까지 반환하는 함수를 정의

In [6]:

# 랜덤하게 도형을 그리고, 실제 정답 값을 생성하는 함수 정의
# 0.png / 1.png / 2.png 파일이 필요함

def make_img_label():
    img = np.zeros((hight_size+400,width_size+400,channel_size))
    label = np.zeros((cell_num,cell_num,label_num))
    num_shape = np.random.randint(1,4)
    i = np.random.choice(range(cell_num),num_shape,replace=False)
    j = np.random.choice(range(cell_num),num_shape,replace=False)
    
    img_0 = cv2.imread('0.png')
    img_1 = cv2.imread('1.png')
    img_2 = cv2.imread('2.png')
    
    for n_h in range(num_shape):
        row = i[n_h]
        col = j[n_h]
        
        shape_type = np.random.randint(0,class_num)
        x_rate = np.random.rand()
        y_rate = np.random.rand()
        w_rate = np.random.rand() * 0.3 +0.1
        h_rate = np.random.rand() * 0.3 +0.1
                
        label[row,col]=[1,x_rate,y_rate,w_rate,h_rate,0,0,0]
        label[row,col,5+shape_type]=1

        x = int(x_rate * width_size/cell_num + col * width_size/cell_num)
        y = int(y_rate * hight_size/cell_num + row * hight_size/cell_num)
        w = int(w_rate * width_size/2) * 2
        h = int(h_rate * hight_size/2) * 2

        if(shape_type==0):
            input_img = cv2.resize(img_0,(w,h))
        if(shape_type==1):
            input_img = cv2.resize(img_1,(w,h))
        if(shape_type==2):
            input_img = cv2.resize(img_2,(w,h))

        img[y-int(h/2)+200 : y+int(h/2)+200, x-int(w/2)+200 : x+int(w/2)+200] =input_img
    img = img[200:200+hight_size,200:200+width_size]        
    
    return img,label

img,label = make_img_label()
cv2_imshow(img)

In [ ]:

# 실습을 위해 생성된 이미지와 클래스(또는 예측 값)를 입력해주면,
# 탐지한 이미지에 박스를 그려주는 함수를 정의
# 함수를 실행하면 경계박스를 찾아서 표시해준다.

In [7]:

# 이미지와 정답(혹은 예측값)을 넣으면 박스를 그려주는 함수 정의
# 임계값 th 설정 (객체가 있다는 확률이 th이상일 때만 박스 생성)
def show_box(img,label,th=0.3):
    b_img = np.zeros((hight_size+400,width_size+400,3))
    b_img[200:200+hight_size,200:200+width_size] = img
    for i in range(cell_num):
        for j in range(cell_num):
            if(label[i,j,0] > th):
                x_rate = label[i,j,1]
                y_rate = label[i,j,2]
                w_rate = label[i,j,3]
                h_rate = label[i,j,4]
                shape_type=np.argmax(label[i,j,5:])
                if(shape_type==0):
                    line_color = [0,0,255]
                if(shape_type==1):
                    line_color = [255,0,0]
                if(shape_type==2):
                    line_color = [0,255,0]
                x = int(x_rate * width_size/3 + j * width_size/3)
                y = int(y_rate * hight_size/3 + i * hight_size/3)
                w = int(w_rate * width_size/2) * 2 + 20
                h = int(h_rate * hight_size/2) * 2 + 20

                cv2.rectangle(b_img,(x-int(w/2)+200,y-int(h/2)+200),(x+int(w/2)+200,y+int(h/2)+200),line_color)
                
    b_img = b_img[200:200+hight_size,200:200+width_size]

    return b_img

cv2_imshow(show_box(img,label))

In [ ]:

# 논문에서 구현하고 있는 NMS(Non-Maximum Suppression)은 여기에 적요하지 않는다.
# NMS : 서로 다른 두 박스가 하나의 객체를 탐지할 경우, 예측 확률이 박스를 지우는 알고리즘.
# 따라서 우리 모델은 특정 도형ㅇ 이미지를 탐지한 모든 경계박스가 표시된다.

In [ ]:

# 객체 탐지 모델이 어느 정도 성능을 갖기 위해서는 복잡한 구조로 구현되어야 한다.
# 전이 학습 방법을 적용하여
# 이미지의 특징을 추출하는데 좋은 성능을 갖는 모델을 기본으로 활용하는 것이 좋다.

In [8]:

# 전이 학습 방법 적용
# VGG16모델을 베이스로 하고
# Conv2D 층과 Dense 레이어를 마지막 객체 탐지 분류가로 설정해준다.
# 모델 구조를 요약하여 확인

# VGG16모델을 베이스로 마지막 부분만 수정한 모델 (전이학습)
vgg_model = tf.keras.applications.VGG16(include_top=False,input_shape=img_size)
vgg_model.trainable=False

i=tf.keras.Input(shape=img_size)
out=tf.keras.layers.Lambda((lambda x : x/255.))(i)
out = vgg_model(out)
out = tf.keras.layers.Conv2D(256,3,padding='same')(out)
out = tf.keras.layers.Conv2D(128,3,padding='same')(out)
out = tf.keras.layers.Conv2D(64,3,padding='same')(out)
out = tf.keras.layers.Flatten()(out)
out = tf.keras.layers.Dense(1024,activation='relu')(out)
out = tf.keras.layers.Dense(3*3*8,activation='sigmoid')(out)
out = tf.keras.layers.Reshape((3,3,8))(out)

yolo_model = tf.keras.Model(inputs=[i],outputs=[out])
opt = tf.keras.optimizers.Adam(0.00001)

# 모델 요약
yolo_model.summary()

Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/vgg16/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5
58892288/58889256 [==============================] - 0s 0us/step
58900480/58889256 [==============================] - 0s 0us/step
Model: "model"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 input_2 (InputLayer)        [(None, 256, 256, 3)]     0         
                                                                 
 lambda (Lambda)             (None, 256, 256, 3)       0         
                                                                 
 vgg16 (Functional)          (None, 8, 8, 512)         14714688  
                                                                 
 conv2d (Conv2D)             (None, 8, 8, 256)         1179904   
                                                                 
 conv2d_1 (Conv2D)           (None, 8, 8, 128)         295040    
                                                                 
 conv2d_2 (Conv2D)           (None, 8, 8, 64)          73792     
                                                                 
 flatten (Flatten)           (None, 4096)              0         
                                                                 
 dense (Dense)               (None, 1024)              4195328   
                                                                 
 dense_1 (Dense)             (None, 72)                73800     
                                                                 
 reshape (Reshape)           (None, 3, 3, 8)           0         
                                                                 
=================================================================
Total params: 20,532,552
Trainable params: 5,817,864
Non-trainable params: 14,714,688
_________________________________________________________________

In [ ]:

# 이미지를 총 9개(3*3)의 셀로 나누고, 셀마다 학습을 진행.
# 객체가 있는 셀의 경우 확률/박스위치 및 크기/클래스의 종류 모두 학습을 진행하고,
# 객체가 없는 셀은 객체가 없는 확률만 학습한다.

# 각 Loss 는 미리 정한 비중을 곱하고, 전체를 더하여 최종 Loss를 만들어 학습시킨다.

In [9]:

# 학습과정을 동영상으로 기록
fcc=cv2.VideoWriter_fourcc(*'DIVX')
out=cv2.VideoWriter('my_yolo.avi',fcc,1.0,(width_size,hight_size))


for e in range(epoch_num):
    img,label = make_img_label()
    img = np.reshape(img,(1,hight_size,width_size,3))
    label = np.reshape(label,(1,3,3,8))
    loss_p_list=[]
    loss_cod_list = []
    loss_c_list = []
    loss_p_no_list = []
    with tf.GradientTape() as tape:
        pred = yolo_model(img)
        # 이미지를 구분한 셀을 탐험
        for i in range(3):
            for j in range(3):
                # 해당 셀에 객체가 있을 경우는 확률, 박스 크기, 클래스까지 모두 Loss로 계산
                if(label[0,i,j,0]==1):
                    loss_p_list.append(tf.square(label[0,i,j,0]-pred[0,i,j,0]))
                    loss_cod_list.append(tf.square(label[0,i,j,1]-pred[0,i,j,1]))
                    loss_cod_list.append(tf.square(label[0,i,j,2]-pred[0,i,j,2]))
                    loss_cod_list.append(tf.square(label[0,i,j,3]-pred[0,i,j,3]))
                    loss_cod_list.append(tf.square(label[0,i,j,4]-pred[0,i,j,4]))
                    loss_c_list.append(tf.square(label[0,i,j,5]-pred[0,i,j,5]))
                    loss_c_list.append(tf.square(label[0,i,j,6]-pred[0,i,j,6]))
                    loss_c_list.append(tf.square(label[0,i,j,7]-pred[0,i,j,7]))
                # 해당 셀에 객체가 없을 경우 객체가 없을 확률만 Loss로 계산
                else:
                    loss_p_no_list.append(tf.square(label[0,i,j,0]-pred[0,i,j,0]))
        loss_p=tf.reduce_mean(loss_p_list)
        loss_cod =tf.reduce_mean(loss_cod_list)
        loss_c = tf.reduce_mean(loss_c_list)
        loss_p_no = tf.reduce_mean(loss_p_no_list)
        # 각 Loss를 비중을 곱해 더해 최종 Loss를 계산
        loss = loss_p_rate * loss_p + loss_cod_rate * loss_cod + loss_c_rate * loss_c + loss_p_no_rate * loss_p_no
    # Loss에 대한 Grad를 구하고, 각 파라미터를 업데이트
    vars = yolo_model.trainable_variables
    grad = tape.gradient(loss, vars)
    opt.apply_gradients(zip(grad, vars))
    # 100번 마다 동영상에 이미지를 기록한다
    if(e%100==0):
        img = np.reshape(img,(256,256,3))
        label = pred.numpy()
        label = np.reshape(label,(3,3,8))
        sample_img = np.uint8(show_box(img,label))
        out.write(sample_img)
    print(e,"완료",loss.numpy())    
out.release()