opencv_vs_torchvision_video.py

import torchvision
import os.path as osp
import numpy as np
import cv2
import torch
import time

def get_video_opencv(fname):

    vid_tensor = []

    vidcap = cv2.VideoCapture(fname)
    success, image = vidcap.read()
    # vid_tensor.append(image)
    while success:
        vid_tensor.append(image)
        success, image = vidcap.read()

    return torch.as_tensor(vid_tensor)

def get_video_torch(fname):
    vid_tensor, _, info = torchvision.io.read_video(fname)
    return vid_tensor

def get_frame_opencv(filevideo, t=0):
    cap = cv2.VideoCapture(filevideo)
    cap.set(propId=1, value=t)
    ret, frame = cap.read()
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    return frame

def get_frame_torch(filevideo, t=0):
    vid_tensor, _, info = torchvision.io.read_video(filevideo)
    return vid_tensor[t]

def main():
    vid_name = osp.join('data/surreal/val', 'ung_122_12', 'ung_122_12_c0002.mp4')

    iteration = 50

    et = []
    for i in range(iteration):
        start = time.time()
        vt = get_frame_opencv(vid_name)
        end = time.time()
        et.append(end-start)
    print(f'Avg execution time: {np.array(et).mean()} s.')


    et = []
    for i in range(iteration):
        start = time.time()
        vt = get_frame_torch(vid_name)
        end = time.time()
        et.append(end - start)
    print(f'Avg execution time: {np.array(et).mean()} s.')

if __name__ == '__main__':
    main()