alexanderkroner / saliency

Contextual Encoder-Decoder Network for Visual Saliency Prediction [Neural Networks 2020]
MIT License
172 stars 46 forks source link

Python video processing #5

Closed iPsych closed 4 years ago

iPsych commented 4 years ago

Is it possible to run the code to video clip, not a through javascript but using Python? Does the video example use specialized model or adopt same pipeline to frame by frame data?

alexanderkroner commented 4 years ago

Hey, yes it's possible to run it on your computer in Python with OpenCV. Here's the code for using a webcam:

import cv2
import numpy as np
import tensorflow as tf

graph_def = tf.GraphDef()

with tf.gfile.Open("model_salicon_gpu.pb", "rb") as file:
    graph_def.ParseFromString(file.read())

input_plhd = tf.placeholder(tf.float32, (None, None, None, 3))

[predicted_maps] = tf.import_graph_def(graph_def,
                                       input_map={"input": input_plhd},
                                       return_elements=["output:0"])

with tf.Session() as sess:
    webcam = cv2.VideoCapture(0)

    while webcam.isOpened():
        _, webcam_img = webcam.read()
        webcam_img = cv2.resize(webcam_img, (320, 240))

        input_img = cv2.cvtColor(webcam_img, cv2.COLOR_BGR2RGB)
        input_img = input_img[np.newaxis, :, :, :]

        saliency = sess.run(predicted_maps,
                            feed_dict={input_plhd: input_img})

        saliency = cv2.cvtColor(saliency.squeeze(),
                                cv2.COLOR_GRAY2BGR)

        saliency = np.uint8(saliency * 255)

        display = cv2.hconcat([webcam_img, saliency])

        cv2.imshow("webcam saliency", display)

        if cv2.waitKey(1) == ord("q"):
            break

    webcam.release()
    cv2.destroyAllWindows()

All you need to do is download a pre-trained model from here and adjust the path. So to answer your second question, this approach makes use of the standard models and predicts saliency one frame at a time. I hope this helps!

iPsych commented 4 years ago

Thanks. It's perfect!