[객체탐지] 왕초보도 가능한 인공지능/딥러닝으로 동영상에서 객체 탐지하기(파이썬, pytorch, YOLO v5)

 

자동으로 이미지/동영상의 객체를 찾아주는
객체 탐지

 

 

객체탐지에서 성능이 좋아 많이 쓰이고 있는 기술(모델)이 YOLO입니다.

그 중 YOLOv5 모델의 코드를 공유합니다. 

 

환경 설정없이 코드를 테스트해보고 싶다면, 아래 파이토치 공식 홈페이지에서 가능합니다. 

 

https://pytorch.org/hub/ultralytics_yolov5/

 

---

 

동영상 탐지하기

.
앞의 코드를 활용하여,

동영상에서 사람을 불러와서 사람을 탐지하고, 결과를 저장하는 코드를 소개합니다.

cap = cv2.VideoCapture('input.mp4')

이 부분에 입력으로 사용할 동영상의 이름을 넣습니다.
그리고, 실행하면 결과로 'output,mp4'로 된 객체 탐지 결과 동영상을 얻을 수 있습니다.

 

import cv2
import torch

yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5l',
                            device='cuda:0' if torch.cuda.is_available() else 'cpu')  # 예측 모델
yolo_model.classes = [0]  # 예측 클래스 (사람)

cap = cv2.VideoCapture('intput.mp4')
ret, frame = cap.read()
out = cv2.VideoWriter("output.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 20.0, (frame.shape[1], frame.shape[0]))

while True:
    ret, frame = cap.read()
    if frame is None:
        break
    results = yolo_model(frame)
    results_refine = results.pandas().xyxy[0].values
    nms_human = len(results_refine)
    if nms_human > 0:
        for bbox in results_refine:
            start_point = (int(bbox[0]), int(bbox[1]))
            end_point = (int(bbox[2]), int(bbox[3]))

            frame = cv2.rectangle(frame, start_point, end_point, (255, 0, 0), 3)
    cv2.imshow("Video streaaming", frame)
    if cv2.waitKey(1) == ord("q"):
        break

    out.write(frame)
out.release()