Skip to content

Instantly share code, notes, and snippets.

@mrbid
Last active May 5, 2024 20:07
Show Gist options
  • Save mrbid/e400fdd9ceaa8b3d83ba67eecd39cdb7 to your computer and use it in GitHub Desktop.
Save mrbid/e400fdd9ceaa8b3d83ba67eecd39cdb7 to your computer and use it in GitHub Desktop.
Linux: Real-Time Neural Style Transfer on the current Active Window.
# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment