Trong bài Phát hiện đối tượng trong video, chúng ta đã thực hành viết ứng dụng sử dụng mô hình Mask-RCNN có sẵn trên TensorFlow hub để phát hiện đối tượng trong một video. Tuy mô hình Mask-RCNN có độ chính xác cao nhưng nó có nhược điểm là tốc độ xử lý chậm, khó đáp ứng yêu cầu trong việc xử lý video thời gian thực. Do đó, trong bài này, chúng ta sẽ sử dụng mô hình YOLOv5 với tốc độ xử lý nhanh để ứng dụng trong bài toán xử lý video thời gian thực.
Trước hết, chúng ta tiến hành cài đặt thư viện pafy để xử lý video trực tiếp tử YouTube url (khi chúng ta sử dụng pip install pafy sẽ có lỗi dislike_count, do đó chúng ta cài đặt trực tiếp từ git):
|
1 2 |
!pip install git+https://github.com/Cupcakus/pafy !pip install youtube_dl |
Tiếp theo chúng ta tiến hành import các thư viện cần thiết sau:
|
1 2 3 4 5 |
import numpy as np import torch from torchvision import transforms from torchvision.utils import draw_bounding_boxes import cv2, pafy |
Hiện tại, TensorFlow hub chưa hỗ trợ mô hình YOLOv5, do đó trong bài này chúng ta sẽ sử dụng Torch thay vì Tensorflow. Ta thực hiện load YOLOv5 như sau:
|
1 2 |
# Load YOLOv5 model detected_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True) |
Tương tự như trong bài Phát hiện đối tượng trong video, chúng ta viết hàm objectDetector(image_array) với input là image_array có shape (H, W, 3) là các frame đọc bởi thư viện opencv:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
# image_array has shape (H, W, 3) def objectDetector(image_array): # Using the detected model to detect the object detected_results = detected_model(image_array) # Get the bounding boxes (xmin, ymin, xmax, ymax) and their labels bbox=detected_results.xyxy[0][:,:4] detected_labels = detected_results.pandas().xyxy[0].name.tolist() # transforms.ToTensor() convert numpy.ndarray (H x W x C) in the range [0, 255] # to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0] img_tensor=transforms.ToTensor()(image_array)*255 # Convert to uint8 img_tensor=img_tensor.type(torch.uint8) # draw bounding boxes on the input image img_bb=draw_bounding_boxes(img_tensor, bbox, width=3, labels=detected_labels,colors= (0,255,0)) # Convert from (C x H x W) to (H x W x C) numpy img_bb=np.transpose(img_bb,(1,2,0)).numpy() return img_bb |
– Torch YOLOv5 nhận đầu vào là URL, Filename, PIL, OpenCV, Numpy và PyTorch (các bạn tham khảo thêm cách sử dụng YOLOv5 trên Torch hub tại đây). Do đó, ta có thể sử dụng trực tiếp image_array mà không cần các bước xử lý dữ liệu thêm.
– Để hiển thị các bounding box, chúng ta sử dụng hàm torchvision.utils.draw_bounding_boxes(). Hàm này nhận đầu vào là image tensor với shape (C x H x W), kiểu dữ liệu là uint8 và giá trị trong khoảng [0, 255]. Do đó, với dữ liệu là numpy array, chúng ta cần sử dụng hàm transforms.ToTensor() để chuyển từ numpy array về Tensor với shape (C x H x W). Tuy nhiên, hàm này chuyển các giá trị về khoảng [0,0 , 1,0] do đó ta cần nhân lại với 255 để đảm bảo đúng giá trị input của hàm torchvision.utils.draw_bounding_boxes() (các bạn tham khảo thêm cách sử dụng hàm transforms.ToTensor() tại đây).
Cuối cùng ta sử dụng thư viện pafy, opencv-python để đọc video và phát hiện đối tượng trong từng frame với hàm objectDetector() như sau:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
# Working directory and data directory working_dir='drive/MyDrive/Deeplearning_Colab/ComputerVision/' data_dir=working_dir+'data/' # Load Youtube video from its url url = "https://www.youtube.com/watch?v=ORrrKXGx2SE&ab_channel=MuhammedJamnas" youtube_video = pafy.new(url).getbest(preftype="mp4") video = cv2.VideoCapture(youtube_video.url) #video = cv2.VideoCapture(0) #use this for webcam # Define the resolutions of saved video frame_width = int(video.get(3)) frame_height = int(video.get(4)) size = (frame_width, frame_height) # Define the codec and create VideoWriter object fourcc = cv2.VideoWriter_fourcc(*'DIVX') out = cv2.VideoWriter(data_dir+'Yolo_output.avi', fourcc, 20, size) count=0 # Process each frame of the video while True: read_success, frame = video.read() if(read_success): # Detect object in each frame objectDetected_frame = objectDetector(frame) # write the frame with detected boxes and labels out.write(objectDetected_frame) else: break # Release resources and destroy all windows video.release() out.release() cv2.destroyAllWindows() |
Sau khi kết thúc việc xử lý video, ta được kết quả thu về là video Yolo_output.avi như sau:
Như vậy chúng ta đã hoàn thành việc xây dựng một ứng dụng phát hiện đối tượng thời gian thực với mô hình YOLOv5. Các bạn có thể tham khảo Colab notebook của bài viết này trên trang Github của ITechSeeker tại đây.