IT

PyTorch: OpenCV와 YOLOv5 모델을 이용한 동영상 객체 탐지

유병혁 2021. 7. 27. 12:50

안녕하세요? 이번 글은 PyTorch에서 OpenCV와 YOLOv5 모델을 이용한 동영상 객체 탐지 과정을 정리해 보겠습니다.

이전 글(PyTorch에서 YOLOv5 모델을 이용한 이미지 객체 탐지) 글은 다음과 같습니다.

 

PyTorch: YOLOv5 모델을 이용한 이미지 객체 탐지

안녕하세요? 이번 글은 PyTorch에서 YOLOv5 모델을 이용한 이미지 객체 탐지 과정을 정리해 보겠습니다. *실습 내용은 아래 글에 포함되어 있는 소스 코드를 그대로 사용했습니다. PyTorch An open source m

blog.daum.net

관련해서 저는 Wells Fargo(웰스 파고)에서 근무하시는 머신러닝 엔지니어 '아카시 아그니호트리(Akash Agnihotri)'님의 아래 블로그 글과 깃허브 소스코드를 활용했습니다. {Thank you for sharing useful articles and source code!}

 

Implementing Real-time Object Detection System using PyTorch and OpenCV

Hands-On Guide to implement real-time object detection system using python

towardsdatascience.com

자, 그럼 시작해볼까요?! 일단 필요한 라이브러리를 호출합니다.

import torch
import numpy as np
import cv2
import pafy
from time import time

기존의 깃허브 소스코드를 활용하여, 객체탐지를 만들기 위한 ObjectDetection 클래스 함수를 정의합니다. ObjectDetect 클래스 함수에서 __init__ 함수는 객체가 생성(예: Video = ObjectDetection())될 때, __call__ 함수는 인스턴스가 생성(예: Video())될 때 호출됩니다. 그밖에 클래스 함수 내부에 정의된 함수는 하나하나 확인해 보겠습니다. 

class ObjectDetection:
    # YouTube 동영상에 YOLOv5 구현
    def __init__(self, url, out_file):
        # 객체 생성 시 호출
        # url: 예측 대상 YouTube URL
        # out_file: 유효한 출력 파일 이름 *.avi
        self._URL = url
        self.model = self.load_model()
        self.classes = self.model.names
        self.out_file = out_file
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
    def get_video_from_url(self):
        # url에서 새 비디오 스트리밍 객체 생성
        play = pafy.new(self._URL).streams[-1]
        assert play is not None
        return cv2.VideoCapture(play.url)
    def load_model(self):
        # YOLOv5 모델 로드
        model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
        return model
    def score_frame(self, frame):
        # frame: 단일 프레임; numpy/list/tuple 형식
        # return: 프레임에서 모델이 감지한 객체의 레이블과 좌표
        self.model.to(self.device)
        frame = [frame]
        results = self.model(frame)
        labels, cord = results.xyxyn[0][:, -1].cpu().numpy(), results.xyxyn[0][:, :-1].cpu().numpy()
        return labels, cord
    def class_to_label(self, x):
        # x 숫자 레이블 -> 문자열 레이블로 반환
        return self.classes[int(x)]
    def plot_boxes(self, results, frame):
        # 경계상자와 레이블을 프레임에 플로팅
        # results: 프레임에서 모델이 감지한 객체의 레이블과 좌표
        # frame: 점수화된 프레임
        # return: 경계 상자와 레이블이 플로팅된 프레임
        labels, cord = results
        n = len(labels)
        x_shape, y_shape = frame.shape[1], frame.shape[0]
        for i in range(n):
            row = cord[i]
            if row[4] >= 0.2:
                x1, y1, x2, y2 = int(row[0]*x_shape), int(row[1]*y_shape), int(row[2]*x_shape), int(row[3]*y_shape)
                bgr = (0, 255, 0)
                cv2.rectangle(frame, (x1, y1), (x2, y2), bgr, 2)
                cv2.putText(frame, self.class_to_label(labels[i])
                            + ': ' + str(x1) + ', ' + str(x2) + ', ' + str(y1) + ', ' + str(y2),
                            (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.9, bgr, 2)
        return frame
    def __call__(self):
        # 인스턴스 생성 시 호출; 프레임 단위로 비디오 로드
        player = self.get_video_from_url()
        assert player.isOpened()
        x_shape = int(player.get(cv2.CAP_PROP_FRAME_WIDTH))
        y_shape = int(player.get(cv2.CAP_PROP_FRAME_HEIGHT))
        four_cc = cv2.VideoWriter_fourcc(*"MJPG")
        out = cv2.VideoWriter(self.out_file, four_cc, 20, (x_shape, y_shape))
        while True:
            start_time = time()
            ret, frame = player.read()
            assert ret
            results = self.score_frame(frame)
            frame = self.plot_boxes(results, frame)
            end_time = time()
            fps = 1/np.round(end_time - start_time, 3)
            print(f"Frames Per Second : {fps}")
            out.write(frame)

__init__ 함수는 예측 대상 YouTube URL과 유효한 출력 파일 이름을 정의합니다.

class ObjectDetection:
    # YouTube 동영상에 YOLOv5 구현
    def __init__(self, url, out_file):
        # 객체 생성 시 호출
        # url: 예측 대상 YouTube URL
        # out_file: 유효한 출력 파일 이름 *.avi
        self._URL = url
        self.model = self.load_model()
        self.classes = self.model.names
        self.out_file = out_file
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

get_video_from_url 함수는 url에서 새 비디오 스트리밍 객체를 생성합니다.

    def get_video_from_url(self):
        # url에서 새 비디오 스트리밍 객체 생성
        play = pafy.new(self._URL).streams[-1]
        assert play is not None
        return cv2.VideoCapture(play.url)

load_model 함수는 PyTorch 허브에서 제공하는 모델(여기서는 YOLOv5)들을 로드합니다.

    def load_model(self):
        # YOLOv5 모델 로드
        model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
        return model

score_frame 함수에서는 프레임에서 모델이 감지한 객체의 레이블과 좌표를 반환합니다.

    def score_frame(self, frame):
        # frame: 단일 프레임; numpy/list/tuple 형식
        # return: 프레임에서 모델이 감지한 객체의 레이블과 좌표
        self.model.to(self.device)
        frame = [frame]
        results = self.model(frame)
        labels, cord = results.xyxyn[0][:, -1].cpu().numpy(), results.xyxyn[0][:, :-1].cpu().numpy()
        return labels, cord

class_to_label 함수는 간단한 역할인데요, 숫자 레이블을 문자열 레이블로 반환해 줍니다.

    def class_to_label(self, x):
        # x 숫자 레이블 -> 문자열 레이블로 반환
        return self.classes[int(x)]

plot_boxes 함수에서는 경계상자와 레이블을 프레임에서 플로팅해주는 역할을 담당합니다.

*프레임 내 감지 객체의 레이블과 좌표를 'label: x1, x2, y1, y2' 형태로 출력하는 방식입니다.

    def plot_boxes(self, results, frame):
        # 경계상자와 레이블을 프레임에 플로팅
        # results: 프레임에서 모델이 감지한 객체의 레이블과 좌표
        # frame: 점수화된 프레임
        # return: 경계 상자와 레이블이 플로팅된 프레임
        labels, cord = results
        n = len(labels)
        x_shape, y_shape = frame.shape[1], frame.shape[0]
        for i in range(n):
            row = cord[i]
            if row[4] >= 0.2:
                x1, y1, x2, y2 = int(row[0]*x_shape), int(row[1]*y_shape), int(row[2]*x_shape), int(row[3]*y_shape)
                bgr = (0, 255, 0)
                cv2.rectangle(frame, (x1, y1), (x2, y2), bgr, 2)
                cv2.putText(frame, self.class_to_label(labels[i])
                            + ': ' + str(x1) + ', ' + str(x2) + ', ' + str(y1) + ', ' + str(y2),
                            (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.9, bgr, 2)
        return frame

__call__ 함수는 인스턴스 생성 시 호출되며, 프레임 단위로 비디오를 로드합니다.

    def __call__(self):
        # 인스턴스 생성 시 호출; 프레임 단위로 비디오 로드
        player = self.get_video_from_url()
        assert player.isOpened()
        x_shape = int(player.get(cv2.CAP_PROP_FRAME_WIDTH))
        y_shape = int(player.get(cv2.CAP_PROP_FRAME_HEIGHT))
        four_cc = cv2.VideoWriter_fourcc(*"MJPG")
        out = cv2.VideoWriter(self.out_file, four_cc, 20, (x_shape, y_shape))
        while True:
            start_time = time()
            ret, frame = player.read()
            assert ret
            results = self.score_frame(frame)
            frame = self.plot_boxes(results, frame)
            end_time = time()
            fps = 1/np.round(end_time - start_time, 3)
            print(f"Frames Per Second : {fps}")
            out.write(frame)

자, 이제 실습을 위한 YouTube 동영상을 선택해 봅니다. 저는 유튜버 My Earth 님의 'Black Bear Bath Time' 동영상을 사용했습니다. {Your videos are really fantastic!}

자, 이제 YouTube URL과 출력 파일(*.avi)로 Video 객체 및 인스턴스를 생성합니다!

Video = ObjectDetection("https://www.youtube.com/watch?v=wCKYg3dP_M0", "Black_Bear_YOLOv5.avi")
Video()

결과는 다음과 같습니다! 이상으로 PyTorch에서 OpenCV와 YOLOv5 모델을 이용하여 동영상 객체를 탐지하는 과정을 정리해 봤습니