Visualizations for Tkinter and PyGame in Colab

I love Colab. It’s a great replacement of a machine with GPU for ML/DL researchers. All you need is a Google account, and you’re good to go. You don’t need to buy a new system; GPUs on Colab get upgraded sometimes; install whatever library you like; you can share it with others or upload/update on Github (some researchers release codes on Colab). If you need space to store data or models, you can link it with your Google drive.

There are some issues with Colab but people come up with tricks for these and sometimes Google updates Colab to solve them. For example, if Colab doesn’t see activity on your notebook, it gets shut down to make resources available to others. This is sometimes annoying when you run a model that trains for a long time. You can find many solutions on Stack Overflow for this. They use javascript to keep clicking on a button to simulate interaction. An instance where Google took initiative was to add an inline image viewer function. Run cv2.imshow in your notebook next time, and it will suggest using cv2_imshow instead (notice the underscore) fromgoogle.colab.patches library.

One thing that kept me bugging was that Colab doesn’t have a default video viewer option. This is needed when I wish to run a reinforcement learning (RL) model. Many of the RL environments are built using PyGame (especially the custom ones) or Tkinter (e.g. UCB’s PacMan project used in their AI course). So I used the cv2_imshow function along with clear_output function from google.colab library.

PyGame

I like PyGame because it is quite easy to use and is documented very well. Also, if you learned Python from Rice University’s course on Coursera, like I did, you’ll be able to adapt to PyGame very quickly. I use PyGame often when I need to develop my own environment.

To make PyGame work on Colab,we need to install PyGame first by running the following in Colab:

!pip install pygame

UPDATE: the above is not needed since pygame is pre-installed on colab now.

Now import the required libraries as the following:

import os
os.environ["SDL_VIDEODRIVER"] = "dummy"
import pygame
from google.colab.patches import cv2_imshow
from google.colab import output

The first two lines create a dummy display environment for PyGame, so that it doesn’t throw an error looking for a display driver while running the game. this is based on a response by Ami F on Stack Overflow. The last two lines help with displaying and refreshing the images.

Now build your game the way you would have normally and extract the image from PyGame surface as the following:

image = pygame.surfarray.array3d(screen)
image = image.transpose([1, 0, 2])

For example, you can use the code from this post on RealPython by John Fincher. Update this code by commenting out or removing Line 20 and Line 26 and paste the lines above at Lines 26–27.

Here, screen is the Canvas object in PyGame, which is used to draw the objects. The second line is needed to switch rows and columns of the image as PyGame has a different height-width convention. If it doesn’t generate the right image for you, remove this line.

The last step is to display the image by running the following (inside the game loop):

output.clear(wait=True)
cv2_imshow(image)

The first line clears the output panel and the second line displays the image. The wait=True argument ensures that output is not cleared unless we have a new output to show. This results in a smooth transition of images. If you want to print something, make sure it is written after output.clear, else it would get cleared too before displaying the image.

Tkinter

We’re going to follow a similar recipe here; set a dummy display environment, extract the image, and display it. The snippets here are based on this notebook by Nextgrid.ai.

We will install some libraries first:

!apt-get install swig cmake libopenmpi-dev zlib1g-dev xvfb x11-utils ghostscript ffmpeg -qq #remove -qq for full output
!pip install pyvirtualdisplay ghostscript

These libraries help with the display environment. Ghostcript is an extra addition here to extract the images from Tkinter.

import os
os.system("Xvfb :1 -screen 0 1024x768x24 &")
os.environ['DISPLAY'] = ':1'
from tkinter import *
from google.colab.patches import cv2_imshow
from google.colab import output
from PIL import Image

Tkinter makes it a bit tricky to directly convert the canvas to an image, so we’ll save the image to a postscript file first, then read it with PIL. Rest is same as with PyGame:

canvas.postscript(file="file_name.eps", colormode='color')
image = Image.open('file_name.eps') 
output.clear(wait=True)
cv2_imshow(image)

After this, you should be able to see your game/video inline. If you get an error about opening the eps file with Image due to clipping, use the following in the last line instead:

cv2_imshow(np.asarray(image))

Here we are converting the image object to numpy array first, which is needed for clipping images while visualizing with cv2_imshow.

Things to Remember

• If your video is too fast, add a time.sleep after cv2_imshow with the appropriate sleep time.
• If you want to print some information (say game score) then prefer to print it below cv2_imshow . Otherwise ,the image may wobble if the printed text is not of teh same length through the iterations.
• You may need to rescale or transpose the image based on how you draw things on canvas. Since image is a numpy array, you can use opencv or numpy functions for these modification.

The tricks I mentioned above follow a simple pattern as mentioned earlier: set a dummy display environment, extract the image, and show the image. The same idea can be used with other visualization environments as well.

References:

• Response by Ami F on Stack Overflow: https://stackoverflow.com/a/54047012
• Notebook for Deep learning labs #5 by Nextgrid.ai: https://colab.research.google.com/drive/12osEZByXOlGy8J-MSpkl3faObhzPGIrB#scrollTo=hyyN-2qyK_T2
• Post on RealPython by John Fincher: PyGame: A Primer on Game Programming in Python