대외활동

내일할거

iseop 2023. 2. 26. 21:36   인쇄용 버전
#
# Derived from: https://pysource.com/2019/06/27/yolo-object-detection-using-opencv-with-python/
#

import cv2
import numpy

# 1. YOLO 신경망 정의
# Weights 및 Config 파일 출처: https://pjreddie.com/darknet/yolo/
# COCO.names 파일 출처: https://github.com/pjreddie/darknet/blob/master/data/coco.names
net = cv2.dnn.readNet("yolov3-608.weights", "yolov3-608.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 - 1] for i in net.getUnconnectedOutLayers()]
# 클래스 갯수만큼 임의의 RGB 배열을 생성
colors = numpy.random.uniform(0, 255, size=(len(classes), 3))


# 2. 이미지 파일 입력
# 이미지 가져오기
img = cv2.imread("IMG.jpg")
height, width, channels = img.shape
# 이미지를 BLOB으로 변환
size = [320, 416, 608]
blob = cv2.dnn.blobFromImage(img, 0.00392, (size[2], size[2]), (0, 0, 0), True, crop=False)
# YOLO에 변환된 이미지 입력
net.setInput(blob)
# YOLO로부터 결과값 받기
outs = net.forward(output_layers)


# 3. 사물 탐지
# 각 사물은 사물의 종류(클래스), 확신도, 상자의 크기 및 위치로 표시함
confidence_threshold = 0.1
class_ids = []
confidences = []
boxes = []
for out in outs:
    for detection in out:
        scores = detection[5:]
        class_id = numpy.argmax(scores)
        confidence = scores[class_id]
        if confidence > confidence_threshold:
            # 탐지한 사물의 중앙 좌표를 찾음
            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)

# 겹쳐져 있는 상자 중 확신도가 가장 높은 상자를 선택함
# NMS: Non-Maximum Suppression
score_threshold = 0.4
nms_threshold = 0.4
indexes = cv2.dnn.NMSBoxes(boxes, confidences, score_threshold, nms_threshold)


# 4. 탐지된 사물에 대한 정보 표시
font = cv2.FONT_HERSHEY_PLAIN
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[i]
        cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
        cv2.putText(img, label, (x, y + 30), font, 3, color, 3)


# 5. 윈도우 출력
cv2.imshow("Image", cv2.resize(img, None, fx=0.3, fy=0.3))
cv2.waitKey(0)
cv2.destroyAllWindows()