Yolo v5, v8, v10 가중치 파일(pt 파일)로 모델 객체 로딩하는 방법

Yolov5 모델 객체 로딩하기

 

Yolov5의 모델 객체 로딩은 이전 포스팅에서 다룬 Yolov7과 동일합니다.

 

아래 코드는 온라인 연결 없이 로컬에서만 로딩하는 방법입니다. Yolov5의 전체 코드를 다운받아서 저장해줘야합니다.

Confidence Threshold값은 따로 설정안하면 디폴트가 0.25입니다. 너무 낮으면 오탐이 많아지니 적당한 선에서 설정해주면 됩니다. 

import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# 가중치 파일 경로
path = r'.\yolov5\save_model\yolov5_medium_head_finetuned_16_30\heads_5m_e30_best.pt' 
model = torch.hub.load('ultralytics/yolov5', 'custom', path, force_reload=True)
model = model.to(device)
model.conf = 0.7 # confidence Threshold 값 설정

 

 

 

로컬에서만 로딩하고 싶으면 아래와 같이 사용하면 됩니다. yolov5의 전체 코드 파일을 깃허브에서 받아서 저장해줘야합니다. 

model = torch.hub.load("욜로5폴더경로", "custom", "가중치파일경로", source="local")

 

 

 

모델 사용법은 이전 포스팅을 참조 부탁드립니다.

https://codevang.tistory.com/356#google_vignette

객체인식 Yolo7(Yolov7) 파인튜닝하는 방법

 

 

 

 

Yolov8 모델 객체 로딩하기

 

버전 8부터는 이전과는 조금 다르게 로딩이 가능합니다. 좀 더 편해졌네요.

 

# v8 모델 영상 테스트
import torch
from ultralytics import YOLO

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
path = 욜로8 가중치파일.pt 경로

model = YOLO(path)
model = model.to(device)

 

 

 

Confidence Threshold 값은 모델에 이미지를 넣을 때(inference) 설정해주면 됩니다. 

 

역시 numpy 어레이로 변환된 이미지를 넣어주면 되고, yolo8부터는 "BGR" 타입으로 넣어줘야합니다. 5나 7같이 "RGB" 타입으로 넣으면 인식률이 떨어집니다.

 

import cv2

img_path = "이미지 경로"
img = cv2.imread(img_path)

result = self.model.predict(img, conf=0.7)

 

 

 

좌표값은 아래와 같이 불러줄 수 있습니다. 역시 이전 버전들과는 약간 구조가 다릅니다. 

 

p_center = result[0].boxes.xywh
p_xyxy = result[0].boxes.xyxy

 

 

 

넘파이 어레이기 때문에 특정 클래스 좌표만 불러오려면 아래와 같이 작성해주면 됩니다. 이전 버전과 다르게 클래스가 좌표와 함께 있지 않고 cls로 따로 분류돼 있습니다. 순서가 같으니 아래와 같이 걸러줄 수 있습니다. 

 

전 요리하기 편하게 바로 리스트로 변환해서 받았습니다. 

 

p_class = 0
cls = result[0].boxes.cls
p_center = p_center[cls==p_class].type(torch.int).tolist()
p_xyxy = p_xyxy[cls==p_class].type(torch.int).tolist()

 

 

 

Yolov10 모델 객체 로딩하기

 

10버전은 아래와 같이 로딩해줍니다.

 

from ultralytics import YOLOv10
import torch

# 최적 가중치 로드
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
best_model_path = yolo10가중치파일.pt 경로
model = YOLOv10(best_model_path)
model = model.to(device)

 

 

Confidence Threshold값 변경 및 결과 반환은 위의 8버전과 동일합니다. 역시 이미지는 꼭 "BGR" 타입으로 넣어줘야합니다.

 

여러 버전을 사용할 때는 아래와 같은 방식으로 코드를 짜주면 상태값 하나만 넣어서 안헷갈리게 쓸 수 있습니다.