반응형
# 나만의 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()
'경기도 인공지능 개발 과정 > Python' 카테고리의 다른 글
[AIFB] pandas 기초 전처리 (0) | 2022.08.08 |
---|---|
[Python] Image segmentation 실습 (0) | 2022.08.04 |
[python] 텐서플로 object_detection 실습 (0) | 2022.08.04 |
[Ptyhon] 딥러닝 활성화 함수, 가중치 정리 (0) | 2022.07.26 |
[Python] AIFB 강의(가입정보를 활용한 고객 데이터 분석) (0) | 2022.07.18 |