# Object_Detection_YouTube.pyimport torch
import numpy as np
import cv2
import pafy
from time import time
classObjectDetection:"""
Class implements Yolo5 model to make inferences on a youtube video using Opencv2.
"""def__init__(self, url, out_file="Labeled_Video.avi"):"""
Initializes the class with youtube url and output file.
:param url: Has to be as youtube URL,on which prediction is made.
:param out_file: A valid output file name.
"""
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'defget_video_from_url(self):"""
Creates a new video streaming object to extract video frame by frame to make prediction on.
:return: opencv2 video capture object, with lowest quality frame available for video.
"""
play = pafy.new(self._URL).streams[-1]assert play isnotNonereturn cv2.VideoCapture(play.url)defload_model(self):"""
Loads Yolo5 model from pytorch hub.
:return: Trained Pytorch model.
"""
model = torch.hub.load('ultralytics/yolov5','yolov5s', pretrained=True)return model
defscore_frame(self, frame):"""
Takes a single frame as input, and scores the frame using yolo5 model.
:param frame: input frame in numpy/list/tuple format.
:return: Labels and Coordinates of objects detected by model in the frame.
"""
self.model.to(self.device)
frame =[frame]
results = self.model(frame)
labels, cord = results.xyxyn[0][:,-1].numpy(), results.xyxyn[0][:,:-1].numpy()return labels, cord
defclass_to_label(self, x):"""
For a given label value, return corresponding string label.
:param x: numeric label
:return: corresponding string label
"""return self.classes[int(x)]defplot_boxes(self, results, frame):