Final3

import skvideo.io
import skvideo.datasets
import tensorflow as tf
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
import cv2
import json
import numpy as np

# enabling eager execution for easier explanation
tf.enable_eager_execution()

model = ResNet50(weights='imagenet')

reader = skvideo.io.FFmpegReader(skvideo.datasets.bigbuckbunny(),
                             inputdict={},
                             outputdict={})
def gen_frames():
    for frame in reader.nextFrame():
        yield frame

dataset = tf.data.Dataset.from_generator(gen_frames, tf.int64)

def preprocess(frame):
    x = tf.image.resize_bilinear(frame, [224, 224])
    x = preprocess_input(x)
    return x, frame

dataset = dataset.batch(64).map(preprocess, 10).prefetch(1)

def predict():
    with tf.device("/gpu:0"):
        for frames, original in dataset:
            yield model.predict(frames.numpy()), original

dataset2 = tf.data.Dataset.from_generator(predict, (tf.float64, tf.int64))

def postprocess(output, original):
    # do some post processing
    return tf.argsort(output)[:3], original

dataset2 = dataset2.apply(tf.data.experimental.unbatch()).map(postprocess, 10)

with open("imagenet_class_index.json") as f:
    CLASS_INDEX = json.load(f)

writer = skvideo.io.FFmpegWriter("output.mp4")
for value, value2 in dataset2:
    indices = value.numpy()
    f = value2.numpy()
    f = np.ascontiguousarray(f, dtype=np.uint8)
    cv2.putText(f, CLASS_INDEX[str(indices[0])][1], (20, 20), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), lineType=cv2.LINE_AA)
    cv2.putText(f, CLASS_INDEX[str(indices[1])][1], (20, 100), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), lineType=cv2.LINE_AA)
    cv2.putText(f, CLASS_INDEX[str(indices[2])][1], (20, 180), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), lineType=cv2.LINE_AA)
    writer.writeFrame(f)
    print(CLASS_INDEX[str(indices[0])][1], CLASS_INDEX[str(indices[1])][1], CLASS_INDEX[str(indices[2])][1])
writer.close()