Wednesday 18 September 2013

screenshots using Python under Linux

I have been working on porting some code to Linux and required a method to grab a screenshot of a window using python. I could just execute a command line application such as scrot, but I wanted to see if there was an easy way in native python.
Turns out it is fairly straight forward.

First I looked at taking a basic image of the whole screen, then tried to modify it to allow me to grab particular windows.

My first attempt seemed to work but the output image was skewed, this was using pygtk (gtk.gdk).

#!/usr/bin/python
'''
Modified from example at http://ubuntuforums.org/showthread.php?t=448160&p=2681009#post2681009
'''

import gtk.gdk
import os
import time
import Image

def screenGrab():
   root_win = gtk.gdk.get_default_root_window()
   size = root_win.get_size()
   pixbuf = gtk.gdk.Pixbuf(gtk.gdk.COLORSPACE_RGB,False,8,size[0],size[1])
   pixbuf = pixbuf.get_from_drawable(root_win,root_win.get_colormap(),0,0,0,0,size[0],size[1])
   if (pixbuf == None):
     return False
   else:
     width,height=size[0],size[1]
     return Image.fromstring("RGB",(width,height),pixbuf.get_pixels())
if __name__ == '__main__':

   screen=screenGrab()
   #screen.show()
   screen.save(os.getcwd()+'/'+str(int(time.time()))+'.png', 'PNG')


This appeared to be using gtk2, was there anyway to use gtk3?

While searching I found a tutorial on PyCairo and this had an example of taking a screenshot using Gdk and using a cairo Image Surface. This seemed to do what I needed.

#!/usr/bin/python

'''
Modified from
ZetCode PyCairo tutorial

This code example takes a screenshot.

Original author: Jan Bodnar
website: zetcode.com
'''

from gi.repository import Gdk
import cairo
import os
import time


def main():
   
    root_win = Gdk.get_default_root_window()

    width = root_win.get_width()

    height = root_win.get_height()   
   
    image_surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)               
    pixbuf = Gdk.pixbuf_get_from_window(root_win, 0, 0, width, height)
       
    cr = cairo.Context(image_surface)   
    Gdk.cairo_set_source_pixbuf(cr, pixbuf, 0, 0)    
    cr.paint()

    image_surface.write_to_png("screenshot-"+str(int(time.time()))+".png")
       
       
if __name__ == "__main__":   
    main()



Now to modify this to capture a particular window. I can get information on any window using xwininfo, included in this is the xid (window id).
If I have the id of the window I want to capture can use the function GdkX11.X11Window.foreign_new_for_display() to switch to this window instead of the root window before creating the pixbuf and image surface.

#!/usr/bin/python

'''
Modified from
ZetCode PyCairo tutorial

original author: Jan Bodnar
website: zetcode.com

This code example takes a screenshot of particular window, id in hex
supplied as first argument

'''

from gi.repository import Gdk
from gi.repository import GdkX11
import cairo
import os
import time
import sys

def main():
    winid = int(sys.argv[1],16)

    root_win = GdkX11.X11Display.get_default()
    win = GdkX11.X11Window.foreign_new_for_display(root_win,winid)

    width = win.get_width()
    height = win.get_height()   
   
    image_surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)               
    pixbuf = Gdk.pixbuf_get_from_window(win, 0, 0, width, height)
       
    cr = cairo.Context(image_surface)   
    Gdk.cairo_set_source_pixbuf(cr, pixbuf, 0, 0)    
    cr.paint()

    image_surface.write_to_png("screenshot-"+str(int(time.time()))+".png")
       
       
if __name__ == "__main__":   
    main()



This works well, however I do not want to have to rely on an external application to get the id. Can I supply a window name and do the same thing?

GdkX11.X11Window.foreign_new_for_display() as well as taking a window id as an argument also takes a window object, so I just need to find a way to identify which window. This was not successful, however if I can identify the window by the name I can just pass the id of that window to GdkX11.X11Window.foreign_new_for_display().
This worked and now I have a script that will take a window name and dump an image of it to a png file.

#!/usr/bin/python

'''
Expanded upon code from
ZetCode PyCairo tutorial

original author: Jan Bodnar
website: zetcode.com

This code example takes a screenshot of a particular window from name
passed as argument.

'''

from gi.repository import Gdk
from gi.repository import GdkX11
import Xlib
import Xlib.display
import cairo
import os
import time
import sys


def get_window(name):
    mydisplay = Xlib.display.Display()
    root_win = mydisplay.screen().root
    window_list = [root_win]

    while len(window_list) != 0:
        win = window_list.pop(0)
        #print win.get_wm_name()
        if win.get_wm_name() == name:
           return win.id
        children = win.query_tree().children
        if children != None:
            window_list += children

    print 'Unable to find window matching - %s\n' % name
    sys.exit(1)
    return None

def main():
   
    root_win = GdkX11.X11Display.get_default()
    browser_win = get_window(sys.argv[1])
    win = GdkX11.X11Window.foreign_new_for_display(root_win,browser_win)


    width = win.get_width()
    height = win.get_height()   
   
    image_surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)               
    pixbuf = Gdk.pixbuf_get_from_window(win, 0, 0, width, height)
       
    cr = cairo.Context(image_surface)   
    Gdk.cairo_set_source_pixbuf(cr, pixbuf, 0, 0)    
    cr.paint()

    image_surface.write_to_png("screenshot-"+str(int(time.time()))+".png")
       
       
if __name__ == "__main__":   

    main()