<commit_before>#!/usr/bin/env python

import sys
import os.path
import getopt
import xmlrpclib

ADMIN_PATH = os.path.normpath(os.path.dirname(__file__))

class SafeTransportWithCert(xmlrpclib.SafeTransport):
    # TODO this should be changed to a config file item
    # TODO add generation of admin certificates in deploy and document it
    __cert_file = os.path.join(ADMIN_PATH, "admin-cert.pem")
    __key_file = os.path.join(ADMIN_PATH, "admin-key.pem")

    def make_connection(self, host):
        host_with_cert = (host, {
            'key_file' : self.__key_file,
            'cert_file' : self.__cert_file
        })
        return xmlrpclib.SafeTransport.make_connection(self, host_with_cert)

def print_configs(keys, show_index):
    print
    i = 0
    for key, value, desc in keys:
        pr = "\n"
        if (show_index):
            pr += "[%02i] " % (i,)
        pr += "%s=%s\n%s" % (key, str(value), desc)
        print pr
        print "---"
        i+=1

def set_key(server, key, old_value, new_value):
    if (new_value == ''):
        raise ValueError("No value given, so not changing value")
    else:
        if type(old_value) is int: # convert the value to the right type
            try:
                new_value = int(new_value)
            except:
                raise ValueError("Not changing the value, because it used to be a number, but I could not change your input into a number.")
                return
        if type(old_value) is bool:
            try:
                new_value = bool(int(new_value))
            except:
                raise ValueError("Not changing the value. Could not convert the value to 1 or 0, True or False respectivly.")
                return

        print "Changing value for '%s' to '%s'..." % (key, new_value)
        server.ChangeConfig(key, new_value)
        print "Changed."

def print_usage():
    print "USAGE: ./config_client.py [--list] [--interactive] [--set KEY=VALUE] AM_HOST:PORT"
    print
    print "When setting boolean VALUEs please use 1 (True) or 0 (False)."
    print "This client will connect with the certificate in called 'admin-cert.pem' / 'admin-key.pem' residing in this script's folder."

DEFAULT_HOST_AND_PORT='127.0.0.1:8001'

if __name__ == '__main__':
    try:
        opts, args = getopt.getopt(sys.argv[1:], 'hlis:', ['help', 'list', 'interactive', 'set='])
        if opts == [] and args == []:
            raise ValueError("Please specify an option.")

        if len(args) == 0:
            print "Assuming default host %s" % (DEFAULT_HOST_AND_PORT,)
            host_name = DEFAULT_HOST_AND_PORT
        else:
            host_name = args[0]
            if host_name.find(':') == -1:
                raise ValueError("Please specify host _and_ port.")

        transport = SafeTransportWithCert()
        server = xmlrpclib.ServerProxy("https://%s/amconfig" % (host_name,), transport=transport)

        for option, opt_arg in opts:
            if option in ['-h', '--help']:
                print_usage()
                sys.exit(0)
            if option in ['-l', '--list']:
                keys = server.ListConfigKeys() # reload the keys every time
                print_configs(keys, False)
                sys.exit(0)
            if option in ['-i', '--interactive']:
                while True:
                    keys = server.ListConfigKeys() # reload the keys every time
                    print_configs(keys, True)
                    user_input = raw_input("Number of config item to change? ('q' for quit) ")
                    print
                    print '>> >> >>'
                    if user_input == 'q':
                        break
                    try:
                        item_index = (int(user_input))
                    except:
                        print 'I can not understand the number you have given me. Please speak a little clearer.'
                        continue

                    key, current_value, desc = keys[item_index]
                    if type(current_value) is bool:
                        new_value = raw_input("new value [0 for False, 1 for True]: ")
                    else:
                        new_value = raw_input("new value [%s]: " % (str(current_value),))

                    try:
                        set_key(server, key, current_value, new_value)
                    except Exception as e:
                        print str(e)
                sys.exit(0)

            if option in ['-s', '--set']:
                if not opt_arg:
                    raise ValueError("When using the --set option, please add an argument after")
                if opt_arg.find('=') == -1:
                    raise ValueError("When using the --set option, please use KEY=VALUE format after.")
                selected_key, new_value = opt_arg.split("=")

                # get the list from the server and find check if it is valid
                keys = server.ListConfigKeys()
                keys = [i for i in keys if i[0] == selected_key]
                if len(keys) != 1:
                    raise ValueError("Could not find the specified key.")
                set_key(server, selected_key, keys[0][1], new_value)
                sys.exit(0)
    except getopt.GetoptError as err: # getopt.GetoptError
        print "ERROR: %s" % (err,)
        print_usage()
        sys.exit(2)

    sys.exit(0)
<commit_msg>more and better option parsing in config client<commit_after>#!/usr/bin/env python

# TODO fix the order of things

import sys
import os.path
import getopt
import xmlrpclib

class SafeTransportWithCert(xmlrpclib.SafeTransport):
    # TODO add generation of admin certificates in deploy and document it

    cert_path = 'unset'
    key_path = 'unset'

    def make_connection(self, host):
        full_cert_path = os.path.abspath(os.path.expanduser(self.cert_path))
        full_key_path = os.path.abspath(os.path.expanduser(self.key_path))
        if not os.path.exists(full_cert_path):
            raise ValueError("The specified certificate does not exist (%s)" % (full_cert_path,))
        if not os.path.exists(full_key_path):
            raise ValueError("The specified key does not exist (%s)" % (full_key_path,))
        host_with_cert = (host, {
            'key_file' : full_key_path,
            'cert_file' : full_cert_path
        })
        return xmlrpclib.SafeTransport.make_connection(self, host_with_cert)

def print_configs(keys, show_index):
    print
    i = 0
    for key, value, desc in keys:
        pr = "\n"
        if (show_index):
            pr += "[%02i] " % (i,)
        pr += "%s=%s\n%s" % (key, str(value), desc)
        print pr
        print "---"
        i+=1

def set_key(server, key, old_value, new_value):
    if (new_value == ''):
        raise ValueError("No value given, so not changing value")
    else:
        if type(old_value) is int: # convert the value to the right type
            try:
                new_value = int(new_value)
            except:
                raise ValueError("Not changing the value, because it used to be a number, but I could not change your input into a number.")
                return
        if type(old_value) is bool:
            try:
                new_value = bool(int(new_value))
            except:
                raise ValueError("Not changing the value. Could not convert the value to 1 or 0, True or False respectivly.")
                return

        print "Changing value for '%s' to '%s'..." % (key, new_value)
        server.ChangeConfig(key, new_value)
        print "Changed."

DEFAULT_HOST_AND_PORT='127.0.0.1:8001'

if __name__ == '__main__':
    import optparse
    parser = optparse.OptionParser(usage = "usage: %prog [options] AM_HOST:PORT")
    parser.add_option("--list", action="store_true", help="Lists all available config items.")
    parser.add_option("--set", help="Sets the config item for the given key with the given value (KEY=VALUE).")
    parser.add_option("--interactive", action="store_true", help="Starts this client in an interactive shell mode.")
    parser.add_option("--cert", help="Specifies the certificate which is used to connect. (in PEM format, defaults to 'admin-cert.pem')", default="admin-cert.pem")
    parser.add_option("--key", help="Specifies the private key used to sign the messages sent. (in PEM format, defaults to 'admin-key.pem')", default="admin-key.pem")
    opts, args = parser.parse_args(sys.argv)

    if len(args) == 1:
        print "Assuming default host %s" % (DEFAULT_HOST_AND_PORT,)
        host_name = DEFAULT_HOST_AND_PORT
    else:
        host_name = args[0]
        if host_name.find(':') == -1:
            raise ValueError("Please specify host _and_ port.")

    transport = SafeTransportWithCert()
    server = xmlrpclib.ServerProxy("https://%s/amconfig" % (host_name,), transport=transport)

    if not (opts.list or opts.interactive or opts.set):
        parser.print_help()

    if opts.cert:
        transport.cert_path = opts.cert
    if opts.key:
        transport.key_path = opts.key

    if opts.list:
        keys = server.ListConfigKeys() # reload the keys every time
        print_configs(keys, False)

    if opts.set:
        if opts.set.find('=') == -1:
            raise ValueError("When using the --set option, please use KEY=VALUE format after.")
        selected_key, new_value = opts.set.split("=")

        # get the list from the server and find check if it is valid
        keys = server.ListConfigKeys()
        keys = [i for i in keys if i[0] == selected_key]
        if len(keys) != 1:
            raise ValueError("Could not find the specified key.")
        set_key(server, selected_key, keys[0][1], new_value)
        sys.exit(0)

    if opts.interactive:
        while True:
            keys = server.ListConfigKeys() # reload the keys every time
            print_configs(keys, True)
            user_input = raw_input("Number of config item to change? ('q' for quit) ")
            print
            print '>> >> >>'
            if user_input == 'q':
                break
            try:
                item_index = (int(user_input))
            except:
                print 'I can not understand the number you have given me. Please speak a little clearer.'
                continue

            key, current_value, desc = keys[item_index]
            if type(current_value) is bool:
                new_value = raw_input("new value [0 for False, 1 for True]: ")
            else:
                new_value = raw_input("new value [%s]: " % (str(current_value),))
            try:
                set_key(server, key, current_value, new_value)
            except Exception as e:
                print str(e)
        sys.exit(0)
    sys.exit(0)
<|endoftext|>


<commit_before># Software License Agreement (BSD License)
#
# Copyright (c) 2009-2014, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software in source and binary forms, with or
# without modification, are permitted provided that the following conditions
# are met:
#
#   Redistributions of source code must retain the above
#   copyright notice, this list of conditions and the
#   following disclaimer.
#
#   Redistributions in binary form must reproduce the above
#   copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other
#   materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Author: vic.iglesias@eucalyptus.com
import re
import copy
import time
from boto.ec2.regioninfo import RegionInfo
import boto
from concurrent.futures import ThreadPoolExecutor
import urllib2
import cookielib
from eutester import Eutester
from boto.ec2.elb.listener import Listener
from boto.ec2.elb.healthcheck import HealthCheck
from os.path import join, abspath

ELBRegionData = {
    'us-east-1': 'elasticloadbalancing.us-east-1.amazonaws.com',
    'us-west-1': 'elasticloadbalancing.us-west-1.amazonaws.com',
    'us-west-2': 'elasticloadbalancing.us-west-2.amazonaws.com',
    'eu-west-1': 'elasticloadbalancing.eu-west-1.amazonaws.com',
    'ap-northeast-1': 'elasticloadbalancing.ap-northeast-1.amazonaws.com',
    'ap-southeast-1': 'elasticloadbalancing.ap-southeast-1.amazonaws.com'}

class ELBops(Eutester):
    @Eutester.printinfo
    def __init__(self, host=None, credpath=None, endpoint=None, aws_access_key_id=None, aws_secret_access_key=None,
                 username="root", region=None, is_secure=False, path='/', port=80, boto_debug=0):
        """

        :param host:
        :param credpath:
        :param endpoint:
        :param aws_access_key_id:
        :param aws_secret_access_key:
        :param username:
        :param region:
        :param is_secure:
        :param path:
        :param port:
        :param boto_debug:
        """
        self.aws_access_key_id = aws_access_key_id
        self.aws_secret_access_key = aws_secret_access_key
        self.account_id = None
        self.user_id = None
        super(ELBops, self).__init__(credpath=credpath)

        self.setup_elb_connection(host=host,
                                  region=region,
                                  endpoint=endpoint,
                                  aws_access_key_id=self.aws_access_key_id,
                                  aws_secret_access_key=self.aws_secret_access_key,
                                  is_secure=is_secure,
                                  path=path,
                                  port=port,
                                  boto_debug=boto_debug)
        self.poll_count = 48
        self.username = username
        self.test_resources = {}
        self.setup_elb_resource_trackers()

    @Eutester.printinfo
    def setup_elb_connection(self, endpoint=None, aws_access_key_id=None, aws_secret_access_key=None, is_secure=True,
                             host=None,
                             region=None, path="/", port=443, boto_debug=0):
        """

        :param endpoint:
        :param aws_access_key_id:
        :param aws_secret_access_key:
        :param is_secure:
        :param host:
        :param region:
        :param path:
        :param port:
        :param boto_debug:
        :raise:
        """
        elb_region = RegionInfo()
        if region:
            self.debug("Check region: " + str(region))
            try:
                if not endpoint:
                    elb_region.endpoint = ELBRegionData[region]
                else:
                    elb_region.endpoint = endpoint
            except KeyError:
                raise Exception('Unknown region: %s' % region)
        else:
            elb_region.name = 'eucalyptus'
            if not host:
                if endpoint:
                    elb_region.endpoint = endpoint
                else:
                    elb_region.endpoint = self.get_elb_ip()
        connection_args = {'aws_access_key_id': aws_access_key_id,
                           'aws_secret_access_key': aws_secret_access_key,
                           'is_secure': is_secure,
                           'debug': boto_debug,
                           'port': port,
                           'path': path,
                           'region': elb_region}

        if re.search('2.6', boto.__version__):
            connection_args['validate_certs'] = False

        try:
            elb_connection_args = copy.copy(connection_args)
            elb_connection_args['path'] = path
            elb_connection_args['region'] = elb_region
            self.debug(
                "Attempting to create load balancer connection to " + elb_region.endpoint + ':' + str(port) + path)
            self.elb = boto.connect_elb(**elb_connection_args)
        except Exception, e:
            self.critical("Was unable to create elb connection because of exception: " + str(e))

    def setup_elb_resource_trackers(self):
        """
        Setup keys in the test_resources hash in order to track artifacts created
        """
        self.test_resources["load_balancers"] = []

    def get_elb_ip(self):
        """Parse the eucarc for the AWS_ELB_URL"""
        elb_url = self.parse_eucarc("AWS_ELB_URL")
        return elb_url.split("/")[2].split(":")[0]

    def get_elb_path(self):
        """Parse the eucarc for the AWS_ELB_URL"""
        elb_url = self.parse_eucarc("AWS_ELB_URL")
        elb_path = "/".join(elb_url.split("/")[3:])
        return elb_path

    def create_listener(self, load_balancer_port=80, protocol="HTTP", instance_port=80, load_balancer=None):
        self.debug(
            "Creating ELB Listener for protocol " + protocol + " and port " + str(load_balancer_port) + "->" + str(
                instance_port))
        listner = Listener(load_balancer=load_balancer,
                           protocol=protocol,
                           load_balancer_port=load_balancer_port,
                           instance_port=instance_port)
        return listner

    def create_healthcheck(self, target="HTTP:80/instance-name", interval=5, timeout=2, healthy_threshold=2,
                           unhealthy_threshold=10):
        self.debug("Creating healthcheck: " + target + " interval=" + str(interval) + " timeout=" + str(timeout) +
                   " healthy threshold=" + str(healthy_threshold) + " unhealthy threshold=" + str(unhealthy_threshold))
        healthcheck = HealthCheck(target=target,
                                  timeout=timeout,
                                  interval=interval,
                                  healthy_threshold=healthy_threshold,
                                  unhealthy_threshold=unhealthy_threshold)
        return healthcheck

    def generate_http_requests(self, url, count=100, worker_threads=20):
        self.debug("Generating {0} http requests against {1}".format(count, url))
        jar = cookielib.FileCookieJar("cookies")
        handler = urllib2.HTTPCookieProcessor(jar)
        opener = urllib2.build_opener(handler)
        urllib2.install_opener(opener)
        response_futures = []
        with ThreadPoolExecutor(max_workers=worker_threads) as executor:
            for _ in range(count):
                response_futures.append(executor.submit(urllib2.urlopen, url))

        responses = []
        for response in response_futures:
            http_response = response.result()
            try:
                http_error_code = http_response.getcode()
                if http_error_code == 200:
                    result = "{0}".format(http_response.read().rstrip())
                    self.debug("Request response: " + result)
                    responses.append(result)
                else:
                    raise Exception("Error code " + http_error_code + " found when sending " +
                                    str(worker_threads) + " concurrent requests to " + url)
            finally:
                http_response.close()
        return responses

    def register_lb_instances(self, name, instances, timeout=360, poll_count=15):
        inst_ids = [inst.id for inst in instances]
        self.debug("Registering instances {0} with lb {1}".format(inst_ids, name))
        self.elb.register_instances(name, inst_ids)
        poll_sleep = timeout / poll_count
        for _ in range(poll_count):
            self.debug("Checking instance health for {0}".format(inst_ids))
            inst_states = self.elb.describe_instance_health(name, instances=inst_ids)
            states = [state.state for state in inst_states]
            if not states or 'OutOfService' in states:
                time.sleep(poll_sleep)
            elif 'InService' in states:
                self.debug("Instances {0} for lb {1} are InService".format(inst_ids, name))
                return
            else:
                # This should never happen
                pass
        raise Exception("Instances {0} failed to enter InService state before timeout".format(inst_ids))

    def create_load_balancer(self, zones, name="test", load_balancer_port=80, instances=None):
        self.debug("Creating load balancer: " + name + " on port " + str(load_balancer_port))
        listener = self.create_listener(load_balancer_port=load_balancer_port)
        self.elb.create_load_balancer(name, zones=zones, listeners=[listener])

        healthcheck = self.create_healthcheck()
        self.elb.configure_health_check(name, healthcheck)

        if instances:
            self.register_instances(name, instances)

        ### Validate the creation of the load balancer
        lbs = self.elb.get_all_load_balancers(load_balancer_names=[name])
        if not "load_balancers" in self.test_resources:
            self.test_resources["load_balancers"] = []

        if len(lbs) == 1:
            self.test_resources["load_balancers"].append(lbs[0])
            return lbs[0]
        else:
            raise Exception("Unable to retrieve load balancer after creation")

    def delete_load_balancers(self, lbs, timeout=60):
        for lb in lbs:
            self.delete_load_balancer(lb, timeout)

    def delete_load_balancer(self, lb, timeout=60, poll_sleep=10):
        self.debug("Deleting Loadbalancer: {0}".format(lb.name))
        self.elb.delete_load_balancer(lb.name)
        poll_count = timeout / poll_sleep
        for _ in range(poll_count):
            lbs = self.elb.get_all_load_balancers(load_balancer_names=[lb.name])
            if lb in lbs:
                time.sleep(poll_sleep)
        if lb in self.test_resources["load_balancers"]:
            self.test_resources["load_balancers"].remove(lb)

    def create_app_cookie_stickiness_policy(self, name, lb_name, policy_name):
        self.debug("Create app cookie stickiness policy: " + str(policy_name))
        self.elb.create_app_cookie_stickiness_policy(name=name,
                                                     lb_name=lb_name,
                                                     policy_name=policy_name)

    def create_lb_cookie_stickiness_policy(self, cookie_expiration_period, lb_name, policy_name):
        self.debug("Create lb cookie stickiness policy: " + str(policy_name))
        self.elb.create_lb_cookie_stickiness_policy(cookie_expiration_period=cookie_expiration_period,
                                                    lb_name=lb_name,
                                                    policy_name=policy_name)

    def create_lb_policy(self, lb_name, policy_name, policy_type, policy_attributes):
        self.debug("Create lb policy: " + str(policy_name))
        self.elb.create_lb_policy(lb_name=lb_name,
                                  policy_name=policy_name,
                                  policy_type=policy_type,
                                  policy_attributes=policy_attributes)

    def set_lb_policy(self, lb_name, lb_port, policy_name=None):
        self.debug("Set policy " + str(policy_name) + " for " + lb_name)
        self.elb.set_lb_policies_of_listener(lb_name=lb_name,
                                             lb_port=lb_port,
                                             policies=policy_name)

    def delete_lb_policy(self, lb_name, policy_name):
        self.debug("Deleting lb policy " + str(policy_name) + " from " + str(lb_name))
        self.elb.delete_lb_policy(lb_name=lb_name,
                                  policy_name=policy_name)

    def describe_lb_policies(self, lb):
        lbs = self.elb.get_all_load_balancers(load_balancer_names=[lb])
        return lbs[0].policies

    def add_lb_listener(self, lb_name, listener):
        self.debug("adding listener")
        self.elb.create_load_balancer_listeners(name=lb_name, listeners=[listener])

    def remove_lb_listener(self, lb_name, port):
        self.debug("removing listener")
        self.elb.delete_load_balancer_listeners(name=lb_name, ports=[port])

    def add_server_cert(self, cert_name, cert_dir="./test_data", cert_file="ssl_server_certs_basics.crt",
                           key_file="ssl_server_certs_basics.pem"):
        cert_body = open(join(cert_dir, cert_file)).read()
        cert_key = open(join(cert_dir, key_file)).read()
        self.upload_server_cert(cert_name=cert_name, cert_body=cert_body, private_key=cert_key)<commit_msg>Change path to server cert dir<commit_after># Software License Agreement (BSD License)
#
# Copyright (c) 2009-2014, Eucalyptus Systems, Inc.
# All rights reserved.
#
# Redistribution and use of this software in source and binary forms, with or
# without modification, are permitted provided that the following conditions
# are met:
#
#   Redistributions of source code must retain the above
#   copyright notice, this list of conditions and the
#   following disclaimer.
#
#   Redistributions in binary form must reproduce the above
#   copyright notice, this list of conditions and the
#   following disclaimer in the documentation and/or other
#   materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
# Author: vic.iglesias@eucalyptus.com
import re
import copy
import time
from boto.ec2.regioninfo import RegionInfo
import boto
from concurrent.futures import ThreadPoolExecutor
import urllib2
import cookielib
from eutester import Eutester
from boto.ec2.elb.listener import Listener
from boto.ec2.elb.healthcheck import HealthCheck
from os.path import join, abspath

ELBRegionData = {
    'us-east-1': 'elasticloadbalancing.us-east-1.amazonaws.com',
    'us-west-1': 'elasticloadbalancing.us-west-1.amazonaws.com',
    'us-west-2': 'elasticloadbalancing.us-west-2.amazonaws.com',
    'eu-west-1': 'elasticloadbalancing.eu-west-1.amazonaws.com',
    'ap-northeast-1': 'elasticloadbalancing.ap-northeast-1.amazonaws.com',
    'ap-southeast-1': 'elasticloadbalancing.ap-southeast-1.amazonaws.com'}

class ELBops(Eutester):
    @Eutester.printinfo
    def __init__(self, host=None, credpath=None, endpoint=None, aws_access_key_id=None, aws_secret_access_key=None,
                 username="root", region=None, is_secure=False, path='/', port=80, boto_debug=0):
        """

        :param host:
        :param credpath:
        :param endpoint:
        :param aws_access_key_id:
        :param aws_secret_access_key:
        :param username:
        :param region:
        :param is_secure:
        :param path:
        :param port:
        :param boto_debug:
        """
        self.aws_access_key_id = aws_access_key_id
        self.aws_secret_access_key = aws_secret_access_key
        self.account_id = None
        self.user_id = None
        super(ELBops, self).__init__(credpath=credpath)

        self.setup_elb_connection(host=host,
                                  region=region,
                                  endpoint=endpoint,
                                  aws_access_key_id=self.aws_access_key_id,
                                  aws_secret_access_key=self.aws_secret_access_key,
                                  is_secure=is_secure,
                                  path=path,
                                  port=port,
                                  boto_debug=boto_debug)
        self.poll_count = 48
        self.username = username
        self.test_resources = {}
        self.setup_elb_resource_trackers()

    @Eutester.printinfo
    def setup_elb_connection(self, endpoint=None, aws_access_key_id=None, aws_secret_access_key=None, is_secure=True,
                             host=None,
                             region=None, path="/", port=443, boto_debug=0):
        """

        :param endpoint:
        :param aws_access_key_id:
        :param aws_secret_access_key:
        :param is_secure:
        :param host:
        :param region:
        :param path:
        :param port:
        :param boto_debug:
        :raise:
        """
        elb_region = RegionInfo()
        if region:
            self.debug("Check region: " + str(region))
            try:
                if not endpoint:
                    elb_region.endpoint = ELBRegionData[region]
                else:
                    elb_region.endpoint = endpoint
            except KeyError:
                raise Exception('Unknown region: %s' % region)
        else:
            elb_region.name = 'eucalyptus'
            if not host:
                if endpoint:
                    elb_region.endpoint = endpoint
                else:
                    elb_region.endpoint = self.get_elb_ip()
        connection_args = {'aws_access_key_id': aws_access_key_id,
                           'aws_secret_access_key': aws_secret_access_key,
                           'is_secure': is_secure,
                           'debug': boto_debug,
                           'port': port,
                           'path': path,
                           'region': elb_region}

        if re.search('2.6', boto.__version__):
            connection_args['validate_certs'] = False

        try:
            elb_connection_args = copy.copy(connection_args)
            elb_connection_args['path'] = path
            elb_connection_args['region'] = elb_region
            self.debug(
                "Attempting to create load balancer connection to " + elb_region.endpoint + ':' + str(port) + path)
            self.elb = boto.connect_elb(**elb_connection_args)
        except Exception, e:
            self.critical("Was unable to create elb connection because of exception: " + str(e))

    def setup_elb_resource_trackers(self):
        """
        Setup keys in the test_resources hash in order to track artifacts created
        """
        self.test_resources["load_balancers"] = []

    def get_elb_ip(self):
        """Parse the eucarc for the AWS_ELB_URL"""
        elb_url = self.parse_eucarc("AWS_ELB_URL")
        return elb_url.split("/")[2].split(":")[0]

    def get_elb_path(self):
        """Parse the eucarc for the AWS_ELB_URL"""
        elb_url = self.parse_eucarc("AWS_ELB_URL")
        elb_path = "/".join(elb_url.split("/")[3:])
        return elb_path

    def create_listener(self, load_balancer_port=80, protocol="HTTP", instance_port=80, load_balancer=None):
        self.debug(
            "Creating ELB Listener for protocol " + protocol + " and port " + str(load_balancer_port) + "->" + str(
                instance_port))
        listner = Listener(load_balancer=load_balancer,
                           protocol=protocol,
                           load_balancer_port=load_balancer_port,
                           instance_port=instance_port)
        return listner

    def create_healthcheck(self, target="HTTP:80/instance-name", interval=5, timeout=2, healthy_threshold=2,
                           unhealthy_threshold=10):
        self.debug("Creating healthcheck: " + target + " interval=" + str(interval) + " timeout=" + str(timeout) +
                   " healthy threshold=" + str(healthy_threshold) + " unhealthy threshold=" + str(unhealthy_threshold))
        healthcheck = HealthCheck(target=target,
                                  timeout=timeout,
                                  interval=interval,
                                  healthy_threshold=healthy_threshold,
                                  unhealthy_threshold=unhealthy_threshold)
        return healthcheck

    def generate_http_requests(self, url, count=100, worker_threads=20):
        self.debug("Generating {0} http requests against {1}".format(count, url))
        jar = cookielib.FileCookieJar("cookies")
        handler = urllib2.HTTPCookieProcessor(jar)
        opener = urllib2.build_opener(handler)
        urllib2.install_opener(opener)
        response_futures = []
        with ThreadPoolExecutor(max_workers=worker_threads) as executor:
            for _ in range(count):
                response_futures.append(executor.submit(urllib2.urlopen, url))

        responses = []
        for response in response_futures:
            http_response = response.result()
            try:
                http_error_code = http_response.getcode()
                if http_error_code == 200:
                    result = "{0}".format(http_response.read().rstrip())
                    self.debug("Request response: " + result)
                    responses.append(result)
                else:
                    raise Exception("Error code " + http_error_code + " found when sending " +
                                    str(worker_threads) + " concurrent requests to " + url)
            finally:
                http_response.close()
        return responses

    def register_lb_instances(self, name, instances, timeout=360, poll_count=15):
        inst_ids = [inst.id for inst in instances]
        self.debug("Registering instances {0} with lb {1}".format(inst_ids, name))
        self.elb.register_instances(name, inst_ids)
        poll_sleep = timeout / poll_count
        for _ in range(poll_count):
            self.debug("Checking instance health for {0}".format(inst_ids))
            inst_states = self.elb.describe_instance_health(name, instances=inst_ids)
            states = [state.state for state in inst_states]
            if not states or 'OutOfService' in states:
                time.sleep(poll_sleep)
            elif 'InService' in states:
                self.debug("Instances {0} for lb {1} are InService".format(inst_ids, name))
                return
            else:
                # This should never happen
                pass
        raise Exception("Instances {0} failed to enter InService state before timeout".format(inst_ids))

    def create_load_balancer(self, zones, name="test", load_balancer_port=80, instances=None):
        self.debug("Creating load balancer: " + name + " on port " + str(load_balancer_port))
        listener = self.create_listener(load_balancer_port=load_balancer_port)
        self.elb.create_load_balancer(name, zones=zones, listeners=[listener])

        healthcheck = self.create_healthcheck()
        self.elb.configure_health_check(name, healthcheck)

        if instances:
            self.register_instances(name, instances)

        ### Validate the creation of the load balancer
        lbs = self.elb.get_all_load_balancers(load_balancer_names=[name])
        if not "load_balancers" in self.test_resources:
            self.test_resources["load_balancers"] = []

        if len(lbs) == 1:
            self.test_resources["load_balancers"].append(lbs[0])
            return lbs[0]
        else:
            raise Exception("Unable to retrieve load balancer after creation")

    def delete_load_balancers(self, lbs, timeout=60):
        for lb in lbs:
            self.delete_load_balancer(lb, timeout)

    def delete_load_balancer(self, lb, timeout=60, poll_sleep=10):
        self.debug("Deleting Loadbalancer: {0}".format(lb.name))
        self.elb.delete_load_balancer(lb.name)
        poll_count = timeout / poll_sleep
        for _ in range(poll_count):
            lbs = self.elb.get_all_load_balancers(load_balancer_names=[lb.name])
            if lb in lbs:
                time.sleep(poll_sleep)
        if lb in self.test_resources["load_balancers"]:
            self.test_resources["load_balancers"].remove(lb)

    def create_app_cookie_stickiness_policy(self, name, lb_name, policy_name):
        self.debug("Create app cookie stickiness policy: " + str(policy_name))
        self.elb.create_app_cookie_stickiness_policy(name=name,
                                                     lb_name=lb_name,
                                                     policy_name=policy_name)

    def create_lb_cookie_stickiness_policy(self, cookie_expiration_period, lb_name, policy_name):
        self.debug("Create lb cookie stickiness policy: " + str(policy_name))
        self.elb.create_lb_cookie_stickiness_policy(cookie_expiration_period=cookie_expiration_period,
                                                    lb_name=lb_name,
                                                    policy_name=policy_name)

    def create_lb_policy(self, lb_name, policy_name, policy_type, policy_attributes):
        self.debug("Create lb policy: " + str(policy_name))
        self.elb.create_lb_policy(lb_name=lb_name,
                                  policy_name=policy_name,
                                  policy_type=policy_type,
                                  policy_attributes=policy_attributes)

    def set_lb_policy(self, lb_name, lb_port, policy_name=None):
        self.debug("Set policy " + str(policy_name) + " for " + lb_name)
        self.elb.set_lb_policies_of_listener(lb_name=lb_name,
                                             lb_port=lb_port,
                                             policies=policy_name)

    def delete_lb_policy(self, lb_name, policy_name):
        self.debug("Deleting lb policy " + str(policy_name) + " from " + str(lb_name))
        self.elb.delete_lb_policy(lb_name=lb_name,
                                  policy_name=policy_name)

    def describe_lb_policies(self, lb):
        lbs = self.elb.get_all_load_balancers(load_balancer_names=[lb])
        return lbs[0].policies

    def add_lb_listener(self, lb_name, listener):
        self.debug("adding listener")
        self.elb.create_load_balancer_listeners(name=lb_name, listeners=[listener])

    def remove_lb_listener(self, lb_name, port):
        self.debug("removing listener")
        self.elb.delete_load_balancer_listeners(name=lb_name, ports=[port])

    def add_server_cert(self, cert_name, cert_dir="./testcases/cloud_user/elb/test_data",
                        cert_file="ssl_server_certs_basics.crt",
                        key_file="ssl_server_certs_basics.pem"):
        cert_body = open(join(cert_dir, cert_file)).read()
        cert_key = open(join(cert_dir, key_file)).read()
        self.upload_server_cert(cert_name=cert_name, cert_body=cert_body, private_key=cert_key)<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# Copyright 2015 Carnegie Mellon University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d

from sklearn.cross_validation import KFold
from sklearn.metrics import accuracy_score

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
plt.style.use('bmh')

import os

import argparse

from scipy import arange

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--workDir', type=str, default='reps')
    parser.add_argument('--lfwPairs', type=str,
                        default=os.path.expanduser("~/openface/data/lfw/pairs.txt"))
    args = parser.parse_args()

    print("Loading embeddings.")
    fname = "{}/labels.csv".format(args.workDir)
    paths = pd.read_csv(fname, header=None).as_matrix()[:, 1]
    paths = map(os.path.basename, paths)  # Get the filename.
    # Remove the extension.
    paths = map(lambda path: os.path.splitext(path)[0], paths)
    fname = "{}/reps.csv".format(args.workDir)
    rawEmbeddings = pd.read_csv(fname, header=None).as_matrix()
    embeddings = dict(zip(*[paths, rawEmbeddings]))

    pairs = loadPairs(args.lfwPairs)
    classifyExp(args.workDir, pairs, embeddings)
    plotClassifyExp(args.workDir)

def loadPairs(pairsFname):
    print("  + Reading pairs.")
    pairs = []
    with open(pairsFname, 'r') as f:
        for line in f.readlines()[1:]:
            pair = line.strip().split()
            pairs.append(pair)
    assert(len(pairs) == 6000)
    return np.array(pairs)

def getEmbeddings(pair, embeddings):
    if len(pair) == 3:
        name1 = "{}_{}".format(pair[0], pair[1].zfill(4))
        name2 = "{}_{}".format(pair[0], pair[2].zfill(4))
        actual_same = True
    elif len(pair) == 4:
        name1 = "{}_{}".format(pair[0], pair[1].zfill(4))
        name2 = "{}_{}".format(pair[2], pair[3].zfill(4))
        actual_same = False
    else:
        raise Exception(
            "Unexpected pair length: {}".format(len(pair)))

    (x1, x2) = (embeddings[name1], embeddings[name2])
    return (x1, x2, actual_same)

def writeROC(fname, thresholds, embeddings, pairsTest):
    with open(fname, "w") as f:
        f.write("threshold,tp,tn,fp,fn,tpr,fpr\n")
        tp = tn = fp = fn = 0
        for threshold in thresholds:
            tp = tn = fp = fn = 0
            for pair in pairsTest:
                (x1, x2, actual_same) = getEmbeddings(pair, embeddings)
                diff = x1 - x2
                dist = np.dot(diff.T, diff)
                predict_same = dist < threshold

                if predict_same and actual_same:
                    tp += 1
                elif predict_same and not actual_same:
                    fp += 1
                elif not predict_same and not actual_same:
                    tn += 1
                elif not predict_same and actual_same:
                    fn += 1

            if tp + fn == 0:
                tpr = 0
            else:
                tpr = float(tp) / float(tp + fn)
            if fp + tn == 0:
                fpr = 0
            else:
                fpr = float(fp) / float(fp + tn)
            f.write(",".join([str(x)
                              for x in [threshold, tp, tn, fp, fn, tpr, fpr]]))
            f.write("\n")
            if tpr == 1.0 and fpr == 1.0:
                # No further improvements.
                f.write(",".join([str(x)
                                  for x in [4.0, tp, tn, fp, fn, tpr, fpr]]))
                return

def evalThresholdAccuracy(embeddings, pairs, threshold):
    y_true = []
    y_predict = []
    for pair in pairs:
        (x1, x2, actual_same) = getEmbeddings(pair, embeddings)
        diff = x1 - x2
        dist = np.dot(diff.T, diff)
        predict_same = dist < threshold
        y_predict.append(predict_same)
        y_true.append(actual_same)

    y_true = np.array(y_true)
    y_predict = np.array(y_predict)
    accuracy = accuracy_score(y_true, y_predict)
    return accuracy

def findBestThreshold(thresholds, embeddings, pairsTrain):
    bestThresh = bestThreshAcc = 0
    for threshold in thresholds:
        accuracy = evalThresholdAccuracy(embeddings, pairsTrain, threshold)
        if accuracy >= bestThreshAcc:
            bestThreshAcc = accuracy
            bestThresh = threshold
        else:
            # No further improvements.
            return bestThresh
    return bestThresh

def classifyExp(workDir, pairs, embeddings):
    print("  + Computing accuracy.")
    folds = KFold(n=6000, n_folds=10, shuffle=False)
    thresholds = arange(0, 4, 0.01)

    if os.path.exists("{}/accuracies.txt".format(workDir)):
        print("{}/accuracies.txt already exists. Skipping processing.".format(workDir))
    else:
        accuracies = []
        with open("{}/accuracies.txt".format(workDir), "w") as f:
            f.write('fold, threshold, accuracy\n')
            for idx, (train, test) in enumerate(folds):
                fname = "{}/l2-roc.fold-{}.csv".format(workDir, idx)
                writeROC(fname, thresholds, embeddings, pairs[test])

                bestThresh = findBestThreshold(
                    thresholds, embeddings, pairs[train])
                accuracy = evalThresholdAccuracy(
                    embeddings, pairs[test], bestThresh)
                accuracies.append(accuracy)
                f.write('{}, {:0.2f}, {:0.2f}\n'.format(
                    idx, bestThresh, accuracy))
            avg = np.mean(accuracies)
            std = np.std(accuracies)
            f.write('\navg, {:0.4f} +/- {:0.4f}\n'.format(avg, std))
            print('    + {:0.4f}'.format(avg))

def getAUC(fprs, tprs):
    sortedFprs, sortedTprs = zip(*sorted(zip(*(fprs, tprs))))
    sortedFprs = list(sortedFprs)
    sortedTprs = list(sortedTprs)
    if sortedFprs[-1] != 1.0:
        sortedFprs.append(1.0)
        sortedTprs.append(sortedTprs[-1])
    return np.trapz(sortedTprs, sortedFprs)

def plotOpenFaceROC(workDir, plotFolds=True, color=None):
    fs = []
    for i in range(10):
        rocData = pd.read_csv("{}/l2-roc.fold-{}.csv".format(workDir, i))
        fs.append(interp1d(rocData['fpr'], rocData['tpr']))
        x = np.linspace(0, 1, 1000)
        if plotFolds:
            foldPlot, = plt.plot(x, fs[-1](x), color='grey', alpha=0.5)
        else:
            foldPlot = None

    fprs = []
    tprs = []
    for fpr in np.linspace(0, 1, 1000):
        tpr = 0.0
        for f in fs:
            v = f(fpr)
            if math.isnan(v):
                v = 0.0
            tpr += v
        tpr /= 10.0
        fprs.append(fpr)
        tprs.append(tpr)
    if color:
        meanPlot, = plt.plot(fprs, tprs, color=color)
    else:
        meanPlot, = plt.plot(fprs, tprs)
    AUC = getAUC(fprs, tprs)
    return foldPlot, meanPlot, AUC

def plotClassifyExp(workDir):
    print("Plotting.")

    fig, ax = plt.subplots(1, 1)

    openbrData = pd.read_csv("comparisons/openbr.v1.1.0.DET.csv")
    openbrData['Y'] = 1 - openbrData['Y']
    # brPlot = openbrData.plot(x='X', y='Y', legend=True, ax=ax)
    brPlot, = plt.plot(openbrData['X'], openbrData['Y'])
    brAUC = getAUC(openbrData['X'], openbrData['Y'])

    foldPlot_v1, meanPlot_v1, AUC_v1 = plotOpenFaceROC("lfw.nn4.v1", False)
    foldPlot_v2, meanPlot_v2, AUC_v2 = plotOpenFaceROC(workDir, color='k')

    humanData = pd.read_table(
        "comparisons/kumar_human_crop.txt", header=None, sep=' ')
    humanPlot, = plt.plot(humanData[1], humanData[0])
    humanAUC = getAUC(humanData[1], humanData[0])

    deepfaceData = pd.read_table(
        "comparisons/deepface_ensemble.txt", header=None, sep=' ')
    dfPlot, = plt.plot(deepfaceData[1], deepfaceData[0], '--',
                       alpha=0.75)
    deepfaceAUC = getAUC(deepfaceData[1], deepfaceData[0])

    baiduData = pd.read_table(
        "comparisons/BaiduIDLFinal.TPFP", header=None, sep=' ')
    bPlot, = plt.plot(baiduData[1], baiduData[0])
    baiduAUC = getAUC(baiduData[1], baiduData[0])

    eigData = pd.read_table(
        "comparisons/eigenfaces-original-roc.txt", header=None, sep=' ')
    eigPlot, = plt.plot(eigData[1], eigData[0])
    eigAUC = getAUC(eigData[1], eigData[0])

    ax.legend([humanPlot, bPlot, dfPlot, brPlot, eigPlot,
               meanPlot_v1, meanPlot_v2, foldPlot_v2],
              ['Human, Cropped [AUC={:.3f}]'.format(humanAUC),
               'Baidu [{:.3f}]'.format(baiduAUC),
               'DeepFace Ensemble [{:.3f}]'.format(deepfaceAUC),
               'OpenBR v1.1.0 [{:.3f}]'.format(brAUC),
               'Eigenfaces (img-restrict) [{:.3f}]'.format(eigAUC),
               'OpenFace nn4.v1 [{:.3f}]'.format(AUC_v1),
               'OpenFace nn4.v2 [{:.3f}]'.format(AUC_v2),
               'OpenFace nn4.v2 folds'],
              loc='lower right')

    plt.plot([0, 1], color='k', linestyle=':')

    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    # plt.ylim(ymin=0,ymax=1)
    plt.xlim(xmin=0, xmax=1)

    plt.grid(b=True, which='major', color='k', linestyle='-')
    plt.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.2)
    plt.minorticks_on()
    fig.savefig(os.path.join(workDir, "roc.pdf"))

if __name__ == '__main__':
    main()
<commit_msg>Change lfw interface.<commit_after>#!/usr/bin/env python3
#
# Copyright 2015 Carnegie Mellon University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import math
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d

from sklearn.cross_validation import KFold
from sklearn.metrics import accuracy_score

import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
plt.style.use('bmh')

import os

import argparse

from scipy import arange

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('tag', type=str, help='The label/tag to put on the ROC curve.')
    parser.add_argument('workDir', type=str,
                        help='The work directory with labels.csv and reps.csv.')
    parser.add_argument('--lfwPairs', type=str,
                        default=os.path.expanduser("~/openface/data/lfw/pairs.txt"))
    args = parser.parse_args()

    print("Loading embeddings.")
    fname = "{}/labels.csv".format(args.workDir)
    paths = pd.read_csv(fname, header=None).as_matrix()[:, 1]
    paths = map(os.path.basename, paths)  # Get the filename.
    # Remove the extension.
    paths = map(lambda path: os.path.splitext(path)[0], paths)
    fname = "{}/reps.csv".format(args.workDir)
    rawEmbeddings = pd.read_csv(fname, header=None).as_matrix()
    embeddings = dict(zip(*[paths, rawEmbeddings]))

    pairs = loadPairs(args.lfwPairs)
    classifyExp(args.workDir, pairs, embeddings)
    plotClassifyExp(args.workDir, args.tag)

def loadPairs(pairsFname):
    print("  + Reading pairs.")
    pairs = []
    with open(pairsFname, 'r') as f:
        for line in f.readlines()[1:]:
            pair = line.strip().split()
            pairs.append(pair)
    assert(len(pairs) == 6000)
    return np.array(pairs)

def getEmbeddings(pair, embeddings):
    if len(pair) == 3:
        name1 = "{}_{}".format(pair[0], pair[1].zfill(4))
        name2 = "{}_{}".format(pair[0], pair[2].zfill(4))
        actual_same = True
    elif len(pair) == 4:
        name1 = "{}_{}".format(pair[0], pair[1].zfill(4))
        name2 = "{}_{}".format(pair[2], pair[3].zfill(4))
        actual_same = False
    else:
        raise Exception(
            "Unexpected pair length: {}".format(len(pair)))

    (x1, x2) = (embeddings[name1], embeddings[name2])
    return (x1, x2, actual_same)

def writeROC(fname, thresholds, embeddings, pairsTest):
    with open(fname, "w") as f:
        f.write("threshold,tp,tn,fp,fn,tpr,fpr\n")
        tp = tn = fp = fn = 0
        for threshold in thresholds:
            tp = tn = fp = fn = 0
            for pair in pairsTest:
                (x1, x2, actual_same) = getEmbeddings(pair, embeddings)
                diff = x1 - x2
                dist = np.dot(diff.T, diff)
                predict_same = dist < threshold

                if predict_same and actual_same:
                    tp += 1
                elif predict_same and not actual_same:
                    fp += 1
                elif not predict_same and not actual_same:
                    tn += 1
                elif not predict_same and actual_same:
                    fn += 1

            if tp + fn == 0:
                tpr = 0
            else:
                tpr = float(tp) / float(tp + fn)
            if fp + tn == 0:
                fpr = 0
            else:
                fpr = float(fp) / float(fp + tn)
            f.write(",".join([str(x)
                              for x in [threshold, tp, tn, fp, fn, tpr, fpr]]))
            f.write("\n")
            if tpr == 1.0 and fpr == 1.0:
                # No further improvements.
                f.write(",".join([str(x)
                                  for x in [4.0, tp, tn, fp, fn, tpr, fpr]]))
                return

def evalThresholdAccuracy(embeddings, pairs, threshold):
    y_true = []
    y_predict = []
    for pair in pairs:
        (x1, x2, actual_same) = getEmbeddings(pair, embeddings)
        diff = x1 - x2
        dist = np.dot(diff.T, diff)
        predict_same = dist < threshold
        y_predict.append(predict_same)
        y_true.append(actual_same)

    y_true = np.array(y_true)
    y_predict = np.array(y_predict)
    accuracy = accuracy_score(y_true, y_predict)
    return accuracy

def findBestThreshold(thresholds, embeddings, pairsTrain):
    bestThresh = bestThreshAcc = 0
    for threshold in thresholds:
        accuracy = evalThresholdAccuracy(embeddings, pairsTrain, threshold)
        if accuracy >= bestThreshAcc:
            bestThreshAcc = accuracy
            bestThresh = threshold
        else:
            # No further improvements.
            return bestThresh
    return bestThresh

def classifyExp(workDir, pairs, embeddings):
    print("  + Computing accuracy.")
    folds = KFold(n=6000, n_folds=10, shuffle=False)
    thresholds = arange(0, 4, 0.01)

    if os.path.exists("{}/accuracies.txt".format(workDir)):
        print("{}/accuracies.txt already exists. Skipping processing.".format(workDir))
    else:
        accuracies = []
        with open("{}/accuracies.txt".format(workDir), "w") as f:
            f.write('fold, threshold, accuracy\n')
            for idx, (train, test) in enumerate(folds):
                fname = "{}/l2-roc.fold-{}.csv".format(workDir, idx)
                writeROC(fname, thresholds, embeddings, pairs[test])

                bestThresh = findBestThreshold(
                    thresholds, embeddings, pairs[train])
                accuracy = evalThresholdAccuracy(
                    embeddings, pairs[test], bestThresh)
                accuracies.append(accuracy)
                f.write('{}, {:0.2f}, {:0.2f}\n'.format(
                    idx, bestThresh, accuracy))
            avg = np.mean(accuracies)
            std = np.std(accuracies)
            f.write('\navg, {:0.4f} +/- {:0.4f}\n'.format(avg, std))
            print('    + {:0.4f}'.format(avg))

def getAUC(fprs, tprs):
    sortedFprs, sortedTprs = zip(*sorted(zip(*(fprs, tprs))))
    sortedFprs = list(sortedFprs)
    sortedTprs = list(sortedTprs)
    if sortedFprs[-1] != 1.0:
        sortedFprs.append(1.0)
        sortedTprs.append(sortedTprs[-1])
    return np.trapz(sortedTprs, sortedFprs)

def plotOpenFaceROC(workDir, plotFolds=True, color=None):
    fs = []
    for i in range(10):
        rocData = pd.read_csv("{}/l2-roc.fold-{}.csv".format(workDir, i))
        fs.append(interp1d(rocData['fpr'], rocData['tpr']))
        x = np.linspace(0, 1, 1000)
        if plotFolds:
            foldPlot, = plt.plot(x, fs[-1](x), color='grey', alpha=0.5)
        else:
            foldPlot = None

    fprs = []
    tprs = []
    for fpr in np.linspace(0, 1, 1000):
        tpr = 0.0
        for f in fs:
            v = f(fpr)
            if math.isnan(v):
                v = 0.0
            tpr += v
        tpr /= 10.0
        fprs.append(fpr)
        tprs.append(tpr)
    if color:
        meanPlot, = plt.plot(fprs, tprs, color=color)
    else:
        meanPlot, = plt.plot(fprs, tprs)
    AUC = getAUC(fprs, tprs)
    return foldPlot, meanPlot, AUC

def plotClassifyExp(workDir, tag):
    print("Plotting.")

    fig, ax = plt.subplots(1, 1)

    openbrData = pd.read_csv("comparisons/openbr.v1.1.0.DET.csv")
    openbrData['Y'] = 1 - openbrData['Y']
    # brPlot = openbrData.plot(x='X', y='Y', legend=True, ax=ax)
    brPlot, = plt.plot(openbrData['X'], openbrData['Y'])
    brAUC = getAUC(openbrData['X'], openbrData['Y'])

    foldPlot, meanPlot, AUC = plotOpenFaceROC(workDir, color='k')

    humanData = pd.read_table(
        "comparisons/kumar_human_crop.txt", header=None, sep=' ')
    humanPlot, = plt.plot(humanData[1], humanData[0])
    humanAUC = getAUC(humanData[1], humanData[0])

    deepfaceData = pd.read_table(
        "comparisons/deepface_ensemble.txt", header=None, sep=' ')
    dfPlot, = plt.plot(deepfaceData[1], deepfaceData[0], '--',
                       alpha=0.75)
    deepfaceAUC = getAUC(deepfaceData[1], deepfaceData[0])

    baiduData = pd.read_table(
        "comparisons/BaiduIDLFinal.TPFP", header=None, sep=' ')
    bPlot, = plt.plot(baiduData[1], baiduData[0])
    baiduAUC = getAUC(baiduData[1], baiduData[0])

    eigData = pd.read_table(
        "comparisons/eigenfaces-original-roc.txt", header=None, sep=' ')
    eigPlot, = plt.plot(eigData[1], eigData[0])
    eigAUC = getAUC(eigData[1], eigData[0])

    ax.legend([humanPlot, bPlot, dfPlot, brPlot, eigPlot,
               meanPlot, foldPlot],
              ['Human, Cropped [AUC={:.3f}]'.format(humanAUC),
               'Baidu [{:.3f}]'.format(baiduAUC),
               'DeepFace Ensemble [{:.3f}]'.format(deepfaceAUC),
               'OpenBR v1.1.0 [{:.3f}]'.format(brAUC),
               'Eigenfaces (img-restrict) [{:.3f}]'.format(eigAUC),
               'OpenFace {} [{:.3f}]'.format(tag, AUC),
               'OpenFace {} folds'.format(tag)],
              loc='lower right')

    plt.plot([0, 1], color='k', linestyle=':')

    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    # plt.ylim(ymin=0,ymax=1)
    plt.xlim(xmin=0, xmax=1)

    plt.grid(b=True, which='major', color='k', linestyle='-')
    plt.grid(b=True, which='minor', color='k', linestyle='-', alpha=0.2)
    plt.minorticks_on()
    # fig.savefig(os.path.join(workDir, "roc.pdf"))
    fig.savefig(os.path.join(workDir, "roc.png"))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>"""
Admin configuration for ``eventkit`` app.
"""

# Define `list_display`, `list_filter` and `search_fields` for each model.
# These go a long way to making the admin more usable.

from dateutil import rrule
import datetime
import json
import six

from django.contrib import admin
from django.core.serializers.json import DjangoJSONEncoder
from django.core.urlresolvers import reverse
from django.http import HttpResponse, JsonResponse
from django.template.response import TemplateResponse
from django.views.decorators.csrf import csrf_exempt
from polymorphic.admin import \
    PolymorphicParentModelAdmin, PolymorphicChildModelAdmin
from timezone import timezone

from eventkit import admin_forms, appsettings, forms, models, plugins

class EventChildAdmin(PolymorphicChildModelAdmin):
    base_form = admin_forms.BaseEventForm
    base_model = models.Event
    formfield_overrides = {
        models.RecurrenceRuleField: {'widget': forms.RecurrenceRuleWidget},
    }

    def save_model(self, request, obj, form, change):
        """
        Propagate changes if requested.
        """
        obj.save(propagate=form.cleaned_data['propagate'])

class EventAdmin(PolymorphicParentModelAdmin):
    base_model = models.Event
    list_filter = ('all_day', 'starts', 'ends', 'is_repeat')
    list_display = ('__str__', 'all_day', 'starts', 'ends', 'is_repeat')
    search_fields = ('title', )

    def get_urls(self):
        """
        Add a calendar URL.
        """
        from django.conf.urls import patterns, url
        urls = super(EventAdmin, self).get_urls()
        my_urls = patterns(
            '',
            url(
                r'^calendar/$',
                self.admin_site.admin_view(self.calendar),
                name='eventkit_event_calendar'
            ),
        )
        return my_urls + urls

    def calendar(self, request):
        """
        Return event data in JSON format for AJAX requests, or a calendar page
        to be loaded in an iframe.
        """
        if not request.is_ajax():
            return TemplateResponse(
                request, 'admin/eventkit/event/calendar.html')
        tz = timezone.get(request.GET.get('timezone'))
        starts = timezone.localize(
            datetime.datetime.strptime(request.GET['start'], '%Y-%m-%d'), tz)
        ends = timezone.localize(
            datetime.datetime.strptime(request.GET['end'], '%Y-%m-%d'), tz)
        events = self.get_queryset(request) \
            .filter(starts__gte=starts, starts__lt=ends) \
            .values_list(
                'id', 'parent', 'title', 'all_day', 'starts', 'ends',
                'is_repeat')
        data = []
        # Get the primary keys for all root and variation events into a list,
        # so we can assign them and their repeat events a consistent colour.
        seen = list(
            models.Event.objects
            .filter(is_repeat=False)
            .values_list('pk', flat=True)
            .order_by('pk'))
        for pk, parent, title, all_day, starts, ends, is_repeat in events:
            id_ = parent if is_repeat else pk
            if id_ not in seen:
                seen.append(id_)
            # Get color based on seen index for parent event, so its repeat
            # events have the same color. Start repeating colors when they have
            # all been used.
            background, color = appsettings.CALENDAR_COLORS[
                seen.index(id_) % len(appsettings.CALENDAR_COLORS)]
            data.append({
                'id': id_,
                'title': title,
                'allDay': all_day,
                'start': timezone.localize(starts),
                'end': timezone.localize(ends),
                'url': reverse('admin:eventkit_event_change', args=[pk]),
                'color': background,
                'textColor': color,
            })
        data = json.dumps(data, cls=DjangoJSONEncoder)
        return HttpResponse(content=data, content_type='applicaton/json')

    def get_child_models(self):
        """
        Get child models from registered ``EventPlugin`` subclasses. Fallback
        to the parent ``Event`` model if no plugins are registered.
        """
        child_models = []
        for plugin in plugins.EventPlugin.plugins:
            child_models.append((plugin.model, plugin.model_admin))
        if not child_models:
            child_models.append((models.Event, EventChildAdmin))
        return child_models

admin.site.register(models.Event, EventAdmin)

class RecurrenceRuleAdmin(admin.ModelAdmin):
    formfield_overrides = {
        models.RecurrenceRuleField: {'widget': forms.RecurrenceRuleWidget},
    }
    model = models.RecurrenceRule

    def get_urls(self):
        """
        Add a preview URL.
        """
        from django.conf.urls import patterns, url
        urls = super(RecurrenceRuleAdmin, self).get_urls()
        my_urls = patterns(
            '',
            url(
                r'^preview/$',
                self.admin_site.admin_view(self.preview),
                name='eventkit_recurrencerule_preview'
            ),
        )
        return my_urls + urls

    @csrf_exempt
    def preview(self, request):
        """
        Return a occurrences in JSON format up until the configured limit.
        """
        recurrence_rule = request.POST.get('recurrence_rule')
        limit = int(request.POST.get('limit', 10))
        try:
            rruleset = rrule.rrulestr(
                recurrence_rule, dtstart=timezone.now(), forceset=True)
        except ValueError as e:
            data = {
                'error': six.text_type(e),
            }
        else:
            data = {
                'occurrences': rruleset[:limit]
            }
        return JsonResponse(data)

admin.site.register(models.RecurrenceRule, RecurrenceRuleAdmin)
<commit_msg>Add missing docstring.<commit_after>"""
Admin configuration for ``eventkit`` app.
"""

# Define `list_display`, `list_filter` and `search_fields` for each model.
# These go a long way to making the admin more usable.

from dateutil import rrule
import datetime
import json
import six

from django.contrib import admin
from django.core.serializers.json import DjangoJSONEncoder
from django.core.urlresolvers import reverse
from django.http import HttpResponse, JsonResponse
from django.template.response import TemplateResponse
from django.views.decorators.csrf import csrf_exempt
from polymorphic.admin import \
    PolymorphicParentModelAdmin, PolymorphicChildModelAdmin
from timezone import timezone

from eventkit import admin_forms, appsettings, forms, models, plugins

class EventChildAdmin(PolymorphicChildModelAdmin):
    """
    Abstract admin class for polymorphic child event models.
    """
    base_form = admin_forms.BaseEventForm
    base_model = models.Event
    formfield_overrides = {
        models.RecurrenceRuleField: {'widget': forms.RecurrenceRuleWidget},
    }

    def save_model(self, request, obj, form, change):
        """
        Propagate changes if requested.
        """
        obj.save(propagate=form.cleaned_data['propagate'])

class EventAdmin(PolymorphicParentModelAdmin):
    base_model = models.Event
    list_filter = ('all_day', 'starts', 'ends', 'is_repeat')
    list_display = ('__str__', 'all_day', 'starts', 'ends', 'is_repeat')
    search_fields = ('title', )

    def get_urls(self):
        """
        Add a calendar URL.
        """
        from django.conf.urls import patterns, url
        urls = super(EventAdmin, self).get_urls()
        my_urls = patterns(
            '',
            url(
                r'^calendar/$',
                self.admin_site.admin_view(self.calendar),
                name='eventkit_event_calendar'
            ),
        )
        return my_urls + urls

    def calendar(self, request):
        """
        Return event data in JSON format for AJAX requests, or a calendar page
        to be loaded in an iframe.
        """
        if not request.is_ajax():
            return TemplateResponse(
                request, 'admin/eventkit/event/calendar.html')
        tz = timezone.get(request.GET.get('timezone'))
        starts = timezone.localize(
            datetime.datetime.strptime(request.GET['start'], '%Y-%m-%d'), tz)
        ends = timezone.localize(
            datetime.datetime.strptime(request.GET['end'], '%Y-%m-%d'), tz)
        events = self.get_queryset(request) \
            .filter(starts__gte=starts, starts__lt=ends) \
            .values_list(
                'id', 'parent', 'title', 'all_day', 'starts', 'ends',
                'is_repeat')
        data = []
        # Get the primary keys for all root and variation events into a list,
        # so we can assign them and their repeat events a consistent colour.
        seen = list(
            models.Event.objects
            .filter(is_repeat=False)
            .values_list('pk', flat=True)
            .order_by('pk'))
        for pk, parent, title, all_day, starts, ends, is_repeat in events:
            id_ = parent if is_repeat else pk
            if id_ not in seen:
                seen.append(id_)
            # Get color based on seen index for parent event, so its repeat
            # events have the same color. Start repeating colors when they have
            # all been used.
            background, color = appsettings.CALENDAR_COLORS[
                seen.index(id_) % len(appsettings.CALENDAR_COLORS)]
            data.append({
                'id': id_,
                'title': title,
                'allDay': all_day,
                'start': timezone.localize(starts),
                'end': timezone.localize(ends),
                'url': reverse('admin:eventkit_event_change', args=[pk]),
                'color': background,
                'textColor': color,
            })
        data = json.dumps(data, cls=DjangoJSONEncoder)
        return HttpResponse(content=data, content_type='applicaton/json')

    def get_child_models(self):
        """
        Get child models from registered ``EventPlugin`` subclasses. Fallback
        to the parent ``Event`` model if no plugins are registered.
        """
        child_models = []
        for plugin in plugins.EventPlugin.plugins:
            child_models.append((plugin.model, plugin.model_admin))
        if not child_models:
            child_models.append((models.Event, EventChildAdmin))
        return child_models

admin.site.register(models.Event, EventAdmin)

class RecurrenceRuleAdmin(admin.ModelAdmin):
    formfield_overrides = {
        models.RecurrenceRuleField: {'widget': forms.RecurrenceRuleWidget},
    }
    model = models.RecurrenceRule

    def get_urls(self):
        """
        Add a preview URL.
        """
        from django.conf.urls import patterns, url
        urls = super(RecurrenceRuleAdmin, self).get_urls()
        my_urls = patterns(
            '',
            url(
                r'^preview/$',
                self.admin_site.admin_view(self.preview),
                name='eventkit_recurrencerule_preview'
            ),
        )
        return my_urls + urls

    @csrf_exempt
    def preview(self, request):
        """
        Return a occurrences in JSON format up until the configured limit.
        """
        recurrence_rule = request.POST.get('recurrence_rule')
        limit = int(request.POST.get('limit', 10))
        try:
            rruleset = rrule.rrulestr(
                recurrence_rule, dtstart=timezone.now(), forceset=True)
        except ValueError as e:
            data = {
                'error': six.text_type(e),
            }
        else:
            data = {
                'occurrences': rruleset[:limit]
            }
        return JsonResponse(data)

admin.site.register(models.RecurrenceRule, RecurrenceRuleAdmin)
<|endoftext|>


<commit_before>import math
import os.path
import pickle
import Sofa

import Flexible.IO
import sys

import SofaPython.Tools
from SofaPython.Tools import listToStr as concat
from SofaPython import Quaternion

__file = __file__.replace('\\', '/')  # windows compatible filename

def insertLinearMapping(node, dofRigidNode=None, dofAffineNode=None, position=None, labelImage=None, labels=None, assemble=True):
    """ insert the correct Linear(Multi)Mapping
    hopefully the template is deduced automatically by the component
    """
    if dofRigidNode is None and dofAffineNode is None:
        print "[Flexible.API.insertLinearMapping] ERROR: no dof given"
    if dofAffineNode is None:
        #TODO
        return None
    else:
        cell = ''
        if not labelImage is None and not labels is None : # use labels to select specific voxels in branching image
            node.createObject(
                "BranchingCellOffsetsFromPositions", template="BranchingImageUC", name="cell",
                position ="@"+SofaPython.Tools.getObjectPath(position)+".position",
                src="@"+SofaPython.Tools.getObjectPath(labelImage.branchingImage), labels=concat(labels))
            cell = "@cell.cell"

        if dofRigidNode is None:
            #TODO
            return None
        else:
            return node.createObject(
                "LinearMultiMapping", cell=cell,
                input1="@"+dofRigidNode.getPathName(), input2="@"+dofAffineNode.getPathName(), output="@.", assemble=assemble)

class Deformable:

    def __init__(self, node, name):
        self.node=node.createChild(name)
        self.name=name
        self.dofs=None
        self.meshLoader=None
        self.topology=None
        self.mass=None
        self.visual=None
        self.mapping=None

    def addMesh(self, meshPath, offset = [0,0,0,0,0,0,1]):
        r = Quaternion.to_euler(offset[3:])  * 180.0 / math.pi
        self.meshLoader = SofaPython.Tools.meshLoader(self.node, meshPath, translation=concat(offset[:3]) , rotation=concat(r))
        self.topology = self.node.createObject("MeshTopology", name="topology", src="@"+self.meshLoader.name )
        self.dofs = self.node.createObject("MechanicalObject", template = "Vec3d", name="dofs", src="@"+self.meshLoader.name)

    def addMass(self,totalMass):
        self.mass = self.node.createObject('UniformMass', totalMass=totalMass)

    def addMapping(self, dofRigidNode=None, dofAffineNode=None, labelImage=None, labels=None, assemble=True):
        self.mapping = insertLinearMapping(self.node, dofRigidNode, dofAffineNode, self.topology, labelImage, labels, assemble)

# TODO: refactor this
    def addSkinning(self, bonesPath, indices, weights, assemble=True):
        self.mapping = self.node.createObject(
            "LinearMapping", template="Rigid3d,Vec3d", name="skinning",
            input="@"+bonesPath, indices=concat(indices), weights=concat(weights), assemble=assemble)

    def addVisual(self, color=[1,1,1,1]):
        self.visual = Deformable.VisualModel(self.node, color)

    class VisualModel:
        def __init__(self, node, color=[1,1,1,1] ):
            self.node = node.createChild("visual")
            self.model = self.node.createObject("VisualModel", name="model", color=concat(color))
            self.mapping = self.node.createObject("IdentityMapping", name="mapping")

class AffineMass:
    def __init__(self, node, dofAffineNode):
        self.node = node # a children node of all nodes and shape function
        self.dofAffineNode = dofAffineNode # where the mechanical state is located

    def massFromDensityImage(self, dofRigidNode, densityImage, lumping='0'):
        node = self.node.createChild('Mass')
        dof = node.createObject('MechanicalObject', name='massPoints', template='Vec3d')
        insertLinearMapping(node, dofRigidNode, self.dofAffineNode, dof, assemble=False)
        densityImage.addBranchingToImage('0') # MassFromDensity on branching images does not exist yet
        massFromDensity = node.createObject('MassFromDensity',  name="MassFromDensity",  template="Affine,ImageD", image="@"+SofaPython.Tools.getObjectPath(densityImage.converter)+".image", transform="@"+SofaPython.Tools.getObjectPath(densityImage.converter)+'.transform', lumping=lumping)
        self.dofAffineNode.createObject('AffineMass', name='mass', massMatrix="@"+SofaPython.Tools.getObjectPath(massFromDensity)+".massMatrix")

    def read(self, directory):
#        with open(os.path.join(directory,"affineMass.pkl"), "r") as f:
#            data = pickle.load(f)
#            self.dofAffineNode.createObject('AffineMass', name='mass', massMatrix=data.mass)
#            print 'Imported Affine Mass from '+directory+"/affineMass.pkl"
        sys.path.insert(0, directory)
        __import__('affineMass').loadMass(self.dofAffineNode)
        print 'Imported Affine Mass from '+directory+"/affineMass.py"

    class InternalData:
        def __init__(self,node):
            self.mass = None

    def write(self, directory):
        self.dofAffineNode.createObject('PythonScriptController', filename=__file__, classname='massExporter', variables=directory)

class massExporter(Sofa.PythonScriptController):
    def bwdInitGraph(self,node):
        directory = self.findData('variables').value[0][0]
#        with open(os.path.join(directory,"affineMass.pkl"), "w") as f:
#            pickle.dump(self.InternalData(node), f)
#        print 'Exported Affine Mass in '+directory+"/affineMass.pkl";
        Flexible.IO.export_AffineMass(node.getObject('mass'), directory+"/affineMass.py")
        print 'Exported Affine Mass in '+directory+"/affineMass.py"
        return 0
    class InternalData:
        def __init__(self,node):
            self.mass = str(node.getObject('mass').massMatrix).replace('\n',' ')

class ShapeFunction:
    """ High-level API to manipulate ShapeFunction
    @todo better handle template
    """
    def __init__(self, node, position=None):
        self.node = node
        self.position = position # component which contains shape function position (spatial coordinates of the parent nodes)
        self.shapeFunction=None

    def addVoronoi(self, image):
        """ Add a Voronoi shape function using position from position component and BranchingImage image
        """
        if self.position is None:
            print "[Flexible.API.ShapeFunction] ERROR: no position"
        imagePath = SofaPython.Tools.getObjectPath(image.branchingImage)
        self.shapeFunction = self.node.createObject(
            "VoronoiShapeFunction", template="ShapeFunctiond,"+"Branching"+image.imageType,
            name="shapeFunction",
            position="@"+SofaPython.Tools.getObjectPath(self.position)+".position",
            src="@"+imagePath, method=0, nbRef=8, bias=True)

    def getFilenameIndices(self, filenamePrefix=None, directory=""):
        _filename=filenamePrefix if not filenamePrefix is None else "SF"
        _filename+="_indices.mhd"
        _filename=os.path.join(directory, _filename)
        return _filename

    def getFilenameWeights(self, filenamePrefix=None, directory=""):
        _filename=filenamePrefix if not filenamePrefix is None else "SF"
        _filename+="_weights.mhd"
        _filename=os.path.join(directory, _filename)
        return _filename

    def addExporter(self, filenamePrefix=None, directory=""):
        if self.shapeFunction is None:
            print "[Flexible.API.ShapeFunction] ERROR: no shapeFunction"
        sfPath = SofaPython.Tools.getObjectPath(self.shapeFunction)
        self.node.createObject(
            "ImageExporter", template="BranchingImageUI", name="exporterIndices",
            image="@"+sfPath+".indices", transform="@"+sfPath+".transform",
            filename=self.getFilenameIndices(filenamePrefix, directory),
            exportAtEnd=True, printLog=True)
        self.node.createObject(
            "ImageExporter", template="BranchingImageD", name="exporterWeights",
            image="@"+sfPath+".weights", transform="@"+sfPath+".transform",
            filename=self.getFilenameWeights(filenamePrefix, directory), exportAtEnd=True, printLog=True)

    def addContainer(self, filenamePrefix=None, directory=""):
        if self.position is None:
            print "[Flexible.API.ShapeFunction] ERROR: no position"
        self.node.createObject(
            "ImageContainer", template="BranchingImageUI", name="containerIndices",
            filename=self.getFilenameIndices(filenamePrefix, directory), drawBB=False)
        self.node.createObject(
            "ImageContainer", template="BranchingImageD", name="containerWeights",
            filename=self.getFilenameWeights(filenamePrefix, directory), drawBB=False)
        self.shapeFunction = self.node.createObject(
            "ImageShapeFunctionContainer", template="ShapeFunctiond,BranchingImageUC", name="shapeFunction", position="@"+SofaPython.Tools.getObjectPath(self.position)+".position",
            transform="@containerWeights.transform",
            weights="@containerWeights.image", indices="@containerIndices.image")

class Behavior:
    """ High level API to add a behavior
    """
    def __init__(self, node, name, type="331", labelImage=None, labels=None):
        self.node = node.createChild(name)
        self.name = name
        self.labelImage = labelImage
        self.labels = labels
        self.type = type
        self.sampler = None
        self.dofs = None
        self.mapping = None

    def addGaussPointSampler(self, shapeFunction, nbPoint):
        shapeFunctionPath = SofaPython.Tools.getObjectPath(shapeFunction.shapeFunction)
        self.sampler = self.node.createObject(
            "ImageGaussPointSampler", template="BranchingImageD,BranchingImageUC", name="sampler",
            indices="@"+shapeFunctionPath+".indices", weights="@"+shapeFunctionPath+".weights", transform="@"+shapeFunctionPath+".transform",
            method="2", order=self.type[2:], targetNumber=nbPoint,
            mask="@"+SofaPython.Tools.getObjectPath(self.labelImage.branchingImage)+".branchingImage", maskLabels=concat(self.labels), clearData=True)

    def addMechanicalObject(self, dofRigidNode=None, dofAffineNode=None, assemble=True):
        if self.sampler is None:
            print "[Flexible.API.Behavior] ERROR: no sampler"
        self.dofs = self.node.createObject("MechanicalObject", template="F"+self.type, name="dofs")
        self.mapping = insertLinearMapping(self.node, dofRigidNode, dofAffineNode, self.sampler, self.labelImage, self.labels, assemble)

    def addHooke(self, strainMeasure="Corotational", youngModulus=0, poissonRatio=0, viscosity=0, assemble=True):
        eNode = self.node.createChild("E")
        eNode.createObject('MechanicalObject',  template="E"+self.type, name="E")
        eNode.createObject(strainMeasure+'StrainMapping', template="F"+self.type+",E"+self.type, assemble=assemble)
        eNode.createObject('HookeForceField',  template="E"+self.type, youngModulus= youngModulus, poissonRatio=poissonRatio, viscosity=viscosity, assemble=assemble, isCompliance=False)

<commit_msg>[Flexible] add TODO<commit_after>import math
import os.path
import pickle
import Sofa

import Flexible.IO
import sys

import SofaPython.Tools
from SofaPython.Tools import listToStr as concat
from SofaPython import Quaternion

__file = __file__.replace('\\', '/')  # windows compatible filename

def insertLinearMapping(node, dofRigidNode=None, dofAffineNode=None, position=None, labelImage=None, labels=None, assemble=True):
    """ insert the correct Linear(Multi)Mapping
    hopefully the template is deduced automatically by the component
    """
    if dofRigidNode is None and dofAffineNode is None:
        print "[Flexible.API.insertLinearMapping] ERROR: no dof given"
    if dofAffineNode is None:
        #TODO
        return None
    else:
        cell = ''
        if not labelImage is None and not labels is None : # use labels to select specific voxels in branching image
            node.createObject(
                "BranchingCellOffsetsFromPositions", template="BranchingImageUC", name="cell",
                position ="@"+SofaPython.Tools.getObjectPath(position)+".position",
                src="@"+SofaPython.Tools.getObjectPath(labelImage.branchingImage), labels=concat(labels))
            cell = "@cell.cell"

        if dofRigidNode is None:
            #TODO
            return None
        else:
            return node.createObject(
                "LinearMultiMapping", cell=cell,
                input1="@"+dofRigidNode.getPathName(), input2="@"+dofAffineNode.getPathName(), output="@.", assemble=assemble)

class Deformable:

    def __init__(self, node, name):
        self.node=node.createChild(name)
        self.name=name
        self.dofs=None
        self.meshLoader=None
        self.topology=None
        self.mass=None
        self.visual=None
        self.mapping=None

    def addMesh(self, meshPath, offset = [0,0,0,0,0,0,1]):
        r = Quaternion.to_euler(offset[3:])  * 180.0 / math.pi
        self.meshLoader = SofaPython.Tools.meshLoader(self.node, meshPath, translation=concat(offset[:3]) , rotation=concat(r))
        self.topology = self.node.createObject("MeshTopology", name="topology", src="@"+self.meshLoader.name )
        self.dofs = self.node.createObject("MechanicalObject", template = "Vec3d", name="dofs", src="@"+self.meshLoader.name)

    def addMass(self,totalMass):
        self.mass = self.node.createObject('UniformMass', totalMass=totalMass)

    def addMapping(self, dofRigidNode=None, dofAffineNode=None, labelImage=None, labels=None, assemble=True):
        self.mapping = insertLinearMapping(self.node, dofRigidNode, dofAffineNode, self.topology, labelImage, labels, assemble)

# TODO: refactor this
    def addSkinning(self, bonesPath, indices, weights, assemble=True):
        self.mapping = self.node.createObject(
            "LinearMapping", template="Rigid3d,Vec3d", name="skinning",
            input="@"+bonesPath, indices=concat(indices), weights=concat(weights), assemble=assemble)

    def addVisual(self, color=[1,1,1,1]):
        self.visual = Deformable.VisualModel(self.node, color)

    class VisualModel:
        def __init__(self, node, color=[1,1,1,1] ):
            self.node = node.createChild("visual")
            self.model = self.node.createObject("VisualModel", name="model", color=concat(color))
            self.mapping = self.node.createObject("IdentityMapping", name="mapping")

class AffineMass:
    def __init__(self, node, dofAffineNode):
        self.node = node # a children node of all nodes and shape function
        self.dofAffineNode = dofAffineNode # where the mechanical state is located

    def massFromDensityImage(self, dofRigidNode, densityImage, lumping='0'):
        node = self.node.createChild('Mass')
        dof = node.createObject('MechanicalObject', name='massPoints', template='Vec3d')
        insertLinearMapping(node, dofRigidNode, self.dofAffineNode, dof, assemble=False) # TODO: ERROR: AssemblyVisitor: if matrix is not assembled, crash when the matrix is assembled
        densityImage.addBranchingToImage('0') # MassFromDensity on branching images does not exist yet
        massFromDensity = node.createObject('MassFromDensity',  name="MassFromDensity",  template="Affine,ImageD", image="@"+SofaPython.Tools.getObjectPath(densityImage.converter)+".image", transform="@"+SofaPython.Tools.getObjectPath(densityImage.converter)+'.transform', lumping=lumping)
        self.dofAffineNode.createObject('AffineMass', name='mass', massMatrix="@"+SofaPython.Tools.getObjectPath(massFromDensity)+".massMatrix")

    def read(self, directory):
#        with open(os.path.join(directory,"affineMass.pkl"), "r") as f:
#            data = pickle.load(f)
#            self.dofAffineNode.createObject('AffineMass', name='mass', massMatrix=data.mass)
#            print 'Imported Affine Mass from '+directory+"/affineMass.pkl"
        sys.path.insert(0, directory)
        __import__('affineMass').loadMass(self.dofAffineNode)
        print 'Imported Affine Mass from '+directory+"/affineMass.py"

    class InternalData:
        def __init__(self,node):
            self.mass = None

    def write(self, directory):
        self.dofAffineNode.createObject('PythonScriptController', filename=__file__, classname='massExporter', variables=directory)

class massExporter(Sofa.PythonScriptController):
    def bwdInitGraph(self,node):
        directory = self.findData('variables').value[0][0]
#        with open(os.path.join(directory,"affineMass.pkl"), "w") as f:
#            pickle.dump(self.InternalData(node), f)
#        print 'Exported Affine Mass in '+directory+"/affineMass.pkl";
        Flexible.IO.export_AffineMass(node.getObject('mass'), directory+"/affineMass.py")
        print 'Exported Affine Mass in '+directory+"/affineMass.py"
        return 0
    class InternalData:
        def __init__(self,node):
            self.mass = str(node.getObject('mass').massMatrix).replace('\n',' ')

class ShapeFunction:
    """ High-level API to manipulate ShapeFunction
    @todo better handle template
    """
    def __init__(self, node, position=None):
        self.node = node
        self.position = position # component which contains shape function position (spatial coordinates of the parent nodes)
        self.shapeFunction=None

    def addVoronoi(self, image):
        """ Add a Voronoi shape function using position from position component and BranchingImage image
        """
        if self.position is None:
            print "[Flexible.API.ShapeFunction] ERROR: no position"
        imagePath = SofaPython.Tools.getObjectPath(image.branchingImage)
        self.shapeFunction = self.node.createObject(
            "VoronoiShapeFunction", template="ShapeFunctiond,"+"Branching"+image.imageType,
            name="shapeFunction",
            position="@"+SofaPython.Tools.getObjectPath(self.position)+".position",
            src="@"+imagePath, method=0, nbRef=8, bias=True)

    def getFilenameIndices(self, filenamePrefix=None, directory=""):
        _filename=filenamePrefix if not filenamePrefix is None else "SF"
        _filename+="_indices.mhd"
        _filename=os.path.join(directory, _filename)
        return _filename

    def getFilenameWeights(self, filenamePrefix=None, directory=""):
        _filename=filenamePrefix if not filenamePrefix is None else "SF"
        _filename+="_weights.mhd"
        _filename=os.path.join(directory, _filename)
        return _filename

    def addExporter(self, filenamePrefix=None, directory=""):
        if self.shapeFunction is None:
            print "[Flexible.API.ShapeFunction] ERROR: no shapeFunction"
        sfPath = SofaPython.Tools.getObjectPath(self.shapeFunction)
        self.node.createObject(
            "ImageExporter", template="BranchingImageUI", name="exporterIndices",
            image="@"+sfPath+".indices", transform="@"+sfPath+".transform",
            filename=self.getFilenameIndices(filenamePrefix, directory),
            exportAtEnd=True, printLog=True)
        self.node.createObject(
            "ImageExporter", template="BranchingImageD", name="exporterWeights",
            image="@"+sfPath+".weights", transform="@"+sfPath+".transform",
            filename=self.getFilenameWeights(filenamePrefix, directory), exportAtEnd=True, printLog=True)

    def addContainer(self, filenamePrefix=None, directory=""):
        if self.position is None:
            print "[Flexible.API.ShapeFunction] ERROR: no position"
        self.node.createObject(
            "ImageContainer", template="BranchingImageUI", name="containerIndices",
            filename=self.getFilenameIndices(filenamePrefix, directory), drawBB=False)
        self.node.createObject(
            "ImageContainer", template="BranchingImageD", name="containerWeights",
            filename=self.getFilenameWeights(filenamePrefix, directory), drawBB=False)
        self.shapeFunction = self.node.createObject(
            "ImageShapeFunctionContainer", template="ShapeFunctiond,BranchingImageUC", name="shapeFunction", position="@"+SofaPython.Tools.getObjectPath(self.position)+".position",
            transform="@containerWeights.transform",
            weights="@containerWeights.image", indices="@containerIndices.image")

class Behavior:
    """ High level API to add a behavior
    """
    def __init__(self, node, name, type="331", labelImage=None, labels=None):
        self.node = node.createChild(name)
        self.name = name
        self.labelImage = labelImage
        self.labels = labels
        self.type = type
        self.sampler = None
        self.dofs = None
        self.mapping = None

    def addGaussPointSampler(self, shapeFunction, nbPoint):
        shapeFunctionPath = SofaPython.Tools.getObjectPath(shapeFunction.shapeFunction)
        self.sampler = self.node.createObject(
            "ImageGaussPointSampler", template="BranchingImageD,BranchingImageUC", name="sampler",
            indices="@"+shapeFunctionPath+".indices", weights="@"+shapeFunctionPath+".weights", transform="@"+shapeFunctionPath+".transform",
            method="2", order=self.type[2:], targetNumber=nbPoint,
            mask="@"+SofaPython.Tools.getObjectPath(self.labelImage.branchingImage)+".branchingImage", maskLabels=concat(self.labels), clearData=True)

    def addMechanicalObject(self, dofRigidNode=None, dofAffineNode=None, assemble=True):
        if self.sampler is None:
            print "[Flexible.API.Behavior] ERROR: no sampler"
        self.dofs = self.node.createObject("MechanicalObject", template="F"+self.type, name="dofs")
        self.mapping = insertLinearMapping(self.node, dofRigidNode, dofAffineNode, self.sampler, self.labelImage, self.labels, assemble)

    def addHooke(self, strainMeasure="Corotational", youngModulus=0, poissonRatio=0, viscosity=0, assemble=True):
        eNode = self.node.createChild("E")
        eNode.createObject('MechanicalObject',  template="E"+self.type, name="E")
        eNode.createObject(strainMeasure+'StrainMapping', template="F"+self.type+",E"+self.type, assemble=assemble)
        eNode.createObject('HookeForceField',  template="E"+self.type, youngModulus= youngModulus, poissonRatio=poissonRatio, viscosity=viscosity, assemble=assemble, isCompliance=False)

<|endoftext|>