[YOLO] 영상 객체 인식

컴퓨터 비전 혹은 머신러닝 등을 공부하면 한 번쯤 접하게 되는 yolo

OpenCV를 활용하여 yolo 포맷을 이용해 영상 속 객체 인식 테스트를 진행해보기로 하겠습니다.

추가적으로 OpenCV cuda를 빌드하여 GPU를 적용시켰으며 이에 대한 빌드는 하단 링크를 참조하면 되겠습니다.

https://prlabhotelshoe.tistory.com/24?category=1004669 

[CUDA&OpenCV] Python 환경 OpenCV GPU 연동 (CUDA backend in OpenCV) - 1. 파일 설치

---

아무래도 yolo의 weights와 cfg 파일 및 names 파일을 사용해야 하므로 직접 깃허브에서 다운로드하거나 모듈을 통해 불러오도록 해야 합니다. 저의 경우 간단히 IDE에서 구현하기 위해 깃허브에서 파일을 통째로 다운로드하여 사용하였습니다. 

 

https://github.com/pjreddie/darknet

GitHub - pjreddie/darknet: Convolutional Neural Networks

 

링크를 타고 들어가 Code -> Download ZIP을 눌러 압축 파일을 다운로드 합니다.

 

압축을 풀면 darknet이란 폴더가 있으며 darknet-master  -> cfg 폴더에 필요한 cfg 파일들이 있습니다.

이번 테스트의 경우 yolov3.cfg 파일을 사용하겠습니다.

 

필요한 names 파일 역시 동일 폴더에 존재합니다. 사용할 파일은 coco.names 입니다.

weights 파일의 경우 darknet-master -> build -> darknet -> x64 -> train 폴더에 위치해 있습니다.

이번 테스트의 경우 yolov3.weights 파일을 사용하겠습니다.

준비된 파일을 따로 복사하여 개별적으로 폴더에 정리하거나 소스코드의 경로에 직접 설정하여도 되겠습니다. 또한 웹캠을 이용하여도 무방하나 다양한 객체를 검출하고싶은 경우 짧은 길이의 영상을 따로 준비해두면 좋겠습다.

 

---

소스코드

import cv2
import numpy as np
import time # -- 프레임 계산을 위해 사용


vedio_path = './video.mp4' #-- 사용할 영상 경로
min_confidence = 0.5

def detectAndDisplay(frame):
    start_time = time.time()
    img = cv2.resize(frame, None, fx=0.8, fy=0.8)
    height, width, channels = img.shape
    #cv2.imshow("Original Image", img)

    #-- 창 크기 설정
    blob = cv2.dnn.blobFromImage(img, 0.00392, (416, 416), (0, 0, 0), True, crop=False)

    net.setInput(blob)
    outs = net.forward(output_layers)

    #-- 탐지한 객체의 클래스 예측 
    class_ids = []
    confidences = []
    boxes = []

    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > min_confidence:
                # 탐지한 객체 박싱
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
               
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)

                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)

    indexes = cv2.dnn.NMSBoxes(boxes, confidences, min_confidence, 0.4)
    font = cv2.FONT_HERSHEY_DUPLEX
    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            label = "{}: {:.2f}".format(classes[class_ids[i]], confidences[i]*100)
            print(i, label)
            color = colors[i] #-- 경계 상자 컬러 설정 / 단일 생상 사용시 (255,255,255)사용(B,G,R)
            cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
            cv2.putText(img, label, (x, y - 5), font, 1, color, 1)
    end_time = time.time()
    process_time = end_time - start_time
    print("=== A frame took {:.3f} seconds".format(process_time))
    cv2.imshow("YOLO test", img)
    
#-- yolo 포맷 및 클래스명 불러오기
model_file = './yolov3.weights' #-- 본인 개발 환경에 맞게 변경할 것
config_file = './yolov3.cfg' #-- 본인 개발 환경에 맞게 변경할 것
net = cv2.dnn.readNet(model_file, config_file)

#-- GPU 사용
#net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
#net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)

#-- 클래스(names파일) 오픈 / 본인 개발 환경에 맞게 변경할 것
classes = []
with open("./coco.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
colors = np.random.uniform(0, 255, size=(len(classes), 3))

#-- 비디오 활성화
cap = cv2.VideoCapture(vedio_path) #-- 웹캠 사용시 vedio_path를 0 으로 변경
if not cap.isOpened:
    print('--(!)Error opening video capture')
    exit(0)
while True:
    ret, frame = cap.read()
    if frame is None:
        print('--(!) No captured frame -- Break!')
        break
    detectAndDisplay(frame)
    #-- q 입력시 종료
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cv2.destroyAllWindows()

 

---

테스트

사과가 잔뜩 나오는 영상으로 테스트한 결과

GPU를 연동하여 빠른 인식과 검출을 볼 수 있으며, 엄청난 성능을 보이진 않지만 그럭저럭 인식하는 모습을 볼 수 있습니다.

 

 

길거리 촬영 영상으로 테스트한 결과

아무래도 위 테스트 결과처럼 color를 random으로 지정하니 너무 정신사나운듯해서 (0,255,0) 값으로 고정한 뒤 출력해 보았습니다. 이 역시도 GPU를 연동하였고 어느정도의 성능을 보여줍니다.