[YOLO] 영상 객체 인식 - 번외: 원하는 객체만 검출

https://prlabhotelshoe.tistory.com/15 

[YOLO] 영상 객체 인식

이전 포스팅에서 OpenCV와 YOLO 포맷을 활용하여 영상 속 객체들을 인식하는 테스트를 진행해 보았습니다. 번외로 YOLO에서 학습된 class 중 원하는 객체만을 검출하는 테스트를 진행해 보겠습니다.

---

class 목록 확인

이전 다운로드 하였던 darknet 폴더에서 darknet-master  -> cfg 의 경로 속에 coco.names를 편집기 등을 통해 확인합니다. 여러가지를 활용해 볼 수 있겠지만 저의 경우 간단하게 VScode로 확인하였습니다.

VScode를 이용할 경우 다음과 같은 화면을 볼 수 있겠습니다. coco.names에는 1 person ~ 80 toothbrush 까지 총 80개의 class가 존재하는데, 여기서 필요한 class 번호를 확인합니다. 저의 경우 사람만 검출하기 위한 테스트를 진행해보기 위해 1 person을 활용할 것입니다.

---

소스 코드

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]
            #-- 원하는 class id 입력 / coco.names의 id에서 -1 할 것 
            if class_id == 0 and 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()

소스 코드의 경우 이전 포스팅과 큰 차이가 없으며, 변경할 점은 객체 인식 부분에 class_id가 원하는 id만을 추출하도록 추가하였습니다. (주석 참고) 주의할 점은 만약 1 person 만을 추출하고 싶을 경우 id에서 -1 을 한 0 값을 코드에 입력하도록 해야합니다.

---

테스트

수정 전

코드 수정 전의 출력 결과

이전 포스팅과 마찬가지의 출력을 볼 수 있습니다.

 

 

수정 후

코드 수정 후 원하는 class인 person 만을 출력하는 결과를 볼 수 있습니다.