Using Neural Style Transfer on videos

sudo apt-get install libgl1-mesa-glx libegl1-mesa libxrandr2 libxrandr2 libxss1 libxcursor1 libxcomposite1 libasound2 libxi6 libxtst6;

source ~/anaconda3/bin/activate;
conda create --yes --name FastStyleTransfer python=3.9;
conda activate FastStyleTransfer;
pip install --upgrade pip;
pip install tensorflow matplotlib tensorflow_hub;

import matplotlib.pylab as plt
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub

from os import listdir
from os.path import isfile, join

import argparse

print("TF Version: ", tf.__version__)
print("TF Hub version: ", hub.__version__)
print("Eager mode enabled: ", tf.executing_eagerly())
print("GPU available: ", tf.config.list_physical_devices('GPU'))

# Parsing command line arguments while making sure they are mandatory/required
parser = argparse.ArgumentParser()
parser.add_argument(
    "--input",
    type=str,
    required=True,
    help="The directory that contains the input video frames.")
parser.add_argument(
    "--output",
    type=str,
    required=True,
    help="The directory that will contain the output video frames.")
parser.add_argument(
    "--style",
    type=str,
    required=True,
    help="The location of the style frame.")


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    
    args = parser.parse_args()
    input_path = args.input + '/'
    output_path = args.output + '/'
    # List all files from the input directory. This directory should contain at least one image/video frame.
    onlyfiles = [f for f in listdir(input_path) if isfile(join(input_path, f))]

    # Loading the input style image.
    style_image_path = args.style  # @param {type:"string"}
    style_image = plt.imread(style_image_path)

    # Convert to float32 numpy array, add batch dimension, and normalize to range [0, 1]. Example using numpy:
    style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.

    # Optionally resize the images. It is recommended that the style image is about
    # 256 pixels (this size was used when training the style transfer network).
    # The content image can be any size.
    style_image = tf.image.resize(style_image, (256, 256))
    
    # Load image stylization module.
    # Enable the following line and disable the next two to load the stylization module from a local folder.
    # hub_module = hub.load('magenta_arbitrary-image-stylization-v1-256_2')
    # Disable the above line and enable these two to load the stylization module from the internet.
    hub_handle = 'https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2'
    hub_module = hub.load(hub_handle)
 

    for inputfile in onlyfiles:
        content_image_path = input_path + inputfile  # @param {type:"string"}
        content_image = plt.imread(content_image_path)
        # Convert to float32 numpy array, add batch dimension, and normalize to range [0, 1]. Example using numpy:
        content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.

        # Stylize image.
        outputs = hub_module(tf.constant(content_image), tf.constant(style_image))
        stylized_image = outputs[0]

        # Saving stylized image to disk.
        content_outimage_path = output_path + inputfile  # @param {type:"string"}
        tf.keras.utils.save_img(content_outimage_path, stylized_image[0])

python3 faster.py --input "$input_frames_folder" --output "$output_frames_folder" --style "$style";

#!/bin/bash
#source ~/anaconda3/bin/activate;
#conda create --yes --name FastStyleTransfer python=3.9;
#pip install --upgrade pip;
#pip install tensorflow matplotlib tensorflow_hub;
#conda activate FastStyleTransfer;

source ~/anaconda3/bin/activate;
conda activate FastStyleTransfer;

input_videos="./input/videos/*";
input_styles="./input/styles/*";
input_frames="./input/frames";
input_audio="./input/audio";
output_frames="./output/frames";
output_videos="./output/videos";

# Loop on each video in the input folder.
for video in $input_videos;
do
  echo "$video";
  videoname=$(basename "$video");

  # Extract all frames from the video file and save them in a new folder using 8-digit numbers with zero padding in an incremental order.
  input_frames_folder="$input_frames/$videoname";
  mkdir -p "$input_frames_folder";
  ffmpeg -v quiet -i "$video" "$input_frames_folder/%08d.ppm";

  # Extract the audio file from the video to the format of an mp3. We will need this audio later to add it to the final product.
  input_audio_folder="$input_audio/$videoname";
  mkdir -p "$input_audio_folder";

  audio="";
  # Only VP8 or VP9 or AV1 video and Vorbis or Opus audio and WebVTT subtitles are supported for WebM.
  if [[ $videoname == *.webm ]]; then
    audio="$input_audio_folder/$videoname.ogg";
    ffmpeg -v quiet -i "$video" -vn -c:a libvorbis -y "$audio";  
  else
    audio="$input_audio_folder/$videoname.mp3";
    ffmpeg -v quiet -i "$video" -vn -c:a libmp3lame -y "$audio";
  fi

  # Retrieve the frame rate from the input video. We will need it to configure the final video later.
  frame_rate=`ffprobe -v 0 -of csv=p=0 -select_streams v:0 -show_entries stream=r_frame_rate "$video"`;

  # Loop on each image style from the input styles folder.
  for style in $input_styles;
  do
    echo "$style";
    stylename=$(basename "$style");
    output_frames_folder="$output_frames/$videoname/$stylename";
    mkdir -p "$output_frames_folder";

    # Stylize all frames using the input image and write all processed frames to the output folder.
    python3 faster.py --input "$input_frames_folder" --output "$output_frames_folder" --style "$style";

    # Combine all stylized video frames and the exported audio into a new video file.
    output_videos_folder="$output_videos/$videoname/$stylename";
    mkdir -p "$output_videos_folder";
    ffmpeg -v quiet -framerate "$frame_rate" -i "$output_frames_folder/%08d.ppm" -i "$audio" -pix_fmt yuv420p -acodec copy -y "$output_videos_folder/$videoname";
    
    rm -rf "$output_frames_folder";
  done
  rm -rf "$output_frames/$videoname";
  rm -rf "$input_frames_folder";
  rm -rf "$input_audio_folder";
done

source ~/anaconda3/bin/activate;
conda activate FastStyleTransfer;
./execute.sh;