[YOLO] 이미지 객체 인식

이전 포스팅에서 OpenCV를 활용해 yolo 포맷을 적용시켜 영상 속 객체 인식에 대해 진행해 보았습니다. 순서에 맞진 않겠지만 이번에는 이미지 속 객체 인식에 대해 테스트 해보겠습니다. 

---

소스코드

import cv2
import numpy as np

def yolo(frame, size, score_threshold, nms_threshold):
    #-- YOLO 포맷 및 클래스명 불러오기
    model_file = './darknet-master/build/darknet/x64/cfg/yolov3.weights' #-- 본인 개발 환경에 맞게 변경할 것
    config_file = './darknet/darknet-master/build/darknet/x64/cfg/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)
    layer_names = net.getLayerNames()
    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

    #-- 클래스의 갯수만큼 랜덤 RGB 배열을 생성
    colors = np.random.uniform(0, 255, size=(len(classes), 3))

    #-- 이미지의 높이, 너비, 채널 받아오기
    height, width, channels = frame.shape

    #-- yolo 포맷 적용 전의 전처리
    blob = cv2.dnn.blobFromImage(frame, 0.00392, (size, size), (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 > 0.1:
                #-- 탐지된 객체경계상자 표기
                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, score_threshold=score_threshold, nms_threshold=nms_threshold)

    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            class_name = classes[class_ids[i]]
            label = f"{class_name} {confidences[i]:.2f}"
            color = colors[class_ids[i]]
           
            cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
            cv2.rectangle(frame, (x - 1, y), (x + len(class_name) * 13 + 65, y - 25), color, -1)
            cv2.putText(frame, label, (x, y - 8), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 0), 2)
            
            #-- 인식한 객체의 정보 출력
            print(f"[{class_name}({i})] conf: {confidences[i]} / x: {x} / y: {y} / width: {w} / height: {h}")

    return frame

#-- 클래스(names파일) 오픈 / 본인 개발 환경에 맞게 변경할 것
classes = []
with open("./coco.names", "r") as f:
    classes = [line.strip() for line in f.readlines()]

#-- 사용할 이미지 경로
img = "/img.jpg" #-- 본인 개발 환경에 맞게 변경할 것

frame = cv2.imread(img) 

#-- yolo 포맷에 입력될 사이즈 리스트
size_list = [320, 416, 608]

frame = yolo(frame=frame, size=size_list[2], score_threshold=0.4, nms_threshold=0.4)
cv2.imshow("yolo_img_test", frame)
cv2.waitKey(0)
cv2.destroyAllWindows()

---

테스트

출력 예시

yolo의 샘플이미지에 적용시킨 결과

생각 이상으로 깔끔하게 인식한 결과를 볼 수 있습니다.

출력 예시

직접 촬영한 사진에 테스트를 진행해 보았습니다. 모니터를 tv로 인식하였고 그외 작은 객체나 식별하기 어려운 객체는 인식하지 못한 모습입니다.