Tensorflow ffmpeg cuda dali

cv2_video.py

import cv2
from tqdm import tqdm
import numpy as np
import tensorflow as tf

VIDEO_PATH = 'ch03_20181228181500.avi'
MODEL = 'ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_2018_07_03/frozen_inference_graph.pb'
OUTPUT_TENSORS = ['num_detections:0', 'detection_boxes:0', 'detection_scores:0', 'detection_classes:0']
BATCH_SIZE = 15

with tf.gfile.GFile(MODEL, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

with tf.Graph().as_default() as graph:
    video_input = tf.placeholder(tf.uint8)
    tf.import_graph_def(graph_def, input_map={"image_tensor:0": video_input}, name='')
    detections = [graph.get_tensor_by_name(name) for name in OUTPUT_TENSORS]
    
sess = tf.Session(graph=graph)


# -----------------

video_capture = cv2.VideoCapture(VIDEO_PATH)
total_frames = int(video_capture.get(cv2.CAP_PROP_FRAME_COUNT))

batch = np.zeros((BATCH_SIZE, 1080, 1920, 3))
for frame_id in tqdm(range(total_frames), desc='Extracting frames'):
    for i in range(BATCH_SIZE):
        success, frame = video_capture.read()
        if success:
            batch[i] = frame
        else:
            continue
    
    r = sess.run(detections, {video_input:batch})
    
    if not success:
        break

dali_tf_video.py

import tensorflow as tf
import nvidia.dali.plugin.tf as dali_tf
import os
import numpy as np

from nvidia.dali.pipeline import Pipeline
import nvidia.dali.ops as ops
import nvidia.dali.types as types

class VideoPipe(Pipeline):
    def __init__(self, batch_size, num_threads, device_id, data, sequence_length):
        super(VideoPipe, self).__init__(batch_size, num_threads, device_id, seed=16)
        self.input = ops.VideoReader(device="gpu", filenames=data, sequence_length=sequence_length)


    def define_graph(self):
        output = self.input(name="Reader")
        return output
        
BATCH_SIZE = 1
video_files = 'sintel_trailer-720p.mp4'
sess = tf.Session()

daliop = dali_tf.DALIIterator()
pipe = VideoPipe(batch_size=BATCH_SIZE, num_threads=1, device_id=0, data=video_files, sequence_length=1)
input_op = daliop(pipe, shapes=[(BATCH_SIZE, 1, 720, 1280, 3)], dtypes = [tf.int32], batch_size=BATCH_SIZE)[0]

r = sess.run(input_op)
print(r.shape)

tf_io_video.py

import tensorflow as tf
import tensorflow_io.video as video_io
from tqdm import tqdm_notebook
tf.reset_default_graph()

VIDEO_PATH = 'ch03_20181228181500.avi'
MODEL = 'ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_2018_07_03/frozen_inference_graph.pb'
OUTPUT_TENSORS = ['num_detections:0', 'detection_boxes:0', 'detection_scores:0', 'detection_classes:0']
BATCH_SIZE = 15

def get_stream_op(path, batch_size):
    video_stream = video_io.VideoDataset(path)
    video_stream = video_stream.batch(batch_size)
    video_stream = video_stream.apply(tf.data.experimental.prefetch_to_device('/GPU:0'))
    video_stream_iterator = video_stream.make_one_shot_iterator()
    video_stream_iterator_next = video_stream_iterator.get_next()
    return video_stream_iterator_next

with tf.gfile.GFile(MODEL, "rb") as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())

with tf.Graph().as_default() as graph:
    video_input = get_stream_op(VIDEO_PATH, BATCH_SIZE)
    tf.import_graph_def(graph_def, input_map={"image_tensor:0": video_input}, name='')
    detections = [graph.get_tensor_by_name(name) for name in OUTPUT_TENSORS]
    
sess = tf.Session(graph=graph)

for _ in tqdm_notebook(range(100000)):
    r = sess.run(detections)