YoloV5 개 고양이 분류

장준휴·2024년 1월 17일
0

딥러닝

목록 보기
2/3

1. 모델 import 받기


모델은 우리가 학습을 시킨 모델을 사용한다.

그러기 위해서는 사이트에서 직접 사진을 가지고 데이터 라벨링을 통해 학습을 시켜야한다.
그 과정은 나중에 추가로 설명할것이다.

일단 Yolo를 사용하는 라이브러리를 받아준다.

import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath

import torch # yolo 사용하는 라이브러리
import numpy as np
import cv2 # 이미지 다루는 라이브러리

그 이후 내가 학습시킨 모델을 사이트에서 다운받고 load한다

model = torch.hub.load('ultralytics/yolov5', 'custom', path='./data/best.pt' , force_reload=True)

이미지 검사


# 이미지 불러오기
dog = cv2.imread('./data/dogs.png')

# 객체 검출하기
result = model(dog)
# ctrl + shift + '-' > 커서 기준으로 셀 분할
# shift + m > 커서 기준으로 셀 분할

# 6개의 값
# x1, y1, x2, y2, 정확도(confidence), class
detect = result.xyxyn[0].numpy()

#기존이미지 크기 
x = dog.shape[1]
y = dog.shape[0]

#여러마리 검출
for i in range(len(detect)):
    # 객체의 위치 계산
    x1 = int(detect[i,0] * x)
    y1 = int(detect[i,1] * y)
    x2 = int(detect[i,2] * x)
    y2 = int(detect[i,3] * y)
    # 객체 검출했으면 이미지에 표시하기
    cv2.rectangle(dog, (x1,y1) , (x2,y2) , (0,0,255), 2)


#이미지 띄우기
cv2.imshow('dog',dog)
cv2.waitKey(0)
cv2.destroyAllWindows()
profile
나는야 토마토

0개의 댓글