Add scripts for distillability modelling

get_screenshots.py

-#!/usr/bin/env python
-# Copyright 2014 The Chromium Authors. All rights reserved.
-# Use of this source code is governed by a BSD-style license that can be
-# found in the LICENSE file.
-
-import argparse
-import json
-import os
-import shutil
-import sys
-import time
-import urllib
-
-try:
-  from selenium import webdriver
-except:
-  print 'ERROR:'
-  print 'Couldn\'t import webdriver. Please run `sudo ./install-build-deps.sh`.'
-  sys.exit(1)
-
-self_dir = os.path.abspath(os.path.dirname(__file__))
-
-def addBuildtoolsToPath():
-  envPath = os.environ['PATH']
-  if not 'buildtools' in envPath:
-    os.environ['PATH'] = '%s/buildtools:%s' % (self_dir, envPath)
-
-def getDistillerUrl(u):
-  params = { 'url': u}
-  return "chrome-distiller://blah/?" + urllib.urlencode(params)
-
-def newDriver():
-  chromeOptions = webdriver.ChromeOptions()
-  chromeOptions.add_argument('--enable-dom-distiller')
-  driver = webdriver.Chrome(chrome_options=chromeOptions)
-  driver.set_window_size(1600, 5000)
-  driver.set_page_load_timeout(5)
-  return driver
-
-def main(argv):
-  parser = argparse.ArgumentParser()
-  parser.add_argument('--out', required=True)
-  parser.add_argument('urls', nargs='*')
-  parser.add_argument('--force', action='store_true')
-  parser.add_argument('--urls-file')
-  parser.add_argument('--restart', action='store_true')
-  options = parser.parse_args(argv)
-
-  outdir = options.out
-  if not options.restart:
-    if os.path.exists(outdir):
-      if not options.force:
-        print outdir + ' exists'
-        return 1
-      shutil.rmtree(outdir, ignore_errors=True)
-    os.makedirs(outdir)
-  else:
-    if not os.path.exists(outdir):
-      print outdir + ' doesn\'t exist'
-      return 1
-
-  addBuildtoolsToPath()
-
-  if options.urls:
-    files = options.urls
-  elif options.urls_file:
-    with open(options.urls_file) as u:
-      files = u.read().splitlines()
-  else:
-    print 'oh no'
-    return 1
-
-  driver = newDriver()
-  output = []
-  startIndex = 0
-  if options.restart:
-    prevfiles = [os.path.join(outdir, f) for f in os.listdir(outdir)]
-    prevfiles = [f for f in prevfiles if os.path.isfile(f) and os.path.splitext(f)[1] == '.info']
-    for f in prevfiles:
-      with open(f) as infofile:
-        info = json.load(infofile)
-        output.append(info)
-    startIndex = max([i['index'] for i in output]) + 1
-    print 'starting at ', startIndex
-
-  try:
-    for i, f in enumerate(files):
-      prefix = '%s/%d' % (outdir, i)
-      if i < startIndex:
-        continue
-      try:
-        ss = '%s.png' % prefix
-        dss = '%s-distilled.png' % prefix
-        driver.get(f)
-        time.sleep(0.1)
-        driver.save_screenshot(ss)
-        driver.get(getDistillerUrl(f))
-        time.sleep(0.1)
-        driver.save_screenshot(dss)
-        data = {
-          'index': i,
-          'url': f,
-          'screenshot': ss,
-          'distilled': dss,
-        }
-        output.append(data)
-        with open('%s.info' % prefix, 'w') as info:
-          json.dump(data, info)
-
-      except Exception as e:
-        driver.quit()
-        driver = newDriver()
-        pass
-
-  finally:
-    driver.quit()
-
-  with open('%s/index' % outdir, 'w') as index:
-    json.dump(output, index)
-  return 0
-
-if __name__ == '__main__':
-  sys.exit(main(sys.argv[1:]))
-