[파이썬] 객체 검출을 위한 YOLO 알고리즘

YOLO (You Only Look Once) 알고리즘은 실시간 객체 검출 알고리즘으로, 이미지나 비디오에서 다양한 객체들을 식별하고 분류하는 데 사용됩니다. 이 알고리즘은 높은 정확도와 빠른 처리 속도를 제공하여 실시간 응용 프로그램에 이상적입니다.

YOLO 알고리즘 개요

YOLO 알고리즘은 다음과 같은 주요 단계로 이루어집니다:

  1. 입력 이미지를 그리드로 분할: 입력 이미지를 S x S 그리드로 나눕니다. 각 그리드 셀은 소위 “anchor box”에 대응됩니다.

  2. 그리드 셀 내에서 객체 감지: 각 그리드 셀은 출력 클래스 및 경계 상자의 확률을 예측합니다.

  3. 필터링 및 경계 상자 정확도 개선: 다중 클래스 예측 결과로 인한 상자의 중복을 제거하고, 경계 상자의 정확도를 개선합니다.

  4. 경계 상자 후보 선정: 경계 상자 후보를 선택하고, NMS (Non-Maximum Suppression) 알고리즘을 사용하여 겹치는 경계 상자를 제거합니다.

  5. 출력: 최종 객체 검출 결과를 출력합니다.

YOLO 알고리즘 구현 예제

아래의 예제 코드는 YOLO 알고리즘을 Python으로 구현한 것입니다. 이 예제는 OpenCVnumpy 라이브러리를 사용합니다. 다음은 예제 코드입니다:

import cv2
import numpy as np

# YOLO 모델 로드
net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")

# 클래스 이름 로드
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))

# 이미지 로드
img = cv2.imread("image.jpg")
height, width, channels = img.shape

# 이미지 전처리 및 객체 검출
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 > 0.5:
            # 객체의 경계 상자 좌표 계산
            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, 0.5, 0.4)

# 객체 검출 결과 출력
font = cv2.FONT_HERSHEY_SIMPLEX
for i in range(len(boxes)):
    if i in indexes:
        x, y, w, h = boxes[i]
        label = f"{classes[class_ids[i]]}: {confidences[i]:.2f}"
        color = colors[class_ids[i]]
        cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
        cv2.putText(img, label, (x, y - 10), font, 0.5, color, 2)

# 결과 출력
cv2.imshow("Output", img)
cv2.waitKey(0)
cv2.destroyAllWindows()

위의 예제 코드를 실행하면 yolov3.weights, yolov3.cfg, coco.names, image.jpg 파일이 필요합니다. yolov3.weightsyolov3.cfg는 사전 트레이닝된 YOLO 모델 가중치 및 설정 파일이고, coco.names은 클래스 이름 파일입니다. image.jpg는 객체 검출을 수행할 이미지입니다.

이제 위의 예제 코드를 사용하여 YOLO 알고리즘을 구현하고 객체를 검출할 수 있습니다.


참고: 예제 코드는 YOLO v3를 기반으로 작성되었습니다. 따라서 YOLO v3 모델 파일과 클래스 이름 파일을 제공해야 합니다.


YOLO 알고리즘은 실시간 객체 검출에 매우 효율적입니다. 따라서 자율 주행, 보안, 모션 감지 등 다양한 영역에서 사용되고 있습니다. YOLO 알고리즘을 사용하면 컴퓨터 비전과 인공 지능을 더욱 실용적으로 적용할 수 있습니다.