Yolo PyTorch custom dataset 학습시키기 [4]

JunSeong An·2022년 6월 29일
0

AI

목록 보기
4/4
post-thumbnail

이번 포스팅에서는 이전 포스팅에서 만들었던 best.pt파일과 last.pt파일을 활용해 실제로 스크롤업 스크롤다운 되는 기능을 구현해보겠습니다.
비단 스크롤업과 스크롤다운 뿐만 아니라, 검출된 정보를 뽑아낼 수 있다면 여러분들이 구현하고자 하는 것대로 커스터마이징할 수 있을 것이라고 생각됩니다.

4. 실시간 검출 진행

실시간 검출을 진행한 후, 검출된 객체의 정보를 활용할 수 있는 방안까지 제시되어있습니다.

4-1 : simple detection

from matplotlib import scale
import torch
import numpy as np
import cv2

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

cap = cv2.VideoCapture(0)

frame_count = 0

while cap.isOpened():
    ret, frame = cap.read()
    
    # Make detections 
    results = model(frame)

	# 검출된 객체화면 출력
    cv2.imshow('YOLO', np.squeeze(results.render()))
    
cap.release()
cv2.destroyAllWindows()

학습시킨 .pt로 model 생성

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

우선 path로 만들어진 last.py(best.pt)의 위치를 넣어줍니다.

model에 frame 넣기

results = model(frame)

다음과 같이 불러온 모델에 captured된 model을 넣어주면 됩니다.


불러온 results에서 실제로 값을 뽑아내고, confidence를 뽑아내고, labels를 뽑아내는 것에 대한 내용은 이 레포지토리를 보시면 어떻게 활용할 지 이해할 수 있습니다.
scroll down and up을 구현한 코드는 여기있습니다.


이상 추가적으로 알게된 내용이 있으면 다시 돌아와 수정하고 배워가는 시간을 가지도록 하겠습니다. 감사합니다.

profile
안녕하세요!

0개의 댓글