프로그래밍/파이썬 Python
[tensorflow] Yolo를 이용한 물체 검출
한별요
2022. 1. 22. 23:32
전체 코드
📢 실행하기 전 아래 코드를 실행하여 opencv라이브러리 설치하기!!
pip install opencv-python
📑 물체를 분류하기 위한 클래스 부류 파일
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()차례대로 영상에 표시, 키보드의 크기가 눌릴 때까지 창 유지, 모든 창을 닫고 프로그램 종료
참고교재 파이썬으로 만드는 인공지능 - 한빛미디어