인공기초 - 0508

기본 yolo 소스 코드

# 모듈 호출
import cv2
import numpy as np

# yolo 함수 생성
def yolo(frame, size, score_threshold, nms_threshold):
    net = cv2.dnn.readNet("yolov3.weights", "yolov3.cfg")   # 이미 학습된 가중치, 학습 파일을 불러다가 사용
    layer_names = net.getLayerNames()   # 레이어 구조 이름을 불러오는 명령어
    output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]   # 새로운 레이어 구조 생성
    colors = np.random.uniform(0, 255, size=(len(classes), 3))   # 랜덤 rgb배열을 생성
    height, width, channels = frame.shape   # 프레임 정보 가져오기

    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)
    print(f"boxes: {boxes}")
    print(f"confidences: {confidences}")
    indexes = cv2.dnn.NMSBoxes(boxes, confidences, score_threshold=score_threshold, nms_threshold=nms_threshold)   # 같은 물체 안에 중복된 박스 제거
    print(f"indexes: ", end='')

    for index in indexes:
        print(index, end=' ')
    print("\n\n============================== classes ==============================")   # 구분선

    # 출력문 제작
    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   # 결과 반환

# 물체의 이름 설정
classes = ["person", "bicycle", "car", "motorcycle",
           "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant",
           "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse",
           "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
           "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis",
           "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard",
           "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife",
           "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog",
           "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table",
           "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
           "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
           "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush"]

office = "office.jpg"   # 이미지 경로
frame = cv2.imread(office)   # 이미지 읽고 저장
size_list = [320, 416, 608]   # 이미지 사이즈 리스트 설정

frame = yolo(frame=frame, size=size_list[2], score_threshold=0.4, nms_threshold=0.4)   # yolo 함수 호출

# 이미지를 띄우기 위한 과정
cv2.imshow("Output_Yolo", frame)   # 적용된 이미지를 띄우기
cv2.waitKey(0)
cv2.destroyAllWindows()

 

 

 

2024.05.02 - [인공지능 기초] - 인공기초-0502

인공기초-0502