Last active
April 10, 2019 10:43
-
-
Save Puzer/60f262f80dea5125ec348bac5e35e68c to your computer and use it in GitHub Desktop.
Tensorflow ffmpeg cuda dali
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| 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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment