Real-Time Neural Style Transfer for ioQuake3 using TensorFlow Keras

ioq3nst.py

# Code adapted from the original article by Orhan G. Yalçın
# https://towardsdatascience.com/fast-neural-style-transfer-in-5-minutes-with-tensorflow-hub-magenta-110b60431dcc
# For use in the following article concerning Real-Time NST in ioQuake3:
# https://james-william-fletcher.medium.com/real-time-neural-style-transfer-in-quake3-71cd5f6c3e4
import tensorflow as tf
import tensorflow_hub as hub
import sys
import time
import os

inim = "/home/vfc/.q3a/baseq3/ioquake3_screenbuffer.jpg"
stim = "style.jpg"

def img_scaler(image, max_dim = 512):
  original_shape = tf.cast(tf.shape(image)[:-1], tf.float32)
  scale_ratio = max_dim / max(original_shape)
  new_shape = tf.cast(original_shape * scale_ratio, tf.int32)
  return tf.image.resize(image, new_shape)

def load_img(path_to_img):
  try:
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_image(img, channels=3)
    img = tf.image.convert_image_dtype(img, tf.float32)
    img = img_scaler(img)
    return img[tf.newaxis, :]
  except Exception:
    pass

hub_module = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/1')
style_image = load_img(stim)

while True:
    try:
      time.sleep(0.001)
      if os.path.isfile(inim):
        content_image = load_img(inim)
        stylized_image = hub_module(tf.constant(content_image), tf.constant(style_image))[0]
        tf.keras.preprocessing.image.save_img("newbuff.bmp", stylized_image[0])
    except Exception:
      pass