Move screenshotting code to its own module