<commit_before>"""
Paramiko run support
"""
from cStringIO import StringIO
from paramiko import ChannelFile

import gevent
import gevent.event
import socket
import pipes
import logging
import shutil

from teuthology.contextutil import safe_while
from teuthology.exceptions import (CommandCrashedError, CommandFailedError,
                                   ConnectionLostError)

log = logging.getLogger(__name__)

class RemoteProcess(object):
    """
    An object to begin and monitor execution of a process on a remote host
    """
    __slots__ = [
        'client', 'args', 'check_status', 'command', 'hostname',
        'stdin', 'stdout', 'stderr',
        '_stdin_buf', '_stdout_buf', '_stderr_buf',
        'returncode', 'exitstatus', 'timeout',
        'greenlets',
        '_wait', 'logger',
        # for orchestra.remote.Remote to place a backreference
        'remote',
        'label',
        ]

    deadlock_warning = "Using PIPE for %s without wait=False would deadlock"

    def __init__(self, client, args, check_status=True, hostname=None,
                 label=None, timeout=None, wait=True, logger=None, cwd=None):
        """
        Create the object. Does not initiate command execution.

        :param client:       paramiko.SSHConnection to run the command with
        :param args:         Command to run.
        :type args:          String or list of strings
        :param check_status: Whether to raise CommandFailedError on non-zero
                             exit status, and . Defaults to True. All signals
                             and connection loss are made to look like SIGHUP.
        :param hostname:     Name of remote host (optional)
        :param label:        Can be used to label or describe what the
                             command is doing.
        :param timeout:      timeout value for arg that is passed to
                             exec_command of paramiko
        :param wait:         Whether self.wait() will be called automatically
        :param logger:       Alternative logger to use (optional)
        :param cwd:          Directory in which the command will be executed
                             (optional)
        """
        self.client = client
        self.args = args
        if isinstance(args, basestring):
            self.command = args
        else:
            self.command = quote(args)

        if cwd:
            self.command = '(cd {cwd} && exec {cmd})'.format(
                           cwd=cwd, cmd=self.command)

        self.check_status = check_status
        self.label = label
        if timeout:
            self.timeout = timeout
        if hostname:
            self.hostname = hostname
        else:
            (self.hostname, port) = client.get_transport().getpeername()[0:2]

        self.greenlets = []
        self.stdin, self.stdout, self.stderr = (None, None, None)
        self.returncode = self.exitstatus = None
        self._wait = wait
        self.logger = logger or log

    def execute(self):
        """
        Execute remote command
        """
        prefix = "Running:"
        if self.label:
            prefix = "Running ({label}):".format(label=self.label)
        log.getChild(self.hostname).info(prefix)
        for line in self.command.split('\n'):
            log.getChild(self.hostname).info('> %s' % line)

        if hasattr(self, 'timeout'):
            (self._stdin_buf, self._stdout_buf, self._stderr_buf) = \
                self.client.exec_command(self.command, timeout=self.timeout)
        else:
            (self._stdin_buf, self._stdout_buf, self._stderr_buf) = \
                self.client.exec_command(self.command)
        (self.stdin, self.stdout, self.stderr) = \
            (self._stdin_buf, self._stdout_buf, self._stderr_buf)

    def add_greenlet(self, greenlet):
        self.greenlets.append(greenlet)

    def setup_stdin(self, stream_obj):
        self.stdin = KludgeFile(wrapped=self.stdin)
        if stream_obj is not PIPE:
            greenlet = gevent.spawn(copy_and_close, stream_obj, self.stdin)
            self.add_greenlet(greenlet)
            self.stdin = None
        elif self._wait:
            # FIXME: Is this actually true?
            raise RuntimeError(self.deadlock_warning % 'stdin')

    def setup_output_stream(self, stream_obj, stream_name):
        if stream_obj is not PIPE:
            # Log the stream
            host_log = self.logger.getChild(self.hostname)
            stream_log = host_log.getChild(stream_name)
            self.add_greenlet(
                gevent.spawn(
                    copy_file_to,
                    getattr(self, stream_name),
                    stream_log,
                    stream_obj,
                )
            )
            setattr(self, stream_name, stream_obj)
        elif self._wait:
            # FIXME: Is this actually true?
            raise RuntimeError(self.deadlock_warning % stream_name)

    def wait(self):
        """
        Block until remote process finishes.

        :returns: self.returncode
        """

        status = self._get_exitstatus()
        if status != 0:
            log.debug("got remote process result: {}".format(status))
        for greenlet in self.greenlets:
            try:
                greenlet.get(block=True,timeout=60)
            except gevent.Timeout:
                log.debug("timed out waiting; will kill: {}".format(greenlet))
                greenlet.kill(block=False)
        for stream in ('stdout', 'stderr'):
            if hasattr(self, stream):
                stream_obj = getattr(self, stream)
                # Despite ChannelFile having a seek() method, it raises
                # "IOError: File does not support seeking."
                if hasattr(stream_obj, 'seek') and \
                        not isinstance(stream_obj, ChannelFile):
                    stream_obj.seek(0)

        self._raise_for_status()
        return status

    def _raise_for_status(self):
        if self.returncode is None:
            self._get_exitstatus()
        if self.check_status:
            if self.returncode in (None, -1):
                # command either died due to a signal, or the connection
                # was lost
                transport = self.client.get_transport()
                if transport is None or not transport.is_active():
                    # look like we lost the connection
                    raise ConnectionLostError(command=self.command,
                                              node=self.hostname)

                # connection seems healthy still, assuming it was a
                # signal; sadly SSH does not tell us which signal
                raise CommandCrashedError(command=self.command)
            if self.returncode != 0:
                raise CommandFailedError(
                    command=self.command, exitstatus=self.returncode,
                    node=self.hostname, label=self.label
                )

    def _get_exitstatus(self):
        """
        :returns: the remote command's exit status (return code). Note that
                  if the connection is lost, or if the process was killed by a
                  signal, this returns None instead of paramiko's -1.
        """
        status = self._stdout_buf.channel.recv_exit_status()
        self.exitstatus = self.returncode = status
        if status == -1:
            status = None
        return status

    @property
    def finished(self):
        gevent.wait(self.greenlets, timeout=0.1)
        ready = self._stdout_buf.channel.exit_status_ready()
        if ready:
            self._get_exitstatus()
        return ready

    def poll(self):
        """
        :returns: self.returncode if the process is finished; else None
        """
        if self.finished:
            self._raise_for_status()
            return self.returncode
        return None

    def __repr__(self):
        return '{classname}(client={client!r}, args={args!r}, check_status={check}, hostname={name!r})'.format(  # noqa
            classname=self.__class__.__name__,
            client=self.client,
            args=self.args,
            check=self.check_status,
            name=self.hostname,
            )

class Raw(object):

    """
    Raw objects are passed to remote objects and are not processed locally.
    """
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return '{cls}({value!r})'.format(
            cls=self.__class__.__name__,
            value=self.value,
            )

    def __eq__(self, value):
        return self.value == value

def quote(args):
    """
    Internal quote wrapper.
    """
    if isinstance(args, basestring):
        return args

    def _quote(args):
        """
        Handle quoted string, testing for raw charaters.
        """
        for a in args:
            if isinstance(a, Raw):
                yield a.value
            else:
                yield pipes.quote(a)
    return ' '.join(_quote(args))

def copy_to_log(f, logger, loglevel=logging.INFO, capture=None):
    """
    Copy line by line from file in f to the log from logger

    :param f: source stream object
    :param logger: the destination logger object
    :param loglevel: the level of logging data
    :param capture: an optional stream object for data copy
    """
    # Work-around for http://tracker.ceph.com/issues/8313
    if isinstance(f, ChannelFile):
        f._flags += ChannelFile.FLAG_BINARY

    for line in f:
        if capture:
            capture.write(line)
        line = line.rstrip()
        # Second part of work-around for http://tracker.ceph.com/issues/8313
        try:
            line = unicode(line, 'utf-8', 'replace').encode('utf-8')
            logger.log(loglevel, line.decode('utf-8'))
        except (UnicodeDecodeError, UnicodeEncodeError):
            logger.exception("Encountered unprintable line in command output")

def copy_and_close(src, fdst):
    """
    copyfileobj call wrapper.
    """
    if src is not None:
        if isinstance(src, basestring):
            src = StringIO(src)
        shutil.copyfileobj(src, fdst)
    fdst.close()

def copy_file_to(src, logger, stream=None):
    """
    Copy file
    :param src: file to be copied.
    :param logger: the logger object
    :param stream: an optional file-like object which will receive a copy of
                   src.
    """
    copy_to_log(src, logger, capture=stream)

def spawn_asyncresult(fn, *args, **kwargs):
    """
    Spawn a Greenlet and pass it's results to an AsyncResult.

    This function is useful to shuffle data from a Greenlet to
    AsyncResult, which then again is useful because any Greenlets that
    raise exceptions will cause tracebacks to be shown on stderr by
    gevent, even when ``.link_exception`` has been called. Using an
    AsyncResult avoids this.
    """
    r = gevent.event.AsyncResult()

    def wrapper():
        """
        Internal wrapper.
        """
        try:
            value = fn(*args, **kwargs)
        except Exception as e:
            r.set_exception(e)
        else:
            r.set(value)
    gevent.spawn(wrapper)

    return r

class Sentinel(object):

    """
    Sentinel -- used to define PIPE file-like object.
    """
    def __init__(self, name):
        self.name = name

    def __str__(self):
        return self.name

PIPE = Sentinel('PIPE')

class KludgeFile(object):

    """
    Wrap Paramiko's ChannelFile in a way that lets ``f.close()``
    actually cause an EOF for the remote command.
    """
    def __init__(self, wrapped):
        self._wrapped = wrapped

    def __getattr__(self, name):
        return getattr(self._wrapped, name)

    def close(self):
        """
        Close and shutdown.
        """
        self._wrapped.close()
        self._wrapped.channel.shutdown_write()

def run(
    client, args,
    stdin=None, stdout=None, stderr=None,
    logger=None,
    check_status=True,
    wait=True,
    name=None,
    label=None,
    timeout=None,
    cwd=None,
    # omit_sudo is used by vstart_runner.py
    omit_sudo=False
):
    """
    Run a command remotely.  If any of 'args' contains shell metacharacters
    that you want to pass unquoted, pass it as an instance of Raw(); otherwise
    it will be quoted with pipes.quote() (single quote, and single quotes
    enclosed in double quotes).

    :param client: SSHConnection to run the command with
    :param args: command to run
    :type args: list of string
    :param stdin: Standard input to send; either a string, a file-like object,
                  None, or `PIPE`. `PIPE` means caller is responsible for
                  closing stdin, or command may never exit.
    :param stdout: What to do with standard output. Either a file-like object,
                   a `logging.Logger`, `PIPE`, or `None` for copying to default
                   log. `PIPE` means caller is responsible for reading, or
                   command may never exit.
    :param stderr: What to do with standard error. See `stdout`.
    :param logger: If logging, write stdout/stderr to "out" and "err" children
                   of this logger. Defaults to logger named after this module.
    :param check_status: Whether to raise CommandFailedError on non-zero exit
                         status, and . Defaults to True. All signals and
                         connection loss are made to look like SIGHUP.
    :param wait: Whether to wait for process to exit. If False, returned
                 ``r.exitstatus`` s a `gevent.event.AsyncResult`, and the
                 actual status is available via ``.get()``.
    :param name: Human readable name (probably hostname) of the destination
                 host
    :param label: Can be used to label or describe what the command is doing.
    :param timeout: timeout value for args to complete on remote channel of
                    paramiko
    :param cwd: Directory in which the command should be executed.
    """
    try:
        transport = client.get_transport()
        if transport:
            (host, port) = transport.getpeername()[0:2]
        else:
            raise ConnectionLostError(command=quote(args), node=name)
    except socket.error:
        raise ConnectionLostError(command=quote(args), node=name)

    if name is None:
        name = host

    if timeout:
        log.info("Running command with timeout %d", timeout)
    r = RemoteProcess(client, args, check_status=check_status, hostname=name,
                      label=label, timeout=timeout, wait=wait, logger=logger,
                      cwd=cwd)
    r.execute()
    r.setup_stdin(stdin)
    r.setup_output_stream(stderr, 'stderr')
    r.setup_output_stream(stdout, 'stdout')
    if wait:
        r.wait()
    return r

def wait(processes, timeout=None):
    """
    Wait for all given processes to exit.

    Raise if any one of them fails.

    Optionally, timeout after 'timeout' seconds.
    """
    if timeout:
        log.info("waiting for %d", timeout)
    if timeout and timeout > 0:
        with safe_while(tries=(timeout / 6)) as check_time:
            not_ready = list(processes)
            while len(not_ready) > 0:
                check_time()
                for proc in list(not_ready):
                    if proc.finished:
                        not_ready.remove(proc)

    for proc in processes:
        proc.wait()
<commit_msg>orchestra/run: more concise logging when we run commands<commit_after>"""
Paramiko run support
"""
from cStringIO import StringIO
from paramiko import ChannelFile

import gevent
import gevent.event
import socket
import pipes
import logging
import shutil

from teuthology.contextutil import safe_while
from teuthology.exceptions import (CommandCrashedError, CommandFailedError,
                                   ConnectionLostError)

log = logging.getLogger(__name__)

class RemoteProcess(object):
    """
    An object to begin and monitor execution of a process on a remote host
    """
    __slots__ = [
        'client', 'args', 'check_status', 'command', 'hostname',
        'stdin', 'stdout', 'stderr',
        '_stdin_buf', '_stdout_buf', '_stderr_buf',
        'returncode', 'exitstatus', 'timeout',
        'greenlets',
        '_wait', 'logger',
        # for orchestra.remote.Remote to place a backreference
        'remote',
        'label',
        ]

    deadlock_warning = "Using PIPE for %s without wait=False would deadlock"

    def __init__(self, client, args, check_status=True, hostname=None,
                 label=None, timeout=None, wait=True, logger=None, cwd=None):
        """
        Create the object. Does not initiate command execution.

        :param client:       paramiko.SSHConnection to run the command with
        :param args:         Command to run.
        :type args:          String or list of strings
        :param check_status: Whether to raise CommandFailedError on non-zero
                             exit status, and . Defaults to True. All signals
                             and connection loss are made to look like SIGHUP.
        :param hostname:     Name of remote host (optional)
        :param label:        Can be used to label or describe what the
                             command is doing.
        :param timeout:      timeout value for arg that is passed to
                             exec_command of paramiko
        :param wait:         Whether self.wait() will be called automatically
        :param logger:       Alternative logger to use (optional)
        :param cwd:          Directory in which the command will be executed
                             (optional)
        """
        self.client = client
        self.args = args
        if isinstance(args, basestring):
            self.command = args
        else:
            self.command = quote(args)

        if cwd:
            self.command = '(cd {cwd} && exec {cmd})'.format(
                           cwd=cwd, cmd=self.command)

        self.check_status = check_status
        self.label = label
        if timeout:
            self.timeout = timeout
        if hostname:
            self.hostname = hostname
        else:
            (self.hostname, port) = client.get_transport().getpeername()[0:2]

        self.greenlets = []
        self.stdin, self.stdout, self.stderr = (None, None, None)
        self.returncode = self.exitstatus = None
        self._wait = wait
        self.logger = logger or log

    def execute(self):
        """
        Execute remote command
        """
        for line in self.command.split('\n'):
            log.getChild(self.hostname).info('%s> %s' % (self.label or '', line))

        if hasattr(self, 'timeout'):
            (self._stdin_buf, self._stdout_buf, self._stderr_buf) = \
                self.client.exec_command(self.command, timeout=self.timeout)
        else:
            (self._stdin_buf, self._stdout_buf, self._stderr_buf) = \
                self.client.exec_command(self.command)
        (self.stdin, self.stdout, self.stderr) = \
            (self._stdin_buf, self._stdout_buf, self._stderr_buf)

    def add_greenlet(self, greenlet):
        self.greenlets.append(greenlet)

    def setup_stdin(self, stream_obj):
        self.stdin = KludgeFile(wrapped=self.stdin)
        if stream_obj is not PIPE:
            greenlet = gevent.spawn(copy_and_close, stream_obj, self.stdin)
            self.add_greenlet(greenlet)
            self.stdin = None
        elif self._wait:
            # FIXME: Is this actually true?
            raise RuntimeError(self.deadlock_warning % 'stdin')

    def setup_output_stream(self, stream_obj, stream_name):
        if stream_obj is not PIPE:
            # Log the stream
            host_log = self.logger.getChild(self.hostname)
            stream_log = host_log.getChild(stream_name)
            self.add_greenlet(
                gevent.spawn(
                    copy_file_to,
                    getattr(self, stream_name),
                    stream_log,
                    stream_obj,
                )
            )
            setattr(self, stream_name, stream_obj)
        elif self._wait:
            # FIXME: Is this actually true?
            raise RuntimeError(self.deadlock_warning % stream_name)

    def wait(self):
        """
        Block until remote process finishes.

        :returns: self.returncode
        """

        status = self._get_exitstatus()
        if status != 0:
            log.debug("got remote process result: {}".format(status))
        for greenlet in self.greenlets:
            try:
                greenlet.get(block=True,timeout=60)
            except gevent.Timeout:
                log.debug("timed out waiting; will kill: {}".format(greenlet))
                greenlet.kill(block=False)
        for stream in ('stdout', 'stderr'):
            if hasattr(self, stream):
                stream_obj = getattr(self, stream)
                # Despite ChannelFile having a seek() method, it raises
                # "IOError: File does not support seeking."
                if hasattr(stream_obj, 'seek') and \
                        not isinstance(stream_obj, ChannelFile):
                    stream_obj.seek(0)

        self._raise_for_status()
        return status

    def _raise_for_status(self):
        if self.returncode is None:
            self._get_exitstatus()
        if self.check_status:
            if self.returncode in (None, -1):
                # command either died due to a signal, or the connection
                # was lost
                transport = self.client.get_transport()
                if transport is None or not transport.is_active():
                    # look like we lost the connection
                    raise ConnectionLostError(command=self.command,
                                              node=self.hostname)

                # connection seems healthy still, assuming it was a
                # signal; sadly SSH does not tell us which signal
                raise CommandCrashedError(command=self.command)
            if self.returncode != 0:
                raise CommandFailedError(
                    command=self.command, exitstatus=self.returncode,
                    node=self.hostname, label=self.label
                )

    def _get_exitstatus(self):
        """
        :returns: the remote command's exit status (return code). Note that
                  if the connection is lost, or if the process was killed by a
                  signal, this returns None instead of paramiko's -1.
        """
        status = self._stdout_buf.channel.recv_exit_status()
        self.exitstatus = self.returncode = status
        if status == -1:
            status = None
        return status

    @property
    def finished(self):
        gevent.wait(self.greenlets, timeout=0.1)
        ready = self._stdout_buf.channel.exit_status_ready()
        if ready:
            self._get_exitstatus()
        return ready

    def poll(self):
        """
        :returns: self.returncode if the process is finished; else None
        """
        if self.finished:
            self._raise_for_status()
            return self.returncode
        return None

    def __repr__(self):
        return '{classname}(client={client!r}, args={args!r}, check_status={check}, hostname={name!r})'.format(  # noqa
            classname=self.__class__.__name__,
            client=self.client,
            args=self.args,
            check=self.check_status,
            name=self.hostname,
            )

class Raw(object):

    """
    Raw objects are passed to remote objects and are not processed locally.
    """
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return '{cls}({value!r})'.format(
            cls=self.__class__.__name__,
            value=self.value,
            )

    def __eq__(self, value):
        return self.value == value

def quote(args):
    """
    Internal quote wrapper.
    """
    if isinstance(args, basestring):
        return args

    def _quote(args):
        """
        Handle quoted string, testing for raw charaters.
        """
        for a in args:
            if isinstance(a, Raw):
                yield a.value
            else:
                yield pipes.quote(a)
    return ' '.join(_quote(args))

def copy_to_log(f, logger, loglevel=logging.INFO, capture=None):
    """
    Copy line by line from file in f to the log from logger

    :param f: source stream object
    :param logger: the destination logger object
    :param loglevel: the level of logging data
    :param capture: an optional stream object for data copy
    """
    # Work-around for http://tracker.ceph.com/issues/8313
    if isinstance(f, ChannelFile):
        f._flags += ChannelFile.FLAG_BINARY

    for line in f:
        if capture:
            capture.write(line)
        line = line.rstrip()
        # Second part of work-around for http://tracker.ceph.com/issues/8313
        try:
            line = unicode(line, 'utf-8', 'replace').encode('utf-8')
            logger.log(loglevel, line.decode('utf-8'))
        except (UnicodeDecodeError, UnicodeEncodeError):
            logger.exception("Encountered unprintable line in command output")

def copy_and_close(src, fdst):
    """
    copyfileobj call wrapper.
    """
    if src is not None:
        if isinstance(src, basestring):
            src = StringIO(src)
        shutil.copyfileobj(src, fdst)
    fdst.close()

def copy_file_to(src, logger, stream=None):
    """
    Copy file
    :param src: file to be copied.
    :param logger: the logger object
    :param stream: an optional file-like object which will receive a copy of
                   src.
    """
    copy_to_log(src, logger, capture=stream)

def spawn_asyncresult(fn, *args, **kwargs):
    """
    Spawn a Greenlet and pass it's results to an AsyncResult.

    This function is useful to shuffle data from a Greenlet to
    AsyncResult, which then again is useful because any Greenlets that
    raise exceptions will cause tracebacks to be shown on stderr by
    gevent, even when ``.link_exception`` has been called. Using an
    AsyncResult avoids this.
    """
    r = gevent.event.AsyncResult()

    def wrapper():
        """
        Internal wrapper.
        """
        try:
            value = fn(*args, **kwargs)
        except Exception as e:
            r.set_exception(e)
        else:
            r.set(value)
    gevent.spawn(wrapper)

    return r

class Sentinel(object):

    """
    Sentinel -- used to define PIPE file-like object.
    """
    def __init__(self, name):
        self.name = name

    def __str__(self):
        return self.name

PIPE = Sentinel('PIPE')

class KludgeFile(object):

    """
    Wrap Paramiko's ChannelFile in a way that lets ``f.close()``
    actually cause an EOF for the remote command.
    """
    def __init__(self, wrapped):
        self._wrapped = wrapped

    def __getattr__(self, name):
        return getattr(self._wrapped, name)

    def close(self):
        """
        Close and shutdown.
        """
        self._wrapped.close()
        self._wrapped.channel.shutdown_write()

def run(
    client, args,
    stdin=None, stdout=None, stderr=None,
    logger=None,
    check_status=True,
    wait=True,
    name=None,
    label=None,
    timeout=None,
    cwd=None,
    # omit_sudo is used by vstart_runner.py
    omit_sudo=False
):
    """
    Run a command remotely.  If any of 'args' contains shell metacharacters
    that you want to pass unquoted, pass it as an instance of Raw(); otherwise
    it will be quoted with pipes.quote() (single quote, and single quotes
    enclosed in double quotes).

    :param client: SSHConnection to run the command with
    :param args: command to run
    :type args: list of string
    :param stdin: Standard input to send; either a string, a file-like object,
                  None, or `PIPE`. `PIPE` means caller is responsible for
                  closing stdin, or command may never exit.
    :param stdout: What to do with standard output. Either a file-like object,
                   a `logging.Logger`, `PIPE`, or `None` for copying to default
                   log. `PIPE` means caller is responsible for reading, or
                   command may never exit.
    :param stderr: What to do with standard error. See `stdout`.
    :param logger: If logging, write stdout/stderr to "out" and "err" children
                   of this logger. Defaults to logger named after this module.
    :param check_status: Whether to raise CommandFailedError on non-zero exit
                         status, and . Defaults to True. All signals and
                         connection loss are made to look like SIGHUP.
    :param wait: Whether to wait for process to exit. If False, returned
                 ``r.exitstatus`` s a `gevent.event.AsyncResult`, and the
                 actual status is available via ``.get()``.
    :param name: Human readable name (probably hostname) of the destination
                 host
    :param label: Can be used to label or describe what the command is doing.
    :param timeout: timeout value for args to complete on remote channel of
                    paramiko
    :param cwd: Directory in which the command should be executed.
    """
    try:
        transport = client.get_transport()
        if transport:
            (host, port) = transport.getpeername()[0:2]
        else:
            raise ConnectionLostError(command=quote(args), node=name)
    except socket.error:
        raise ConnectionLostError(command=quote(args), node=name)

    if name is None:
        name = host

    if timeout:
        log.info("Running command with timeout %d", timeout)
    r = RemoteProcess(client, args, check_status=check_status, hostname=name,
                      label=label, timeout=timeout, wait=wait, logger=logger,
                      cwd=cwd)
    r.execute()
    r.setup_stdin(stdin)
    r.setup_output_stream(stderr, 'stderr')
    r.setup_output_stream(stdout, 'stdout')
    if wait:
        r.wait()
    return r

def wait(processes, timeout=None):
    """
    Wait for all given processes to exit.

    Raise if any one of them fails.

    Optionally, timeout after 'timeout' seconds.
    """
    if timeout:
        log.info("waiting for %d", timeout)
    if timeout and timeout > 0:
        with safe_while(tries=(timeout / 6)) as check_time:
            not_ready = list(processes)
            while len(not_ready) > 0:
                check_time()
                for proc in list(not_ready):
                    if proc.finished:
                        not_ready.remove(proc)

    for proc in processes:
        proc.wait()
<|endoftext|>


<commit_before>
# Copyright (c) 2015, 2014 Computational Molecular Biology Group, Free University
# Berlin, 14195 Berlin, Germany.
# All rights reserved.
#
# Redistribution and use 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 HOLDER 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__ = 'noe'

import numpy as _np
import matplotlib.pylab as _plt

def plot_implied_timescales(ITS, ax=None, outfile=None, xlog=False, ylog=True, confidence=0.95, refs=None, units='steps',
                            dt=1.):
    r""" Generate a pretty implied timescale plot

    Parameters
    ----------
    ITS : implied timescales object.
        Object whose data will be plotted. Must provide the functions: get_timescales() and get_timescales(i) where i is the
        the property samples_available
    ax : matplotlib Axes object, optional, default = None
        the axes to plot to. When set to None the default Axes object will be used.
    outfile : str, optional, default = None
        output file to write the figure to. When not given, the plot will be displayed
    xlog : bool, optional, default = False
        Iff true, the x-Axis is logarithmized
    ylog : bool, optional, default = True
        Iff true, the y-Axis is logarithmized
    confidence : float, optional, default = 0.95
        The confidence interval for plotting error bars (if available)
    refs : ndarray((m), dtype=float), optional, default = None
        Reference (exact solution or other reference) timescales if known. The number of timescales must match those
        in the ITS object
    units: str, optional, default = 'steps'
        Affects the labeling of the axes. Used with :py:obj:`dt`, allows for changing the physical units of the axes.
        Accepts simple LaTeX math strings, eg. '$\mu$s'
    dt: float, optional, default = 1.0
        Physical meaning of the string in units. E.g, if you know that each frame corresponds to .010 ns, you can use the
        combination of parameters :py:obj:`dt` =0.01, :py:obj:`units` ='ns' to display the implied timescales
        in ns (instead of frames)

    Returns
    -------
    ax : Axes object containing the plot

    """
    # check input
    if (ax is None):
        ax = _plt.gca()
    colors = ['blue','red','green','cyan','purple','orange','violet']
    lags = ITS.lagtimes
    xmax = _np.max(lags)
    #ymin = min(_np.min(lags), _np.min(ITS.get_timescales()))
    #ymax = 1.5*_np.min(ITS.get_timescales())
    for i in range(ITS.number_of_timescales):
        # plot estimate
        ax.plot(lags*dt, ITS.get_timescales(process=i)*dt, color = colors[i % len(colors)])
        # sample available?
        if (ITS.samples_available and ITS.sample_number_of_timescales > i):
            # plot sample mean
            ax.plot(ITS.sample_lagtimes*dt, ITS.get_sample_mean(process=i)*dt, marker='o', color = colors[i % len(colors)], linestyle = 'dashed')
            (lconf, rconf) = ITS.get_sample_conf(confidence, i)
            ax.fill_between(ITS.sample_lagtimes*dt, lconf*dt, rconf*dt, alpha=0.2, color=colors[i % len(colors)])
        # reference available?
        if (refs is not None):
            tref = refs[i]
            ax.plot([0,min(tref,xmax)]*dt, [tref,tref]*dt, color='black', linewidth=1)
    # cutoff
    ax.fill_between(lags*dt, ax.get_ylim()[0]*_np.ones(len(lags))*dt, lags*dt, alpha=0.5, color='grey')
    ax.plot(lags*dt, lags*dt, linewidth=2, color='black')
    ax.set_xlim([1*dt,xmax*dt])
    #ax.set_ylim([ymin,ymax])
    # formatting
    ax.set_xlabel('lag time / %s'%units)
    ax.set_ylabel('timescale / %s'%units)
    if (xlog):
        ax.set_xscale('log')
    if (ylog):
        ax.set_yscale('log')

    # show or save
    # if outfile is None:
    #    _plt.show()
    if outfile is not None:
        _plt.savefig(outfile)

    return ax<commit_msg>minor changes<commit_after>
# Copyright (c) 2015, 2014 Computational Molecular Biology Group, Free University
# Berlin, 14195 Berlin, Germany.
# All rights reserved.
#
# Redistribution and use 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 HOLDER 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__ = 'noe'

import numpy as _np
import matplotlib.pylab as _plt

def plot_implied_timescales(ITS, ax=None, outfile=None, xlog=False, ylog=True, confidence=0.95, refs=None, units='steps',
                            dt=1.):
    r""" Generate a pretty implied timescale plot

    Parameters
    ----------
    ITS : implied timescales object.
        Object whose data will be plotted. Must provide the functions: get_timescales() and get_timescales(i) where i is the
        the property samples_available
    ax : matplotlib Axes object, optional, default = None
        the axes to plot to. When set to None the default Axes object will be used.
    outfile : str, optional, default = None
        output file to write the figure to. When not given, the plot will be displayed
    xlog : bool, optional, default = False
        Iff true, the x-Axis is logarithmized
    ylog : bool, optional, default = True
        Iff true, the y-Axis is logarithmized
    confidence : float, optional, default = 0.95
        The confidence interval for plotting error bars (if available)
    refs : ndarray((m), dtype=float), optional, default = None
        Reference (exact solution or other reference) timescales if known. The number of timescales must match those
        in the ITS object
    units: str, optional, default = 'steps'
        Affects the labeling of the axes. Used with :py:obj:`dt`, allows for changing the physical units of the axes.
        Accepts simple LaTeX math strings, eg. '$\mu$s'
    dt: float, optional, default = 1.0
        Physical time between frames, expressed the units given in :py:obj:`units`. E.g, if you know that each
        frame corresponds to .010 ns, you can use the combination of parameters :py:obj:`dt` =0.01,
        :py:obj:`units` ='ns' to display the implied timescales in ns (instead of frames)

    Returns
    -------
    ax : Axes object containing the plot

    """
    # check input
    if (ax is None):
        ax = _plt.gca()
    colors = ['blue','red','green','cyan','purple','orange','violet']
    lags = ITS.lagtimes
    xmax = _np.max(lags)
    #ymin = min(_np.min(lags), _np.min(ITS.get_timescales()))
    #ymax = 1.5*_np.min(ITS.get_timescales())
    for i in range(ITS.number_of_timescales):
        # plot estimate
        ax.plot(lags*dt, ITS.get_timescales(process=i)*dt, color = colors[i % len(colors)])
        # sample available?
        if (ITS.samples_available and ITS.sample_number_of_timescales > i):
            # plot sample mean
            ax.plot(ITS.sample_lagtimes*dt, ITS.get_sample_mean(process=i)*dt, marker='o', color = colors[i % len(colors)], linestyle = 'dashed')
            (lconf, rconf) = ITS.get_sample_conf(confidence, i)
            ax.fill_between(ITS.sample_lagtimes*dt, lconf*dt, rconf*dt, alpha=0.2, color=colors[i % len(colors)])
        # reference available?
        if (refs is not None):
            tref = refs[i]
            ax.plot([0,min(tref,xmax)]*dt, [tref,tref]*dt, color='black', linewidth=1)
    # cutoff
    ax.fill_between(lags*dt, ax.get_ylim()[0]*_np.ones(len(lags))*dt, lags*dt, alpha=0.5, color='grey')
    ax.plot(lags*dt, lags*dt, linewidth=2, color='black')
    ax.set_xlim([1*dt,xmax*dt])
    #ax.set_ylim([ymin,ymax])
    # formatting
    ax.set_xlabel('lag time / %s'%units)
    ax.set_ylabel('timescale / %s'%units)
    if (xlog):
        ax.set_xscale('log')
    if (ylog):
        ax.set_yscale('log')

    # show or save
    # if outfile is None:
    #    _plt.show()
    if outfile is not None:
        _plt.savefig(outfile)

    return ax<|endoftext|>


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

import sys
import io
import os

import time
from time import sleep

import numpy
import matplotlib
from PIL import Image
from PyQt5.QtCore import *
import PyQt5.QtCore
from PyQt5.QtGui import *

from PyQt5.QtWidgets import *

from pyqt import FreakingQtImageViewer
import scipy.misc

from skimage import color
from skimage import transform as tf
#from skimage import data
from skimage.feature import canny
from skimage.transform import hough_ellipse
from skimage.draw import ellipse_perimeter

from skimage.transform import hough_circle
from skimage.feature import peak_local_max
from skimage.draw import circle_perimeter
from skimage.util import img_as_ubyte

import scipy.misc

import picamera
import picamera.array

from cv import circle_detection
from _thread import start_new_thread

class DraughtsGameWindow(QWidget):

	def __init__(self):
		super().__init__()
		self.corners = []
		self.CELLS_PER_ROW = 8
		self.BLACK_COLOR = QColor(0,0,0,255)
		self.WHITE_COLOR = QColor(255,255,255,255)
		self.BLUE_COLOR = QColor(0,0,255,255)
		self.RED_COLOR = QColor(255,0,0,255)
		self.GREEN_COLOR = QColor(0,255,0,255)
		self.DEFAULT_PEN = QPen(self.BLUE_COLOR)
		self.DEFAULT_PEN.setWidth(20)
		self.PIX_PADDING = 25
		self.BORDER_RADIUS = 10

		self.rg_bg = (5.5, 5.5)

		self.initUI()

	def keyPressEvent(self, event):
		if event.key() == Qt.Key_Escape:
			self.close()
		elif event.key() == Qt.Key_C:
			start_new_thread(self.calibrate, (0, 0,))
		elif event.key() == Qt.Key_V:

			self.vbox.setContentsMargins(self.upper_left[0] + 10,self.PIX_PADDING + 10,240,112)
			self.pixmap = QPixmap("checkers_randlos.jpg")
			self.pixmap = self.pixmap.scaledToWidth( self.upper_right[0]-self.upper_left[0] )
			'''
			self.painter.drawEllipse(self.upper_left[0]-10,self.upper_left[1]-10,20,20)
			self.painter = QPainter(self.pixmap)
			self.DEFAULT_PEN.setColor(self.WHITE_COLOR)
			self.painter.setPen(self.DEFAULT_PEN)
			self.painter.drawEllipse(self.upper_right[0]-10,self.upper_right[1]-10,20,20)
			self.DEFAULT_PEN.setColor(self.RED_COLOR)
			self.painter.setPen(self.DEFAULT_PEN)
			self.painter.drawEllipse(self.lower_right[0]-10,self.lower_right[1]-10,20,20)
			self.DEFAULT_PEN.setColor(self.GREEN_COLOR)
			self.painter.setPen(self.DEFAULT_PEN)
			self.painter.drawEllipse(self.lower_left[0]-10,self.lower_left[1]-10,20,20)
			'''

			self.lbl.setPixmap(self.pixmap)
			self.lbl.repaint()

		elif event.key() == Qt.Key_S:
			start_new_thread(self.findTiles, (0, 0,))

		elif event.key() == Qt.Key_P:
			camera = picamera.PiCamera()

			# test manual white balance
			camera.awb_mode = 'off'
			camera.awb_gains = self.rg_bg
			# end

			with picamera.array.PiRGBArray(camera) as stream:
				camera.capture(stream, format='rgb')
				im = Image.fromarray(stream.array)
				im.save('./preview.png')
				camera.close()
				os.system('xdg-open ./preview.png')

	def initUI(self):

		self.WIDTH, self.HEIGHT, self.CELL_SIZE = screen_rect.width(), screen_rect.height(), screen_rect.width() / self.CELLS_PER_ROW
		self.pixmap = QPixmap(self.WIDTH-20,self.HEIGHT-20)

		self.upper_left = (self.PIX_PADDING + self.WIDTH / 2 - self.pixmap.height() / 2, self.PIX_PADDING)
		self.upper_right = (- self.PIX_PADDING + self.WIDTH / 2 + self.pixmap.height() / 2, self.PIX_PADDING + 0)
		self.lower_right = (-self.PIX_PADDING + self.WIDTH / 2 + self.pixmap.height() / 2, -self.PIX_PADDING + self.pixmap.height())
		self.lower_left = (self.PIX_PADDING + self.WIDTH / 2 - self.pixmap.height() / 2, -self.PIX_PADDING + self.pixmap.height())

		self.DRAW_SIZE = self.upper_right[0] - self.upper_left[0]
		print(self.DRAW_SIZE)

		self.painter = QPainter(self.pixmap)
		self.painter.setPen(self.DEFAULT_PEN)

		self.painter.fillRect(0, 0,self.WIDTH, self.HEIGHT, self.BLACK_COLOR)
		self.painter.drawEllipse(self.upper_left[0]-10,self.upper_left[1]-10,20,20)
		self.DEFAULT_PEN.setColor(self.WHITE_COLOR)
		self.painter.setPen(self.DEFAULT_PEN)
		self.painter.drawEllipse(self.upper_right[0]-10,self.upper_right[1]-10,20,20)
		self.DEFAULT_PEN.setColor(self.RED_COLOR)
		self.painter.setPen(self.DEFAULT_PEN)
		self.painter.drawEllipse(self.lower_right[0]-10,self.lower_right[1]-10,20,20)
		self.DEFAULT_PEN.setColor(self.GREEN_COLOR)
		self.painter.setPen(self.DEFAULT_PEN)
		self.painter.drawEllipse(self.lower_left[0]-10,self.lower_left[1]-10,20,20)

		self.vbox = QVBoxLayout(self)
		self.lbl = QLabel(self)
		self.lbl.setPixmap(self.pixmap)
		self.vbox.addWidget(self.lbl)

		self.setStyleSheet("background-color:black;")

		self.setLayout(self.vbox)
		self.vbox.setAlignment(Qt.AlignHCenter)
		self.showFullScreen()
		self.setWindowTitle('DraughtsCV')
		self.show()
		#calibrate()

	def checkCoords(self, img, coords):
		self.redC = []
		self.greenC = []
		self.blueC = []
		self.whiteC = []
		for index, coord in enumerate(coords):
			#print(img[coord[0], coord[1]]," - ", )
			pixel = img[coord[0], coord[1]]
			sum = 0
			sum += pixel[0]
			sum += pixel[1]
			sum += pixel[2]
			'''if sum > 300:
				print()
				#print("delete ", coord, ' with ', pixel)
			else:'''
			print("pixel: ", pixel, "@",coord)
			if pixel[0] > 150 and pixel[1] > 150 and pixel[2] > 150:
				self.whiteC.append(coord)
			elif pixel[0] > 110:
				self.redC.append(coord)
			elif pixel[1] > 110:
				self.greenC.append(coord)
			elif pixel[2] > 110:
				self.blueC.append(coord)

		#final_red = (int(numpy.sum(self.redC[0])/len(self.redC)),int(numpy.sum(self.redC[1])/len(self.redC[])))	#(sum(self.redC[0]), sum(self.redC[1]))
		final_white = numpy.average(self.whiteC, axis=0)
		final_red = numpy.average(self.redC, axis=0)
		final_green = numpy.average(self.greenC, axis=0)
		final_blue = numpy.average(self.blueC, axis=0)

		self.lbl.setPixmap(self.pixmap)
		if len(self.whiteC) == 0 or len(self.blueC) == 0 or len(self.greenC) == 0 or len(self.redC) == 0 :
			print('less than 4 corners for calibration detected, returning -1')
			return -1

		final_white = (int(final_white[1]),int(final_white[0]))
		final_red = (int(final_red[1]),int(final_red[0]))
		final_green = (int(final_green[1]),int(final_green[0]))
		final_blue = (int(final_blue[1]),int(final_blue[0]))

		print('avg blue ',final_blue)
		print('avg green',final_green)
		print('avg red ',final_red)
		print('avg white ',final_white)

		return numpy.array((final_blue, final_green, final_red, final_white))
		#return finalCoords

	def calibrate(self, ignore1, ignore2): # pythonmaster
		camera = picamera.PiCamera()

		# test manual white balance
		camera.awb_mode = 'off'
		camera.awb_gains = self.rg_bg
		# end

		self.calibrationPoints = -1
		#while True:
		with picamera.array.PiRGBArray(camera) as stream:
			camera.capture(stream, format='rgb')
			img = stream.array
			camera.close()

			img, coords, circleDebug = circle_detection(img, 20, 25)
			im = Image.fromarray(img) #.convert('LA')
			input_img = im
			imDeb = Image.fromarray(circleDebug) #.convert('LA')
			im.save('./tmp.png')
			imDeb.save('./deb.png')

			self.calibrationPoints = self.checkCoords(img, coords)
			if isinstance(self.calibrationPoints, (numpy.ndarray, numpy.generic) ):
				src = numpy.array((
					(0, 0), #upper left
					(0, 800), #lower left
					(800, 800), #lright
					(800, 0) #uright
				))

				transformer = tf.ProjectiveTransform()
				transformer.estimate(src, self.calibrationPoints)
				transformed_image = tf.warp(input_img, transformer, output_shape = (800,800))

				#print("Transformed Image: ",transformed_image)

				scipy.misc.imsave('./transformed.png', transformed_image)
				#break
			else: print('ups')
		print(self.corners)

	def findTiles(self, ignore1, ignore2): # pythonmaster

		if isinstance(self.calibrationPoints, (numpy.ndarray, numpy.generic) ):

			camera = picamera.PiCamera()

			# test manual white balance
			camera.awb_mode = 'off'
			camera.awb_gains = self.rg_bg
			# end

			with picamera.array.PiRGBArray(camera) as stream:
				camera.capture(stream, format='rgb')
				img = stream.array
				camera.close()

				src = numpy.array((
					(0, 0), #upper left
					(0, 800), #lower left
					(800, 800), #lright
					(800, 0) #uright
				))

				transformer = tf.ProjectiveTransform()
				transformer.estimate(src, self.calibrationPoints)
				board_image = tf.warp(img, transformer, output_shape = (800,800))

				scipy.misc.imsave('./board.png', board_image)

			print('board image captured')

		else:
			print('calibrationPoints not yet set')

if __name__ == '__main__':

	app = QApplication(sys.argv)

	screen_rect = app.desktop().screenGeometry()
	viewer = DraughtsGameWindow() # init ui
	sys.exit(app.exec_())<commit_msg>And it goes on ...<commit_after>#!/usr/bin/env python3

import sys
import io
import os

import time
from time import sleep

import numpy
import matplotlib
from PIL import Image
from PyQt5.QtCore import *
import PyQt5.QtCore
from PyQt5.QtGui import *

from PyQt5.QtWidgets import *

from pyqt import FreakingQtImageViewer
import scipy.misc

from skimage import color
from skimage import transform as tf
#from skimage import data
from skimage.feature import canny
from skimage.transform import hough_ellipse
from skimage.draw import ellipse_perimeter

from skimage.transform import hough_circle
from skimage.feature import peak_local_max
from skimage.draw import circle_perimeter
from skimage.util import img_as_ubyte

import scipy.misc

import picamera
import picamera.array

from cv import circle_detection
from _thread import start_new_thread

class DraughtsGameWindow(QWidget):

	def __init__(self):
		super().__init__()
		self.corners = []
		self.CELLS_PER_ROW = 8
		self.BLACK_COLOR = QColor(0,0,0,255)
		self.WHITE_COLOR = QColor(255,255,255,255)
		self.BLUE_COLOR = QColor(0,0,255,255)
		self.RED_COLOR = QColor(255,0,0,255)
		self.GREEN_COLOR = QColor(0,255,0,255)
		self.DEFAULT_PEN = QPen(self.BLUE_COLOR)
		self.DEFAULT_PEN.setWidth(20)
		self.PIX_PADDING = 25
		self.BORDER_RADIUS = 10

		self.rg_bg = (4, 4)

		self.initUI()

	def keyPressEvent(self, event):
		if event.key() == Qt.Key_Escape:
			self.close()
		elif event.key() == Qt.Key_C:
			start_new_thread(self.calibrate, (0, 0,))
		elif event.key() == Qt.Key_V:

			self.vbox.setContentsMargins(self.upper_left[0] + 10,self.PIX_PADDING + 10,240,112)
			self.pixmap = QPixmap("checkers_randlos.jpg")
			self.pixmap = self.pixmap.scaledToWidth( self.upper_right[0]-self.upper_left[0] )
			'''
			self.painter.drawEllipse(self.upper_left[0]-10,self.upper_left[1]-10,20,20)
			self.painter = QPainter(self.pixmap)
			self.DEFAULT_PEN.setColor(self.WHITE_COLOR)
			self.painter.setPen(self.DEFAULT_PEN)
			self.painter.drawEllipse(self.upper_right[0]-10,self.upper_right[1]-10,20,20)
			self.DEFAULT_PEN.setColor(self.RED_COLOR)
			self.painter.setPen(self.DEFAULT_PEN)
			self.painter.drawEllipse(self.lower_right[0]-10,self.lower_right[1]-10,20,20)
			self.DEFAULT_PEN.setColor(self.GREEN_COLOR)
			self.painter.setPen(self.DEFAULT_PEN)
			self.painter.drawEllipse(self.lower_left[0]-10,self.lower_left[1]-10,20,20)
			'''

			self.lbl.setPixmap(self.pixmap)
			self.lbl.repaint()

		elif event.key() == Qt.Key_S:
			start_new_thread(self.findTiles, (0, 0,))

		elif event.key() == Qt.Key_P:
			camera = picamera.PiCamera()

			# test manual white balance
			camera.awb_mode = 'off'
			camera.awb_gains = self.rg_bg
			# end

			with picamera.array.PiRGBArray(camera) as stream:
				camera.capture(stream, format='rgb')
				im = Image.fromarray(stream.array)
				im.save('./preview.png')
				camera.close()
				os.system('xdg-open ./preview.png')

	def initUI(self):

		self.WIDTH, self.HEIGHT, self.CELL_SIZE = screen_rect.width(), screen_rect.height(), screen_rect.width() / self.CELLS_PER_ROW
		self.pixmap = QPixmap(self.WIDTH-20,self.HEIGHT-20)

		self.upper_left = (self.PIX_PADDING + self.WIDTH / 2 - self.pixmap.height() / 2, self.PIX_PADDING)
		self.upper_right = (- self.PIX_PADDING + self.WIDTH / 2 + self.pixmap.height() / 2, self.PIX_PADDING + 0)
		self.lower_right = (-self.PIX_PADDING + self.WIDTH / 2 + self.pixmap.height() / 2, -self.PIX_PADDING + self.pixmap.height())
		self.lower_left = (self.PIX_PADDING + self.WIDTH / 2 - self.pixmap.height() / 2, -self.PIX_PADDING + self.pixmap.height())

		self.DRAW_SIZE = self.upper_right[0] - self.upper_left[0]
		print(self.DRAW_SIZE)

		self.painter = QPainter(self.pixmap)
		self.painter.setPen(self.DEFAULT_PEN)

		self.painter.fillRect(0, 0,self.WIDTH, self.HEIGHT, self.BLACK_COLOR)
		self.painter.drawEllipse(self.upper_left[0]-10,self.upper_left[1]-10,20,20)
		self.DEFAULT_PEN.setColor(self.WHITE_COLOR)
		self.painter.setPen(self.DEFAULT_PEN)
		self.painter.drawEllipse(self.upper_right[0]-10,self.upper_right[1]-10,20,20)
		self.DEFAULT_PEN.setColor(self.RED_COLOR)
		self.painter.setPen(self.DEFAULT_PEN)
		self.painter.drawEllipse(self.lower_right[0]-10,self.lower_right[1]-10,20,20)
		self.DEFAULT_PEN.setColor(self.GREEN_COLOR)
		self.painter.setPen(self.DEFAULT_PEN)
		self.painter.drawEllipse(self.lower_left[0]-10,self.lower_left[1]-10,20,20)

		self.vbox = QVBoxLayout(self)
		self.lbl = QLabel(self)
		self.lbl.setPixmap(self.pixmap)
		self.vbox.addWidget(self.lbl)

		self.setStyleSheet("background-color:black;")

		self.setLayout(self.vbox)
		self.vbox.setAlignment(Qt.AlignHCenter)
		self.showFullScreen()
		self.setWindowTitle('DraughtsCV')
		self.show()
		#calibrate()

	def checkCoords(self, img, coords):
		self.redC = []
		self.greenC = []
		self.blueC = []
		self.whiteC = []
		for index, coord in enumerate(coords):
			#print(img[coord[0], coord[1]]," - ", )
			pixel = img[coord[0], coord[1]]
			sum = 0
			sum += pixel[0]
			sum += pixel[1]
			sum += pixel[2]
			'''if sum > 300:
				print()
				#print("delete ", coord, ' with ', pixel)
			else:'''
			print("pixel: ", pixel, "@",coord)
			if pixel[0] > 150 and pixel[1] > 150 and pixel[2] > 150:
				self.whiteC.append(coord)
			elif pixel[0] > 110:
				self.redC.append(coord)
			elif pixel[1] > 110:
				self.greenC.append(coord)
			elif pixel[2] > 110:
				self.blueC.append(coord)

		#final_red = (int(numpy.sum(self.redC[0])/len(self.redC)),int(numpy.sum(self.redC[1])/len(self.redC[])))	#(sum(self.redC[0]), sum(self.redC[1]))
		final_white = numpy.average(self.whiteC, axis=0)
		final_red = numpy.average(self.redC, axis=0)
		final_green = numpy.average(self.greenC, axis=0)
		final_blue = numpy.average(self.blueC, axis=0)

		self.lbl.setPixmap(self.pixmap)
		if len(self.whiteC) == 0 or len(self.blueC) == 0 or len(self.greenC) == 0 or len(self.redC) == 0 :
			print('less than 4 corners for calibration detected, returning -1')
			return -1

		final_white = (int(final_white[1]),int(final_white[0]))
		final_red = (int(final_red[1]),int(final_red[0]))
		final_green = (int(final_green[1]),int(final_green[0]))
		final_blue = (int(final_blue[1]),int(final_blue[0]))

		print('avg blue ',final_blue)
		print('avg green',final_green)
		print('avg red ',final_red)
		print('avg white ',final_white)

		return numpy.array((final_blue, final_green, final_red, final_white))
		#return finalCoords

	def calibrate(self, ignore1, ignore2): # pythonmaster
		camera = picamera.PiCamera()

		# test manual white balance
		camera.awb_mode = 'off'
		camera.awb_gains = self.rg_bg
		# end

		self.calibrationPoints = -1
		#while True:
		with picamera.array.PiRGBArray(camera) as stream:
			camera.capture(stream, format='rgb')
			img = stream.array
			camera.close()

			img, coords, circleDebug = circle_detection(img, 20, 25)
			im = Image.fromarray(img) #.convert('LA')
			input_img = im
			imDeb = Image.fromarray(circleDebug) #.convert('LA')
			im.save('./tmp.png')
			imDeb.save('./deb.png')

			self.calibrationPoints = self.checkCoords(img, coords)
			if isinstance(self.calibrationPoints, (numpy.ndarray, numpy.generic) ):
				src = numpy.array((
					(0, 0), #upper left
					(0, 800), #lower left
					(800, 800), #lright
					(800, 0) #uright
				))

				transformer = tf.ProjectiveTransform()
				transformer.estimate(src, self.calibrationPoints)
				transformed_image = tf.warp(input_img, transformer, output_shape = (800,800))

				#print("Transformed Image: ",transformed_image)

				scipy.misc.imsave('./transformed.png', transformed_image)
				#break
			else: print('ups')
		print(self.corners)

	def findTiles(self, ignore1, ignore2): # pythonmaster

		if isinstance(self.calibrationPoints, (numpy.ndarray, numpy.generic) ):

			camera = picamera.PiCamera()

			# test manual white balance
			camera.awb_mode = 'off'
			camera.awb_gains = self.rg_bg
			# end

			with picamera.array.PiRGBArray(camera) as stream:
				camera.capture(stream, format='rgb')
				img = stream.array
				camera.close()

				src = numpy.array((
					(0, 0), #upper left
					(0, 800), #lower left
					(800, 800), #lright
					(800, 0) #uright
				))

				transformer = tf.ProjectiveTransform()
				transformer.estimate(src, self.calibrationPoints)
				board_image = tf.warp(img, transformer, output_shape = (800,800))

				scipy.misc.imsave('./board.png', board_image)

			print('board image captured')

		else:
			print('calibrationPoints not yet set')

if __name__ == '__main__':

	app = QApplication(sys.argv)

	screen_rect = app.desktop().screenGeometry()
	viewer = DraughtsGameWindow() # init ui
	sys.exit(app.exec_())<|endoftext|>


<commit_before>import pygame
from colors import *
from hitDetection import HitDetector
from random import randint

Running, Player1Win, Player2Win = range(3)
Keys = [pygame.K_DOWN, pygame.K_UP]

class PongPos(object):
    """docstring for PongPos"""
    def __init__(self, x=0, y=0):
        super(PongPos, self).__init__()
        self.x = x
        self.y = y

class PongSize(object):
    """docstring for PongSize"""
    def __init__(self, width=0, height=0):
        super(PongSize, self).__init__()
        self.width = width
        self.height = height

class PongObject(object):
    """docstring for PongObject"""
    def __init__(self, name, pos=PongPos(), size=PongSize(), color=black):
        super(PongObject, self).__init__()
        self.name = name
        self.pos = PongPos(pos.x, pos.y)
        self.history = []
        self.size = size
        self.color = color

    def left(self):
        return self.pos.x

    def right(self):
        return self.pos.x + self.size.width

    def top(self):
        return self.pos.y

    def bottom(self):
        return self.pos.y + self.size.height

class Pad(PongObject):
    """docstring for Pad"""
    defaultHeight = 100
    defaultWidth = 10

    def __init__(self, name, pos=PongPos()):
        super(Pad, self).__init__(name, pos, PongSize(Pad.defaultWidth, Pad.defaultHeight), red)
        self.velocity = PongPos()
        self.speed = 10

    def startMove(self, directionDown):
        if directionDown:
            self.velocity.y = -self.speed
        else:
            self.velocity.y = self.speed

    def endMove(self):
        self.velocity.y = 0

class Ball(PongObject):
    """docstring for Ball"""
    defaultSize = 20

    def __init__(self, pos=PongPos()):
        super(Ball, self).__init__("ball", pos, PongSize(Ball.defaultSize, Ball.defaultSize), green)

        # Randomly choose one of four starting directions
        self.directions = [(-1, -1), (-1, 1), (1, -1), (1, 1)]
        self.speed = 10
        self.restart(pos)

    def restart(self, pos=PongPos()):
        i = randint(0, 3)
        self.velocity = PongPos(self.speed*self.directions[i][0], self.speed*self.directions[i][1])
        self.pos = PongPos(pos.x, pos.y)

class PongLogic(object):
    """docstring for PongLogic"""
    def __init__(self, width, height):
        super(PongLogic, self).__init__()

        self.status = Running

        self.width = width
        self.height = height

        self.ballSize = Ball.defaultSize
        self.ballStart = PongPos((self.width - self.ballSize) / 2,
                                 (self.height - self.ballSize) / 2)

        self.padStartOffsets = PongPos(50, 10)

        self.pad1Start = self.padStartOffsets
        self.pad2Start = PongPos(self.width - self.padStartOffsets.x - Pad.defaultWidth,
                                 self.height - self.padStartOffsets.y - Pad.defaultHeight)

        self.ball = Ball(self.ballStart)
        self.pad1 = Pad("1", self.pad1Start)
        self.pad2 = Pad("2", self.pad2Start)

        self.score = [0, 0]
        self.scoreMax = 10

        self.hitDetector = HitDetector(0, 0, 600, 400)
        self.pads = [self.pad1, self.pad2]

    def pressKey(self, key):
        for i in range(len(Keys)):
            if key == Keys[i]:
                self.pad1.startMove(i)
                break

    def releaseKey(self, key):
        for i in range(len(Keys)):
            if key == Keys[i]:
                self.pad1.endMove()
                break

    def update(self):
        if self.status == Running:
            if self.hitDetector.didHitXBoundary(self.ball):
                print("Boundary X Hit")
                if self.width - self.ball.pos.x <= Ball.defaultSize:
                    self.score[0] += 1
                    if self.score[0] >= self.scoreMax:
                        self.status = Player1Win
                        print("Player 1 Wins")
                else:
                    self.score[1] += 1
                    if self.score[1] >= self.scoreMax:
                        self.status = Player2Win
                        print("Player 2 Wins")
                print("Player 1: " + str(self.score[0]) + ", Player 2: " + str(self.score[1]))
                self.ball.restart(self.ballStart)

            else:
                if self.hitDetector.collision(self.ball, self.pad1) or self.hitDetector.collision(self.ball, self.pad2):
                    self.ball.velocity.x = -self.ball.velocity.x
                    print("Pad Hit")
                elif self.hitDetector.didHitYBoundary(self.ball):
                    self.ball.velocity.y = -self.ball.velocity.y
                    print("Boundary Y Hit")

                self.ball.pos.x += self.ball.velocity.x
                self.ball.pos.y += self.ball.velocity.y

            if self.hitDetector.didHitYBoundary(self.pad1):
                print("Pad1 boundary Y hit")
                self.pad1.pos.y += -self.pad1.velocity.y
                self.pad1.endMove()
            else:
                self.pad1.pos.y += self.pad1.velocity.y
<commit_msg>Can play as either player (or both), but there are still bugs in collision detection<commit_after>import pygame
from colors import *
from hitDetection import HitDetector
from random import randint

Running, Player1Win, Player2Win = range(3)
Keys = [pygame.K_s, pygame.K_w, pygame.K_DOWN, pygame.K_UP]

class PongPos(object):
    """docstring for PongPos"""
    def __init__(self, x=0, y=0):
        super(PongPos, self).__init__()
        self.x = x
        self.y = y

class PongSize(object):
    """docstring for PongSize"""
    def __init__(self, width=0, height=0):
        super(PongSize, self).__init__()
        self.width = width
        self.height = height

class PongObject(object):
    """docstring for PongObject"""
    def __init__(self, name, pos=PongPos(), size=PongSize(), color=black):
        super(PongObject, self).__init__()
        self.name = name
        self.pos = PongPos(pos.x, pos.y)
        self.history = []
        self.size = size
        self.color = color

    def left(self):
        return self.pos.x

    def right(self):
        return self.pos.x + self.size.width

    def top(self):
        return self.pos.y

    def bottom(self):
        return self.pos.y + self.size.height

class Pad(PongObject):
    """docstring for Pad"""
    defaultHeight = 100
    defaultWidth = 10

    def __init__(self, name, pos=PongPos()):
        super(Pad, self).__init__(name, pos, PongSize(Pad.defaultWidth, Pad.defaultHeight), red)
        self.velocity = PongPos()
        self.speed = 10

    def startMove(self, directionDown):
        if directionDown:
            self.velocity.y = -self.speed
        else:
            self.velocity.y = self.speed

    def endMove(self):
        self.velocity.y = 0

class Ball(PongObject):
    """docstring for Ball"""
    defaultSize = 20

    def __init__(self, pos=PongPos()):
        super(Ball, self).__init__("ball", pos, PongSize(Ball.defaultSize, Ball.defaultSize), green)

        # Randomly choose one of four starting directions
        self.directions = [(-1, -1), (-1, 1), (1, -1), (1, 1)]
        self.speed = 10
        self.restart(pos)

    def restart(self, pos=PongPos()):
        i = randint(0, 3)
        self.velocity = PongPos(self.speed*self.directions[i][0], self.speed*self.directions[i][1])
        self.pos = PongPos(pos.x, pos.y)

class PongLogic(object):
    """docstring for PongLogic"""
    def __init__(self, width, height):
        super(PongLogic, self).__init__()

        self.status = Running

        self.width = width
        self.height = height

        self.ballSize = Ball.defaultSize
        self.ballStart = PongPos((self.width - self.ballSize) / 2,
                                 (self.height - self.ballSize) / 2)

        self.padStartOffsets = PongPos(50, 10)

        self.pad1Start = self.padStartOffsets
        self.pad2Start = PongPos(self.width - self.padStartOffsets.x - Pad.defaultWidth,
                                 self.height - self.padStartOffsets.y - Pad.defaultHeight)

        self.ball = Ball(self.ballStart)
        self.pad1 = Pad("1", self.pad1Start)
        self.pad2 = Pad("2", self.pad2Start)

        self.score = [0, 0]
        self.scoreMax = 10

        self.hitDetector = HitDetector(0, 0, 600, 400)
        self.pads = [self.pad1, self.pad2]

    def pressKey(self, key):
        for i in range(len(Keys)):
            if key == Keys[i]:
                if i < 2:
                    self.pad1.startMove(i)
                else:
                    self.pad2.startMove(i - 2)
                break

    def releaseKey(self, key):
        for i in range(len(Keys)):
            if key == Keys[i]:
                if i < 2:
                    self.pad1.endMove()
                else:
                    self.pad2.endMove()
                break

    def update(self):
        if self.status == Running:
            if self.hitDetector.didHitXBoundary(self.ball):
                print("Boundary X Hit")
                if self.width - self.ball.pos.x <= Ball.defaultSize:
                    self.score[0] += 1
                    if self.score[0] >= self.scoreMax:
                        self.status = Player1Win
                        print("Player 1 Wins")
                else:
                    self.score[1] += 1
                    if self.score[1] >= self.scoreMax:
                        self.status = Player2Win
                        print("Player 2 Wins")
                print("Player 1: " + str(self.score[0]) + ", Player 2: " + str(self.score[1]))
                self.ball.restart(self.ballStart)

            else:
                if self.hitDetector.collision(self.ball, self.pad1) or self.hitDetector.collision(self.ball, self.pad2):
                    self.ball.velocity.x = -self.ball.velocity.x
                    print("Pad Hit")
                elif self.hitDetector.didHitYBoundary(self.ball):
                    self.ball.velocity.y = -self.ball.velocity.y
                    print("Boundary Y Hit")

                self.ball.pos.x += self.ball.velocity.x
                self.ball.pos.y += self.ball.velocity.y

            for pad in self.pads:
                self.updatePad(pad)

    def updatePad(self, pad):
        if self.hitDetector.didHitYBoundary(pad):
            print(pad.name + " boundary Y hit")
            if pad.pos.y + pad.size.height >= self.height:
                pad.pos.y -= pad.speed
            else:
                pad.pos.y += pad.speed
        else:
            pad.pos.y += pad.velocity.y
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
    Stanza Processor

    This module starts a stanza processor that is responsible for
    processing all stanzas addressed to the local domain(s)

    :copyright: 2011 by the pyfire Team, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""

import cPickle
import zmq
from zmq.eventloop import ioloop, zmqstream
import xml.etree.ElementTree as ET

from pyfire.logger import Logger
from pyfire import configuration as config
from pyfire.stream.stanzas import iq, message, presence
from pyfire.stream.stanzas.errors import StanzaError, FeatureNotImplementedError

log = Logger(__name__)

class StanzaProcessor(object):
    """Holds a stanza handler for local domains"""

    def __init__(self, local_domains=("localhost")):
        self.local_domains = local_domains
        self.current_topic = None
        self.loop = ioloop.IOLoop.instance()
        self.ctx = zmq.Context()

        log.debug('Registering StanzaProcessor at forwarder..')
        router = self.ctx.socket(zmq.REQ)
        router.connect(config.get('ipc', 'forwarder_command_channel'))

        # TODO: add auth for authenticating us at the forwarder when it supports it
        router.send(' ')
        self.pub_url, self.sub_url = router.recv_json()
        router.close()

        self.pub_socket = self.ctx.socket(zmq.PUB)
        self.pub_socket.bind(self.pub_url)
        sub_socket = self.ctx.socket(zmq.SUB)
        sub_socket.connect(self.sub_url)
        # TODO: only subscribe to local domains
        sub_socket.setsockopt(zmq.SUBSCRIBE, '')
        substream = zmqstream.ZMQStream(sub_socket, self.loop)
        substream.on_recv(self.handle_stanza)

        # init the handlers
        self.stanza_handlers = {
                'iq': iq.Iq(),
                'message': message.Message(),
                'presence': presence.Presence()
            }

    def start(self):
        """Starts the handling of the bundles IOLoop"""
        self.loop.start()

    def handle_stanza(self, msg):
        """This actually handles the incomming stamzas"""
        if self.current_topic is None:
            self.current_topic = msg[0]
            return

        # TODO: check if we really want to handle the topis set..
        tree = cPickle.loads(msg[0])
        log.debug("Received stanza to handle: " + ET.tostring(tree))

        try:
            if tree.tag not in self.stanza_handlers:
                raise FeatureNotImplementedError(tree)

            response = self.stanza_handlers[tree.tag].handle(tree)
            if response is not None:
                self.pub_socket.send_multipart((tree.get("from"), cPickle.dumps(response)))
        except StanzaError, e:
            # send cought errors back to sender
            self.pub_socket.send_multipart((tree.get("from"), unicode(e)))
<commit_msg>Fix unpickle, zmq encodes newlines<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
    Stanza Processor

    This module starts a stanza processor that is responsible for
    processing all stanzas addressed to the local domain(s)

    :copyright: 2011 by the pyfire Team, see AUTHORS for more details.
    :license: BSD, see LICENSE for more details.
"""

import cPickle
import zmq
from zmq.eventloop import ioloop, zmqstream
import xml.etree.ElementTree as ET

from pyfire.logger import Logger
from pyfire import configuration as config
from pyfire.stream.stanzas import iq, message, presence
from pyfire.stream.stanzas.errors import StanzaError, FeatureNotImplementedError

log = Logger(__name__)

class StanzaProcessor(object):
    """Holds a stanza handler for local domains"""

    def __init__(self, local_domains=("localhost")):
        self.local_domains = local_domains
        self.current_topic = None
        self.loop = ioloop.IOLoop.instance()
        self.ctx = zmq.Context()

        log.debug('Registering StanzaProcessor at forwarder..')
        router = self.ctx.socket(zmq.REQ)
        router.connect(config.get('ipc', 'forwarder_command_channel'))

        # TODO: add auth for authenticating us at the forwarder when it supports it
        router.send(' ')
        self.pub_url, self.sub_url = router.recv_json()
        router.close()

        self.pub_socket = self.ctx.socket(zmq.PUB)
        self.pub_socket.bind(self.pub_url)
        sub_socket = self.ctx.socket(zmq.SUB)
        sub_socket.connect(self.sub_url)
        # TODO: only subscribe to local domains
        sub_socket.setsockopt(zmq.SUBSCRIBE, '')
        substream = zmqstream.ZMQStream(sub_socket, self.loop)
        substream.on_recv(self.handle_stanza)

        # init the handlers
        self.stanza_handlers = {
                'iq': iq.Iq(),
                'message': message.Message(),
                'presence': presence.Presence()
            }

    def start(self):
        """Starts the handling of the bundles IOLoop"""
        self.loop.start()

    def handle_stanza(self, msg):
        """This actually handles the incomming stamzas"""
        if self.current_topic is None:
            self.current_topic = msg[0]
            return

        # TODO: check if we really want to handle the topis set..
        tree = cPickle.loads(msg[0].split('\n'))
        log.debug("Received stanza to handle: " + ET.tostring(tree))

        try:
            if tree.tag not in self.stanza_handlers:
                raise FeatureNotImplementedError(tree)

            response = self.stanza_handlers[tree.tag].handle(tree)
            if response is not None:
                self.pub_socket.send_multipart((tree.get("from"), cPickle.dumps(response)))
        except StanzaError, e:
            # send cought errors back to sender
            self.pub_socket.send_multipart((tree.get("from"), unicode(e)))
<|endoftext|>