digits-walkthrough

#!/usr/bin/python
# Copyright (c) 2014-2015, NVIDIA CORPORATION.  All rights reserved.

import os
import sys
import time
from urlparse import urlparse
import requests
import json
import argparse
import socket

from selenium import webdriver
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.common.action_chains import ActionChains

wait_time = 2
def wait(s=wait_time):
    time.sleep(s)

def get_page(driver, url):
    driver.get(url)
    wait()

def create_dataset(driver, name, folder):
    dropdown_elements = driver.find_elements_by_class_name('dropdown-toggle')
    dropdown_dataset = dropdown_elements[0]
    dropdown_dataset.click()
    wait()

    dropdown_menu = driver.find_elements_by_class_name('dropdown-menu')
    dropdown_menu_dataset = dropdown_menu[0]
    dataset_link = dropdown_menu_dataset.find_element_by_tag_name('a')
    dataset_link.click()
    wait()

    folder_train_tooltip = driver.find_element_by_name('folder_train_explanation')
    folder_train_tooltip.click()
    wait()

    folder_train = driver.find_element_by_name('folder_train')
    folder_train.send_keys(folder)
    wait()

    resize_channels_tooltip = driver.find_element_by_name('resize_channels_explanation')
    resize_channels_tooltip.click()
    wait()

    image_type = driver.find_element_by_name('resize_channels')
    for option in image_type.find_elements_by_tag_name('option'):
        if option.text == 'Grayscale':
            image_type.click()
            wait()
            option.click()
            break
    wait()

    resize_width_tooltip = driver.find_element_by_name('resize_dims_explanation')
    resize_width_tooltip.click()
    wait()

    resize_width = driver.find_element_by_name('resize_width')
    resize_width.clear()
    resize_width.send_keys('28')
    wait()

    resize_height = driver.find_element_by_name('resize_height')
    resize_height.clear()
    resize_height.send_keys('28')
    wait()

    dataset_name = driver.find_element_by_name('dataset_name')
    dataset_name.click()
    dataset_name.send_keys(name)
    wait()

    create_button = driver.find_element_by_name('create-dataset')
    create_button.click()

    job_url = driver.current_url.replace('datasets', 'jobs')
    status_url = job_url + '/status'
    done = False
    while not done:
        r = requests.get(status_url)
        status = json.loads(r.content)
        done = status['status'] == 'Done'
        wait()
    wait()

def create_model(driver, name, dataset_name, test_image):
    dropdown_elements = driver.find_elements_by_class_name('dropdown-toggle')
    dropdown_model = dropdown_elements[1]
    dropdown_model.click()
    wait()

    dropdown_menu = driver.find_elements_by_class_name('dropdown-menu')
    dropdown_menu_model = dropdown_menu[1]
    model_link = dropdown_menu_model.find_element_by_tag_name('a')
    model_link.click()
    # move to 0,0 so we don't accidentally select a hover element
    body = driver.find_element_by_css_selector('body')
    body.click()
    wait()

    dataset_tooltip = driver.find_element_by_name('dataset_explanation')
    dataset_tooltip.click()
    wait()

    datasets = driver.find_element_by_name('dataset')
    for option in datasets.find_elements_by_tag_name('option'):
        if option.text == dataset_name:
            option.click()
            break
    wait()

    standard_networks = driver.find_elements_by_name('standard_networks')
    lenet = standard_networks[0]
    lenet.click()
    wait()

    model_name = driver.find_element_by_name('model_name')
    model_name.click()
    model_name.send_keys(name)
    wait()

    create_button = driver.find_element_by_name('create-model')
    create_button.click()

    job_url = driver.current_url.replace('models', 'jobs')
    status_url = job_url + '/status'
    done = False
    while not done:
        r = requests.get(status_url)
        status = json.loads(r.content)
        done = status['status'] == 'Done'
        wait()
        #driver.refresh()

    # test image
    print 'Testing...'
    image_path = driver.find_element_by_name('image_url')
    image_path.send_keys(test_image)
    wait()

    test_button = driver.find_element_by_name('test-one')
    test_button.click()


def main(argv):
    parser = argparse.ArgumentParser(description='Run a Selenium demo of DiGiTS')

    ### Positional arguments

    parser.add_argument('mnist_image_folder',
            type=str,
            help='Path to the MNIST dataset folder')
    parser.add_argument('test_image',
            type=str,
            help='Image to test with')

    ### Optional arguments

    parser.add_argument('-p', '--port',
            type=int,
            default=80,
            help='Port the server is running on (default 80)')

    args = vars(parser.parse_args())

    home_page = 'http://0.0.0.0:%d/' % args['port']
    dataset_path = args['mnist_image_folder']
    dataset_name = 'MNIST Dataset'
    model_name = 'MNIST Model'
    test_image = args['test_image']

    r = requests.get(home_page)
    assert r.status_code == requests.codes.ok, 'page "%s" does not exist - are you looking on the wrong port?' % home_page

    # Start selenium driver.
    driver = webdriver.Firefox()
    print 'Firefox webdriver started.'
    mouse = webdriver.ActionChains(driver)

    try:
        driver.maximize_window()

        get_page(driver, home_page)

        print 'Creating dataset...'
        create_dataset(driver, dataset_name, dataset_path)

        get_page(driver, home_page)

        print 'Creating model...'
        create_model(driver, model_name, dataset_name, test_image)

        print 'Done.'
        
        #display an alert message
        get_page(driver, "javascript:alert('Completed Walkthrough!');void(0);")

        # wait forever
        while True:
            wait(60)

    except KeyboardInterrupt:
        pass
    finally:
        print 'Cleaning up.'
        try:
            driver.close()
        except socket.error:
            pass

if __name__ == '__main__':
    main(sys.argv)