Linux: Real-Time Neural Style Transfer on the current Active Window.

rtnst_linux.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

# Reddit poster "Jonno_FTW" suggested I use a pygame window to display images
# https://www.reddit.com/r/learnmachinelearning/comments/paqzqj/tutorial_realtime_neural_style_transfer_in_quake3/ha8zuo5?utm_source=share&utm_medium=web2x&context=3

# This tutorial was used to setup a pygame window
# https://www.geeksforgeeks.org/python-display-images-with-pygame/

# Grabbing screen captures on linux using python was obtained from this askubuntu answer
# https://askubuntu.com/questions/1011507/screenshot-of-an-active-application-using-python

# YouTube video of the script in action
# https://www.youtube.com/watch?v=MVwFTn0Og7U

# All you need to do is put the style image in the same directory as this python script
# which is currently set to load "style.jpg" via the style_image variable.

# Dependencies:
# python3 -m pip install tensorflow
# python3 -m pip install tensorflow_hub
# python3 -m pip install pygame

import tensorflow as tf
import tensorflow_hub as hub
import os
import pygame
import gi
gi.require_version('Gtk', '3.0')
gi.require_version('Gdk', '3.0')
from gi.repository import Gdk

style_image = "style.jpg"

pygame.init()
display_surface = pygame.display.set_mode((512, 512))
pygame.display.set_caption('NST')

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

inim = "/dev/shm/nstss.bmp"
hub_module = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/1')
style_image = load_img(style_image)
fps_clock = pygame.time.Clock()

while True:
    try:
      fps_clock.tick(60)
      screen = Gdk.get_default_root_window().get_screen()
      w = screen.get_active_window()
      pb = Gdk.pixbuf_get_from_window(w, *w.get_geometry())
      pb.savev(inim, "bmp", (), ())
      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("/dev/shm/nstbuff.bmp", stylized_image[0])
        image = pygame.image.load('/dev/shm/nstbuff.bmp')
        display_surface.blit(image, (0, 0))
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
        pygame.display.update() 
    except Exception:
      pass