[tensorflow] Yolo를 이용한 물체 검출

전체 코드


📢 실행하기 전 아래 코드를 실행하여 opencv라이브러리 설치하기!!

pip install opencv-python

 

📑 물체를 분류하기 위한 클래스 부류 파일

coco.names.txt
0.00MB

Yolo로 물체 검출

import numpy as np
import cv2

classes = []
f = open('/content/coco.names.txt','r')
classes = [line.strip() for line in f.readlines()]
colors = np.random.uniform(0,255, size=(len(classes),3))

img=cv2.imread('/content/dog.jpg')
height, width,channels = img.shape
blob=cv2.dnn.blobFromImage(img, 1.0/256, (448,448), (0,0,0), swapRB=True, crop=False)

yolo_model=cv2.dnn.readNet('./yolov3.weights','/content/yolov3.cfg')
layer_names=yolo_model.getLayerNames()
out_layers=[layer_names[i[0]-1] for i in yolo_model.getUnconnectedOutLayers()]

yolo_model.setInput(blob)
output3=yolo_model.forward(out_layers)

class_ids,confidences,boxes=[],[],[]
for output in output3:
    for vec85 in output:
        scores=vec85[5:]
        class_id=np.argmax(scores)
        confidence=scores[class_id]
        if confidence>0.5:
            centerx, centery= int(vec85[0]*width), int(vec85[1]*height)
            w,h=int(vec85[2]*width), int(vec85[3]*height)
            x,y= int(centerx-w/2), int(centery-h/2)
            boxes.append([x,y,w,h])
            confidences.append(float(confidence))
            class_ids.append(class_id)

indexes=cv2.dnn.NMSBoxes(boxes,confidences,0.5,0.4)

for i in range(len(boxes)):
    if i in indexes:
        x,y,w,h=boxes[i]
        text=str(classes[class_ids[i]])+"%.3f"%confidences[i]
        cv2.rectangle(img,(x,y),(x+w,y+h), colors[class_ids[i]],2)
        cv2.putText(img,text,(x,y+30),cv2.FONT_HERSHEY_PLAIN,2,colors[class_ids[i]],2)


cv2.imshow('Object detection', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

(좌) 테스트 이미지 (우) 실행 결과

코드 해설


colors = np.random.uniform(0,255, size=(len(classes),3))

위는 각 클래스를 분류하기 위해서 난수로 색깔을 지정한 것이다. (나중에 박스 표시할 때 이 색깔로 구분, R,G,B)

img=cv2.imread('/content/dog.jpg')
height, width,channels = img.shape
blob=cv2.dnn.blobFromImage(img, 1.0/256, (448,448), (0,0,0), swapRB=True, crop=False)

테스트 영상을 읽고 blobFromImage함수를 이용해 전처리를 수행한다.

  • [0,255]범위의 화소를 [0,1]로 정규화
  • 448*448크기 영상으로 변환
  • BGR순서로 되어 있는 영상을 RGB 순으로 바꾼다.
yolo_model=cv2.dnn.readNet('./yolov3.weights','/content/yolov3.cfg')

yolo 모델을 불러온다.

이때 필요한 가중치와 cfg파일(신경망 구조 정보)은 다운 받아서 사용한다.

layer_names=yolo_model.getLayerNames()
out_layers=[layer_names[i[0]-1] for i in yolo_model.getUnconnectedOutLayers()]

신경망의 출력을 담당하는 3개의 층, yolo_82, yolo_94, yolo_106을 알아낸다.

yolo_model.setInput(blob)
output3=yolo_model.forward(out_layers)

테스트 영상을 입력하고 출력을 받아 output3객체에 저장한다.

output3 객체는 14*14*3*85 텐서, 28*28*3*85 텐서, 56*56*3*85 텐서를 리스트에 담고 있다.

85-차원의 벡터는 (x,y,w,h,o,p1,p2,⋯,p80)을 표현하며, 앞 4개는 바운딩 박스 정보, o는 물체일 가능성, 이후 80개는 물체 부류 확률

class_ids,confidences,boxes=[],[],[]
for output in output3:
    for vec85 in output:
        scores=vec85[5:]
        class_id=np.argmax(scores)
        confidence=scores[class_id]
        if confidence>0.5: #신뢰도 > 50%인 경우만 취함
            centerx, centery= int(vec85[0]*width), int(vec85[1]*height)
            				#[0,1]표현을 영상 크기로 변환
            w,h=int(vec85[2]*width), int(vec85[3]*height)
            x,y= int(centerx-w/2), int(centery-h/2)
            boxes.append([x,y,w,h])
            confidences.append(float(confidence))
            class_ids.append(class_id)

최고 부류 확률이 0.5를 넘는 바운딩 박스 정보를 모은다.

indexes=cv2.dnn.NMSBoxes(boxes,confidences,0.5,0.4)

비최대억제(NMS) 알고리즘을 적용해 주위 바운딩 박스에 비해 최대를 유지한 것만 남긴다.

for i in range(len(boxes)):
    if i in indexes:
        x,y,w,h=boxes[i]
        text=str(classes[class_ids[i]])+"%.3f"%confidences[i]
        cv2.rectangle(img,(x,y),(x+w,y+h), colors[class_ids[i]],2)
        cv2.putText(img,text,(x,y+30),cv2.FONT_HERSHEY_PLAIN,2,colors[class_ids[i]],2)

비최대 억제에서 살아남은 바운딩 박스를 부류 이름, 부류 확률 정보와 함께 영상에 표시한다.

cv2.imshow('Object detection', img)
cv2.waitKey(0)
cv2.destroyAllWindows()

차례대로 영상에 표시, 키보드의 크기가 눌릴 때까지 창 유지, 모든 창을 닫고 프로그램 종료

 

 


참고교재 파이썬으로 만드는 인공지능 - 한빛미디어