yolo v8 모델을 이용해서 야생동물 식별 탐지하기 - (3) 심화 예측

이번 글은 이전 글에서 설명한 기본적인 예측( 2023.04.20 - [머신러닝 & 딥러닝] - yolo v8 모델을 이용해서 야생동물 식별탐지하기 - (2) 기본 예측 )에서 좀 더 사용자가 수정 가능한 형태로 yolo v8 버전을 이용해서

야생동물 탐지를 위해 학습데이터를 훈련하고 나온 최적의 탐지모델을 이용해서

야생동물 이미지와 동영상을 예측하는 파이썬 코드와 그 결과로 생성되는 데이터 내용을 설명하고자 한다.

 

가상환경을 활성화하고 코딩 에디터인 VS code를 다음같이 실행한다.

 

 

YOLO v8 버전에서 나온 결과를 사용자가 조정해서 원하는 방식으로 사용하기 위해서 이 글에서는 YOLO v8 버전을 실행한 결과를 설명하고자 한다.

 

다음 구문에서 예측하는 구문에서 results를 출력한 결과물 데이터는

LIST 행태로 나온다. 크게 9가지 유형의 정보를 포함하고 있다.

def run_img():            
    model = YOLO('yolov8n.pt')
    results = model(img_path, save=True)
    print(results)

출력 결과를 보면 boxes, masks, names, orig_img, orig_shape, path, probs, speed 등에 대한 정보를 알 수 있다.

이 정보 중에서 boxes, names, orig_img에 대한 정보를 좀 더 자세하게 살펴보기 위해서

for 문을 이용해서 세부 정보를 가져와야 한다.

 

boxes 세부정보는 다음 파이썬 코드로 알 수 있다.

def run_img():            
    model = YOLO('yolov8n.pt')
    results = model(img_path, save=True)
    # print(results)

    print(results[0].boxes)

위의 코드를 실행하여 boxes에 대한 출력 결과이다.

크게 7가지 정보를 알 수 있는데  맨 처음 boxes에 대한 자료만 보면 탐지된 물체를 둘러싸고 있는 바운딩 박스의 xmin, ymin, xmax, ymax,  신뢰도, 클래스 아이디를 알 수 있다.

이를 세부적으로 cls, conf, data, xyxy로 다시 구분하여 정보를 보여주고 있다.

 

1. 바운딩 박스 정보

이때 맨 위의 boxes에 있는 정보만 있으면 모든 정보를 알 수 있기 때문에 

다음 구문으로 실행하면 맨 위의 boxes의 정보만을 얻어낼 수 있다.

print(results[0].boxes.boxes)

아래 결과를 보면 results[0].boxes.data를 써야한다고 한다. 앞으로는 이 구문을 사용하기로 한다.

여기서 xmin, ymin, xmax, ymax와 맨 마지막 클래스 id는 int 형태로 사용해야 하고 신뢰도는 float 형태로 사용해야 한다. 

2. 이미지 정보

이미지에 대한 정보는 다음 구문으로 알 수 있고 데이터 형태는 ndarray 배열 형태이다. 

이 데이터를 이용하면 cv2 패키지의 imshow()함수와 matplotlib 패키지의 imshow()함수로 화면으로 출력할 수 있다.

 

img_array = results[0].orig_img

 

3. 예측결과 출력하기

바운딩 박스와 이미지 정보를 이용해서 사용자가 예측결과를 출력할 수 있다.

이렇게 하는 이유는 이미지에서 탐지된 물체의 개수나 종류를 추출해서 또 다른 분석에 활용할 수 있다.

 

img_array = results[0].orig_img

for i, bndbox in enumerate(bndboxs):

        xmin = int(bndbox[0])
        ymin = int(bndbox[1])
        xmax = int(bndbox[2])
        ymax = int(bndbox[3])
        conf = float(bndbox[4])
        class_id = int(bndbox[5])
        class_name = names[class_id]

        print(xmin, ymin, xmax, ymax)

        text = f"{class_name}-{round(conf, 2)}"       

        cv2.rectangle(img_array, (xmin, ymin), (xmax, ymax), (0,255,0), 2)        
        cv2.putText(img_array, text, (xmin, ymin-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)

결과를 보면 사용자가 지정하여 사각형 색상은 녹색으로 하고 글자는 노란색으로 출력된 것으로 알 수 있다.

이와 같이 사용자가 예측된 결과로 부터 정보를 쉽게 접근하여 사용할 수 있게 된다.

 

전체 코드는 다음과 같다.

from ultralytics import YOLO
import pandas as pd
import cv2
import matplotlib.pyplot as plt

img_path = 'bus.jpg'

def run_img_custom():

    model = YOLO('yolov8n.pt')
    results = model(img_path)
    print(results)

    bndboxs = results[0].boxes.data
    class_id = results[0].boxes.cls
    conf = results[0].boxes.conf

    img_array = results[0].orig_img
    
    names =  results[0].names 
    print(names)

    for i, bndbox in enumerate(bndboxs):

        xmin = int(bndbox[0])
        ymin = int(bndbox[1])
        xmax = int(bndbox[2])
        ymax = int(bndbox[3])
        conf = float(bndbox[4])
        class_id = int(bndbox[5])
        class_name = names[class_id]

        print(xmin, ymin, xmax, ymax)

        text = f"{class_name}-{round(conf, 2)}"       

        cv2.rectangle(img_array, (xmin, ymin), (xmax, ymax), (0,255,0), 2)        
        cv2.putText(img_array, text, (xmin, ymin-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 255), 2)
     
     
    cv2.imshow('animal detection',img_array) 
    cv2.waitKey(0)
    cv2.imwrite("test.png", img_array) 
    

if __name__ == '__main__':
    run_img_custom()