티스토리 뷰

반응형

yolo v8 버전을 이용해서 이미지와 동영상을 예측하는 기본으로 사용되는 파이썬 코드를 설명하고자 한다.

yolo v8 버전은 다음 그림처럼 분류, 물체탐지, 물체 상세분할 등 3가지 유형의 작업을 수행할 수 있다.

실제로 이미지를 예측하기 위해서 가상환경을 활성화하고 내가 자주 쓰는 코딩 에디터인 VS code를 다음 같이 실행한다.

 

1. 이미지 예측

다음 파이썬 스크립트가 이미지를 예측한 기본 구조이다.

여기서는 버스이미지를 yolo v8 버전에서 기본적으로 제공하는 yolo8n.pt모델을 이용해서 예측하는 경우이다. 

save=True 옵션이 예측한 결과를 저장하는 기능을 하고 이 옵션이 없으면 결과를 확인할 수 없다.

from ultralytics import YOLO

# 버스 이미지
img_path = 'bus.jpg'

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

예측한 결과는 실행파일이 있는 동일한 폴더 안 runs\detect\predict4 폴더에 저장된다.

여기서 숫자는 실행한 수만큼 숫자가 증가한다. 즉, 4번째 예측한 결과임을 뜻한다.

다음 왼쪽 그림을 예측한 결과가 오른쪽 그림이다.

2. 동영상 예측

동영상 예측은 이미지를 예측하는 코드와 큰 차이는 없다.

단지 차이는 입력 데이터가 이미지인지 동영상인지가 차이의 핵심이다.

 

from ultralytics import YOLO

mp4_data = 'wildlife1_fox_test_image.mp4'
mp4_data1 = '너구리.MOV'

def run_vid():
    model = YOLO('yolov8m.pt')
    results = model(mp4_data1, save=True)  

   
if __name__ == '__main__':
    run_vid()

3. 사용자 모델 기반의 예측

다음 파이썬 구문은 사용자가 직접 생산한 학습데이터를 이용해서 개발된 모델 best.pt를 이용해서

동영상 자료와 이미지 자료를 동시에 예측하는 파이썬 구문이다.

from ultralytics import YOLO

# 사용자 모델
best_model_path ='model/best.pt'

# 동영상 입력자료
mp4_data = 'wildlife1_fox_test_image.mp4'
mp4_data1 = '너구리.MOV'

# 이미지 입력자료
img_path = 'bus.jpg'

# 사용자 모델기반의 동영상 예측
def run_custom_vid():
    model = YOLO(best_model_path)  # load a pretrained YOLOv8n model    
    results = model(source=mp4_data, save=True)  # predict on an image
    print(results)

# 사용자 모델기반의 이미지 예측
def run_custom_img():            
    model = YOLO(best_model_path)
    results = model(img_path, save=True)
    print(results)
    
if __name__ == '__main__': 
    run_custom_vid()    
    run_custom_img()

 

 

 

 

 

 

반응형
댓글