"""
:Authors: - Wilker Aziz
"""
import logging
import argparse
import sys

"""

:Authors: - Wilker Aziz
"""

from os.path import splitext
import subprocess as sp
import shlex
import argparse
import logging
import sys
import itertools
import os
import numpy as np
import traceback
from multiprocessing import Pool
from functools import partial
from collections import deque

from grasp.loss.fast_bleu import DecodingBLEU
from grasp.loss.fast_bleu import doc_bleu, stream_doc_bleu

import grasp.ptypes as ptypes

from grasp.recipes import smart_ropen, smart_wopen, make_unique_directory, pickle_it, unpickle_it, traceit

from grasp.scoring.scorer import TableLookupScorer, StatelessScorer, StatefulScorer
from grasp.scoring.util import make_models
from grasp.scoring.util import read_weights

from grasp.mt.cdec_format import load_grammar
from grasp.mt.util import GoalRuleMaker
from grasp.mt.util import save_forest, save_ffs, load_ffs, make_dead_srule, make_batches, number_of_batches
from grasp.mt.segment import SegmentMetaData
from grasp.mt.input import make_pass_grammar

import grasp.semiring as semiring
from grasp.semiring.operator import FixedLHS, FixedRHS

from grasp.formal.scfgop import output_projection
from grasp.formal.fsa import make_dfa, make_dfa_set, make_dfa_set2
from grasp.formal.scfgop import make_hypergraph_from_input_view, output_projection
from grasp.formal.scfgop import lookup_components, stateless_components
from grasp.formal.topsort import AcyclicTopSortTable
from grasp.formal.traversal import bracketed_string, yield_string
from grasp.formal.wfunc import TableLookupFunction, ConstantFunction, derivation_weight

from grasp.cfg.model import DummyConstant
from grasp.cfg.symbol import Nonterminal
from grasp.cfg.symbol import Terminal
from grasp.cfg.srule import OutputView

from grasp.alg.deduction import NederhofParser, EarleyParser, EarleyRescorer
from grasp.alg.inference import viterbi_derivation, AncestralSampler
from grasp.alg.value import acyclic_value_recursion, acyclic_reversed_value_recursion, compute_edge_expectation
from grasp.alg.rescoring import weight_edges
from grasp.alg.rescoring import SlicedRescoring
from grasp.alg.rescoring import stateless_rescoring
from grasp.alg.chain import apply_filters, group_by_identity, group_by_projection
from grasp.alg.expectation import expected_components

from grasp.scoring.frepr import FComponents

from grasp.io.results import save_mcmc_yields, save_mcmc_derivations, save_markov_chain

from random import shuffle
from numpy import linalg as LA
from scipy.optimize import minimize
from time import time, strftime
from types import SimpleNamespace

import grasp.mt.pipeline2 as pipeline

def npvec2str(nparray, fnames=None, separator=' '):
    """converts an array of feature values into a string (fnames can be provided)"""
    if fnames is None:
        return separator.join(repr(fvalue) for fvalue in nparray)
    else:
        return separator.join('{0}={1}'.format(fname, repr(fvalue)) for fname, fvalue in zip(fnames, nparray))

def cmd_optimisation(parser):
    # Optimisation
    parser.add_argument("--maxiter", '-M', type=int, default=10,
                        help="Maximum number of iterations")
    parser.add_argument('--mode', type=str, default='10',
                        help="use 'all' for all data, use 'online' for online updates, "
                             "use 0-100 to specify batch size in percentage")
    parser.add_argument('--shuffle',
                        action='store_true',
                        help='shuffle training instances')
    parser.add_argument('--temperature', type=float, default=1.0,
                        help='scales the initial model')
    parser.add_argument('--proxy-init', type=str, default='uniform',
                        help="use 'uniform' for uniform weights, 'random' for random weights, or choose a default weight")
    parser.add_argument('--target-init', type=str, default='uniform',
                        help="use 'uniform' for uniform weights, 'random' for random weights, or choose a default weight")
    parser.add_argument("--resume", type=int, default=0,
                        help="Resume from a certain iteration (requires the config file of the preceding run)")
    parser.add_argument('--merge', type=int, default=0,
                        help="how many iterations should we consider in estimating Z(x) (use 0 or less for all)")
    parser.add_argument("--sgd", type=int, nargs=2, default=[10, 10],
                        help="Number of iterations and function evaluations for target optimisation")
    parser.add_argument("--tol", type=float, nargs=2, default=[1e-9, 1e-9],
                        help="f-tol and g-tol in target optimisation")
    parser.add_argument("--L2", type=float, default=0.0,
                        help="Weight of L2 regulariser in target optimisation")

def cmd_logging(parser):
    parser.add_argument('--save-d',
                        action='store_true', default=0,
                        help='store sampled derivations (after MCMC filters apply)')
    parser.add_argument('--save-y',
                        action='store_true', default=0,
                        help='store sampled translations (after MCMC filters apply)')
    parser.add_argument('--verbose', '-v',
                        action='count', default=0,
                        help='increase the verbosity level')

def cmd_loss(group):
    group.add_argument('--bleu-order',
                       type=int, default=4,
                       metavar='N',
                       help="longest n-gram feature for sentence-level IBM-BLEU")
    group.add_argument('--bleu-smoothing',
                       type=float, default=1.0,
                       metavar='F',
                       help="add-p smoothing for sentence-level IBM-BLEU")

def cmd_parser(group):
    group.add_argument('--goal',
                       type=str, default='GOAL', metavar='LABEL',
                       help='default goal symbol (root after parsing/intersection)')
    group.add_argument('--framework',
                       type=str, default='exact', choices=['exact', 'slice'],
                       metavar='FRAMEWORK',
                       help="inference framework: 'exact', 'slice' sampling")

def cmd_grammar(group):
    group.add_argument('--start', '-S',
                       type=str, default='S',
                       metavar='LABEL',
                       help='default start symbol')
    group.add_argument("--dev-grammars", type=str,
                       help="grammars for the dev set")
    group.add_argument("--devtest-grammars", type=str,
                       help="grammars for the devtest set")
    group.add_argument('--extra-grammar',
                       action='append', default=[], metavar='PATH',
                       help="path to an additional grammar (multiple allowed)")
    group.add_argument('--glue-grammar',
                       action='append', default=[], metavar='PATH',
                       help="glue rules are only applied to initial states (multiple allowed)")
    group.add_argument('--pass-through',
                       action='store_true',
                       help="add pass-through rules for every input word (and an indicator feature for unknown words)")
    group.add_argument('--default-symbol', '-X',
                       type=str, default='X', metavar='LABEL',
                       help='default nonterminal (used for pass-through rules and automatic glue rules)')

def cmd_sampler(group):
    group.add_argument('--samples',
                       type=int, default=100,
                       metavar='N',
                       help="number of samples from proxy")

def get_argparser():
    parser = argparse.ArgumentParser(description='Training by MLE',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    #parser.add_argument('config', type=str, help="configuration file")
    parser.add_argument("workspace",
                        type=str, default=None,
                        help="where samples can be found and where decisions are placed")
    parser.add_argument("proxy", type=str,
                        help="proxy model description")
    parser.add_argument("target", type=str,
                        help="target model description")
    parser.add_argument("dev", type=str,
                        help="development set")
    parser.add_argument('--experiment',
                        type=str,
                        help='folder within the workspace where results are stored'
                             'by default we use a timestamp and a random suffix')
    parser.add_argument("--proxy-weights", '-Q', type=str,
                        help="proxy weights")
    parser.add_argument("--target-weights", '-P', type=str,
                        help="target weights")
    parser.add_argument("--proxy-temperature", '-Tq', type=float, default=1.0,
                        help="scales the model (the bigger the more uniform)")
    parser.add_argument("--target-temperature", '-Tp', type=float, default=1.0,
                        help="scales the model (the bigger the more uniform)")
    parser.add_argument("--jobs", type=int, default=2, help="number of processes")
    parser.add_argument('--dev-alias', type=str, default='dev',
            help='Change the alias of the dev set')
    parser.add_argument("--devtest", type=str,
                        help="devtest set")
    parser.add_argument('--devtest-alias', type=str, default='devtest',
            help='Change the alias of the devtest set')
    parser.add_argument('--redo', action='store_true',
                        help='overwrite already computed files (by default we do not repeat computation)')
    cmd_parser(parser.add_argument_group('Parser'))
    cmd_grammar(parser.add_argument_group('Grammar'))
    cmd_optimisation(parser.add_argument_group('Parameter optimisation by SGD'))
    cmd_loss(parser.add_argument_group('Loss'))
    cmd_sampler(parser.add_argument_group('Importance sampler'))
    cmd_logging(parser.add_argument_group('Logging'))
    # General

    return parser

def make_dirs(args, exist_ok=True):
    """
    Make output directories and saves the command line arguments for documentation purpose.

    :param args: command line arguments
    :return: main output directory within workspace (prefix is a timestamp and suffix is a unique random string)
    """

    # create the workspace if missing
    logging.info('Workspace: %s', args.workspace)
    if not os.path.exists(args.workspace):
        os.makedirs(args.workspace, exist_ok=exist_ok)

    # create a unique experiment area or reuse a given one
    if not args.experiment:
        outdir = make_unique_directory(args.workspace)
    else:
        outdir = '{0}/{1}'.format(args.workspace, args.experiment)
        os.makedirs(outdir, exist_ok=exist_ok)
    logging.info('Writing files to: %s', outdir)

    devdir = '{0}/{1}'.format(outdir, args.dev_alias)
    os.makedirs(devdir, exist_ok=exist_ok)

    if args.devtest:
        devtestdir = '{0}/{1}'.format(outdir, args.devtest_alias)
        os.makedirs(devtestdir, exist_ok=exist_ok)

    dynamicdir = '{0}/iterations'.format(outdir)
    os.makedirs(dynamicdir, exist_ok=exist_ok)

    return outdir, devdir

@traceit
def pass0_to_pass2(seg, options, workingdir, model, redo, log):
    saving = {'pass2.forest': '{0}/{1}.q-forest'.format(workingdir, seg.id),
              'pass2.components': '{0}/{1}.q-components'.format(workingdir, seg.id)}

    if pipeline.all_steps_complete(saving, redo):
        return True

    forest, components = pipeline.pass0_to_pass2(seg,
                                                 options,
                                                 model.lookup,
                                                 model.stateless,
                                                 model.stateful,
                                                 saving=saving, redo=redo, log=log)

    return forest.n_nodes() > 0

def make_pass0_to_pass2_options(args):
    options = SimpleNamespace()
    options.extra_grammars = args.extra_grammar
    options.glue_grammars = args.glue_grammar
    options.pass_through = args.pass_through
    options.default_symbol = args.default_symbol
    options.goal = args.goal
    options.start = args.start
    return options

def parse_training(args, staticdir, model, segments):
    logging.info('Parsing %d training instances using %d workers', len(segments), args.jobs)

    with Pool(args.jobs) as workers:
        feedback = workers.map(partial(pass0_to_pass2,
                                       options=make_pass0_to_pass2_options(args),
                                       workingdir=staticdir,
                                       model=model,
                                       redo=args.redo,
                                       log=logging.info),
                               segments)
    return tuple([seg for seg, status in zip(segments, feedback) if status])

def make_impsamp_options(args):
    options = make_pass0_to_pass2_options(args)
    options.samples = args.samples
    options.bleu_order = args.bleu_order
    options.bleu_smoothing = args.bleu_smoothing
    return options

@traceit
def importance_sample(seg, options, staticdir, workingdir, proxy, target, redo, log):

    saving = {'is.samples': '{0}/samples/{1}.is'.format(workingdir, seg.id),
              'pass2.forest': '{0}/{1}.q-forest'.format(staticdir, seg.id),
              'pass2.components': '{0}/{1}.q-components'.format(staticdir, seg.id)}

    # TODO:
    #   1. normalise q(d) \propto g(d) exactly?
    #   2. use sample frequency for q(d)?
    #   3. use unnormalised g(d)

    samples = pipeline.importance_sample(seg,
                                         options,
                                         proxy,
                                         target,
                                         saving=saving, redo=redo, log=log)

    # support
    Y = [None] * len(samples)
    # posterior
    Q = np.zeros(len(samples), dtype=ptypes.weight)
    P = np.zeros(len(samples), dtype=ptypes.weight)
    # compute posterior
    for i, sample in enumerate(samples):
        Y[i] = sample.y.split()
        D = sample.D
        qy = 0.0
        py = 0.0
        #py = semiring.inside.zero
        for d in D:
            f = target.score(d.p_comps)
            g = proxy.score(d.q_comps)  # TODO: consider normalising g exactly
            w = semiring.inside.divide(f, g)
            qy += float(d.count) / len(samples)
            py += d.count * semiring.inside.as_real(w)
            #py = semiring.inside.plus(semiring.inside.times(semiring.inside.from_real(d.count), w), py)
        #P[i] = semiring.inside.as_real(py)
        Q[i] = qy
        P[i] = py
    P /= P.sum()
    # compute consensus loss
    bleu = DecodingBLEU(Y, P, max_order=options.bleu_order, smoothing=options.bleu_smoothing)
    L = [bleu.loss(y) for y in Y]
    ranking = sorted(range(len(Y)), key=lambda i: (L[i], -P[i]))

    with smart_wopen('{0}/samples/{1}.ranking.gz'.format(workingdir, seg.id)) as fo:
        print('# L ||| p(y) ||| q(y) ||| y', file=fo)
        for i in ranking:
            print('{0} ||| {1} ||| {2} ||| {3}'.format(L[i], P[i], Q[i], samples[i].y), file=fo)

    return samples[i].y, P[i], L[i]

def sample_and_decode(args, staticdir, workingdir, proxy, target, segments):
    logging.info('Decoding %d segments using %d workers', len(segments), args.jobs)
    os.makedirs('{0}/samples'.format(workingdir), exist_ok=True)
    with Pool(args.jobs) as workers:
        decisions = workers.map(partial(importance_sample,
                                        options=make_impsamp_options(args),
                                        staticdir=staticdir,
                                        workingdir=workingdir,
                                        proxy=proxy,
                                        target=target,
                                        redo=args.redo,
                                        log=logging.info),
                                segments)
    return decisions

def mteval(args, workspace, iteration, proxy, target, segments, alias):
    decisions = sample_and_decode(args,
                                  '{0}/{1}'.format(workspace, alias),
                                  '{0}/iterations/{1}/{2}'.format(workspace, iteration, alias),
                                  proxy, target, segments)
    evaldir = '{0}/iterations/{1}/{2}'.format(workspace, iteration, alias)
    os.makedirs(evaldir, exist_ok=True)
    with smart_wopen('{0}/hyps'.format(evaldir)) as fo:
        for y, p, l in decisions:
            print(y, file=fo)
    bleu, pn, bp = stream_doc_bleu(smart_ropen('{0}/hyps'.format(evaldir)),
                                   smart_ropen('{0}/{1}/refs'.format(workspace, alias)),
                                   max_order=args.bleu_order,
                                   smoothing=args.bleu_smoothing)
    logging.info('BLEU %s: %.4f', alias, bleu)
    return bleu

def sanity_checks(args):
    failed = False
    if not os.path.exists(args.dev):
        logging.error('Training set not found: %s', args.dev)
        failed = True
    if args.devtest and not os.path.exists(args.devtest):
        logging.error('Validation set not found: %s', args.devtest)
        failed = True
    if not os.path.exists(args.proxy):
        logging.error('Proxy model description not found: %s', args.proxy)
        failed = True
    if not os.path.exists(args.target):
        logging.error('Target model description not found: %s', args.target)
        failed = True
    if args.proxy_weights and not os.path.exists(args.proxy_weights):
        logging.error('Proxy model weights not found: %s', args.proxy_weights)
        failed = True
    if args.target_weights and not os.path.exists(args.target_weights):
        logging.error('Target model weights not found: %s', args.target_weights)
        failed = True
    return not failed

def core(args):

    workspace, devdir = make_dirs(args)

    if not sanity_checks(args):
        raise FileNotFoundError('One or more files could not be found')

    proxy = pipeline.load_model(args.proxy, args.proxy_weights, args.proxy_init, args.proxy_temperature)
    logging.info('Proxy:\n%s', proxy)
    target = pipeline.load_model(args.target, args.target_weights, args.target_init, args.target_temperature)
    logging.info('Target:\n%s', target)

    # 2. Parse data
    dev = pipeline.read_segments_from_file(args.dev, args.dev_grammars)
    dev = parse_training(args, devdir, proxy, dev)
    logging.info(' %d training instances', len(dev))

    # store references for evaluation purposes
    pipeline.save_references('{0}/{1}/refs'.format(workspace, args.dev_alias), dev)

    # Validation set
    if args.devtest is None:
        args.devtest = args.dev
        args.devtest_alias = args.dev_alias
        args.devtest_grammars = args.dev_grammars
        devtest = dev
    else:
        devtest = pipeline.read_segments_from_file(args.devtest, args.devtest_grammars)
        devtest = parse_training(args, '{0}/{1}'.format(workspace, args.devtest_alias), proxy, devtest)
        logging.info(' %d validation instances', len(devtest))
        pipeline.save_references('{0}/{1}/refs'.format(workspace, args.devtest_alias), devtest)

    # evaluate the initial model
    mteval(args, workspace, 0, proxy, target, devtest, args.devtest_alias)
    ##print('{0} ||| init ||| {1}={2} ||| {3}'.format(0, args.devtest_alias, bleu, npvec2str(model.weights().densify(), fnames)))

    # 3. Optimise
    #dimensionality = len(fnames)

def main():
    args = get_argparser().parse_args()

    if args.verbose == 1:
        logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s', datefmt='%H:%M:%S')
    elif args.verbose > 1:
        logging.basicConfig(level=logging.DEBUG, format='%(asctime)s %(levelname)s %(message)s', datefmt='%H:%M:%S')

    core(args)

if __name__ == '__main__':
    main()


#!/usr/bin/env python

import math
from subprocess import Popen, PIPE
import re
import traceback
import numpy
import rospy
from std_srvs.srv import Trigger
from sensor_msgs.msg import Image, CameraInfo, NavSatFix, Imu, Range
from mavros_msgs.msg import State, OpticalFlowRad
from mavros_msgs.srv import ParamGet
from geometry_msgs.msg import PoseStamped, TwistStamped, PoseWithCovarianceStamped
import tf.transformations as t
from aruco_pose.msg import MarkerArray

# TODO: roscore is running
# TODO: clever.service is running
# TODO: check attitude is present
# TODO: disk free space
# TODO: map, base_link, body
# TODO: rc service
# TODO: perform commander check, ekf2 status on PX4
# TODO: check if FCU params setter succeed
# TODO: selfcheck ROS service (with blacklists for checks)

rospy.init_node('selfcheck')

failures = []

def failure(text, *args):
    failures.append(text % args)

def check(name):
    def inner(fn):
        def wrapper(*args, **kwargs):
            failures[:] = []
            try:
                fn(*args, **kwargs)
                for f in failures:
                    rospy.logwarn('%s: %s', name, f)
            except Exception as e:
                for f in failures:
                    rospy.logwarn('%s: %s', name, f)
                traceback.print_exc()
                rospy.logwarn('%s: exception occurred', name)
                return
            if not failures:
                rospy.loginfo('%s: OK', name)
        return wrapper
    return inner

param_get = rospy.ServiceProxy('mavros/param/get', ParamGet)

def get_param(name):
    res = param_get(param_id=name)
    if not res.success:
        failure('Unable to retrieve PX4 parameter %s', name)
    else:
        if res.value.integer != 0:
            return res.value.integer
        return res.value.real

@check('FCU')
def check_fcu():
    try:
        state = rospy.wait_for_message('mavros/state', State, timeout=3)
        if not state.connected:
            failure('no connection to the FCU (check wiring)')

        est = get_param('SYS_MC_EST_GROUP')
        if est == 1:
            rospy.loginfo('Selected estimator: LPE')
            fuse = get_param('LPE_FUSION')
            if fuse & (1 << 4):
                rospy.loginfo('LPE_FUSION: land detector fusion is enabled')
            else:
                rospy.loginfo('LPE_FUSION: land detector fusion is disabled')
            if fuse & (1 << 7):
                rospy.loginfo('LPE_FUSION: barometer fusion is enabled')
            else:
                rospy.loginfo('LPE_FUSION: barometer fusion is disabled')

        elif est == 2:
            rospy.loginfo('Selected estimator: EKF2')
        else:
            failure('Unknown selected estimator: %s', est)

    except rospy.ROSException:
        failure('no MAVROS state (check wiring)')

@check('Camera')
def check_camera(name):
    try:
        img = rospy.wait_for_message(name + '/image_raw', Image, timeout=1)
    except rospy.ROSException:
        failure('%s: no images (is the camera connected properly?)', name)
        return
    try:
        info = rospy.wait_for_message(name + '/camera_info', CameraInfo, timeout=1)
    except rospy.ROSException:
        failure('%s: no calibration info', name)
        return

    if img.width != info.width:
        failure('%s: calibration width doesn\'t match image width (%d != %d)', name, info.width, img.width)
    if img.height != info.height:
        failure('%s: calibration height doesn\'t match image height (%d != %d))', name, info.height, img.height)

@check('ArUco detector')
def check_aruco():
    try:
        rospy.wait_for_message('aruco_detect/markers', MarkerArray, timeout=1)
    except rospy.ROSException:
        failure('no markers detection')
        return
    try:
        rospy.wait_for_message('aruco_map/pose', PoseWithCovarianceStamped, timeout=1)
    except rospy.ROSException:
        failure('no map detection')

@check('Vision position estimate')
def check_vpe():
    try:
        vis = rospy.wait_for_message('mavros/vision_pose/pose', PoseStamped, timeout=1)
    except rospy.ROSException:
        try:
            vis = rospy.wait_for_message('mavros/mocap/pose', PoseStamped, timeout=1)
        except rospy.ROSException:
            failure('no VPE or MoCap messages')
            return

    # check PX4 settings
    est = get_param('SYS_MC_EST_GROUP')
    if est == 1:
        ext_yaw = get_param('ATT_EXT_HDG_M')
        if ext_yaw != 1:
            failure('vision yaw is disabled, change ATT_EXT_HDG_M parameter')
        vision_yaw_w = get_param('ATT_W_EXT_HDG')
        if vision_yaw_w == 0:
            failure('vision yaw weight is zero, change ATT_W_EXT_HDG parameter')
        else:
            rospy.loginfo('Vision yaw weight: %.2f', vision_yaw_w)
        fuse = get_param('LPE_FUSION')
        if not fuse & (1 << 2):
            failure('vision position fusing is disabled, change LPE_FUSION parameter')
        delay = get_param('LPE_VIS_DELAY')
        if delay != 0:
            failure('LPE_VIS_DELAY parameter is %s, but it should be zero', delay)
        rospy.loginfo('LPE_VIS_XY is %.2f m, LPE_VIS_Z is %.2f m', get_param('LPE_VIS_XY'), get_param('LPE_VIS_Z'))
    elif est == 2:
        fuse = get_param('EKF2_AID_MASK')
        if not fuse & (1 << 3):
            failure('vision position fusing is disabled, change EKF2_AID_MASK parameter')
        if not fuse & (1 << 4):
            failure('vision yaw fusing is disabled, change EKF2_AID_MASK parameter')
        delay = get_param('EKF2_EV_DELAY')
        if delay != 0:
            failure('EKF2_EV_DELAY is %.2f, but it should be zero', delay)
        rospy.loginfo('EKF2_EVA_NOISE is %.3f, EKF2_EVP_NOISE is %.3f',
            get_param('EKF2_EVA_NOISE'),
            get_param('EKF2_EVP_NOISE'))

    # check vision pose and estimated pose inconsistency
    try:
        pose = rospy.wait_for_message('mavros/local_position/pose', PoseStamped, timeout=1)
    except:
        return
    horiz = math.hypot(vis.pose.position.x - pose.pose.position.x, vis.pose.position.y - pose.pose.position.y)
    if horiz > 0.5:
        failure('horizontal position inconsistency: %.2f m', horiz)
    vert = vis.pose.position.z - pose.pose.position.z
    if abs(vert) > 0.5:
        failure('vertical position inconsistency: %.2f m', vert)
    op = pose.pose.orientation
    ov = vis.pose.orientation
    yawp, _, _ = t.euler_from_quaternion((op.x, op.y, op.z, op.w), axes='rzyx')
    yawv, _, _ = t.euler_from_quaternion((ov.x, ov.y, ov.z, ov.w), axes='rzyx')
    yawdiff = yawp - yawv
    yawdiff = math.degrees((yawdiff + 180) % 360 - 180)
    if abs(yawdiff) > 8:
        failure('yaw inconsistency: %.2f deg', yawdiff)

@check('Simple offboard node')
def check_simpleoffboard():
    try:
        rospy.wait_for_service('navigate', timeout=3)
        rospy.wait_for_service('get_telemetry', timeout=3)
        rospy.wait_for_service('land', timeout=3)
    except rospy.ROSException:
        failure('no simple_offboard services')

@check('IMU')
def check_imu():
    try:
        rospy.wait_for_message('mavros/imu/data', Imu, timeout=1)
    except rospy.ROSException:
        failure('no IMU data (check flight controller calibration)')

@check('Local position')
def check_local_position():
    try:
        pose = rospy.wait_for_message('mavros/local_position/pose', PoseStamped, timeout=1)
        o = pose.pose.orientation
        _, pitch, roll = t.euler_from_quaternion((o.x, o.y, o.z, o.w), axes='rzyx')
        MAX_ANGLE = math.radians(2)
        if abs(pitch) > MAX_ANGLE:
            failure('pitch is %.2f deg; place copter horizontally or redo level horizon calib',
                    math.degrees(pitch))
        if abs(roll) > MAX_ANGLE:
            failure('roll is %.2f deg; place copter horizontally or redo level horizon calib',
                    math.degrees(roll))

    except rospy.ROSException:
        failure('no local position')

@check('Velocity estimation')
def check_velocity():
    try:
        velocity = rospy.wait_for_message('mavros/local_position/velocity', TwistStamped, timeout=1)
        horiz = math.hypot(velocity.twist.linear.x, velocity.twist.linear.y)
        vert = velocity.twist.linear.z
        if abs(horiz) > 0.1:
            failure('horizontal velocity estimation is %.2f m/s; is copter staying still?' % horiz)
        if abs(vert) > 0.1:
            failure('vertical velocity estimation is %.2f m/s; is copter staying still?' % vert)

        angular = velocity.twist.angular
        ANGULAR_VELOCITY_LIMIT = 0.01
        if abs(angular.x) > ANGULAR_VELOCITY_LIMIT:
            failure('pitch rate estimation is %.2f rad/s (%.2f deg/s); is copter staying still?',
                    angular.x, math.degrees(angular.x))
        if abs(angular.y) > ANGULAR_VELOCITY_LIMIT:
            failure('pitch rate estimation is %.2f rad/s (%.2f deg/s); is copter staying still?',
                    angular.y, math.degrees(angular.y))
        if abs(angular.z) > ANGULAR_VELOCITY_LIMIT:
            failure('pitch rate estimation is %.2f rad/s (%.2f deg/s); is copter staying still?',
                    angular.z, math.degrees(angular.z))
    except rospy.ROSException:
        failure('no velocity estimation')

@check('Global position (GPS)')
def check_global_position():
    try:
        rospy.wait_for_message('mavros/global_position/global', NavSatFix, timeout=1)
    except rospy.ROSException:
        failure('no global position')

@check('Optical flow')
def check_optical_flow():
    # TODO:check FPS!
    try:
        rospy.wait_for_message('mavros/px4flow/raw/send', OpticalFlowRad, timeout=0.5)

        # check PX4 settings
        rot = get_param('SENS_FLOW_ROT')
        if rot != 0:
            failure('SENS_FLOW_ROT parameter is %s, but it should be zero', rot)
        est = get_param('SYS_MC_EST_GROUP')
        if est == 1:
            fuse = get_param('LPE_FUSION')
            if not fuse & (1 << 1):
                failure('optical flow fusing is disabled, change LPE_FUSION parameter')
            if not fuse & (1 << 1):
                failure('flow gyro compensation is disabled, change LPE_FUSION parameter')
            scale = get_param('LPE_FLW_SCALE')
            if not numpy.isclose(scale, 1.0):
                failure('LPE_FLW_SCALE parameter is %.2f, but it should be 1.0', scale)

            rospy.loginfo('LPE_FLW_QMIN is %s, LPE_FLW_R is %.4f, LPE_FLW_RR is %.4f, SENS_FLOW_MINHGT is %.3f, SENS_FLOW_MAXHGT is %.3f',
                get_param('LPE_FLW_QMIN'),
                get_param('LPE_FLW_R'),
                get_param('LPE_FLW_RR'),
                get_param('SENS_FLOW_MINHGT'),
                get_param('SENS_FLOW_MAXHGT'))
        elif est == 2:
            fuse = get_param('EKF2_AID_MASK')
            if not fuse & (1 << 1):
                failure('optical flow fusing is disabled, change EKF2_AID_MASK parameter')
            delay = get_param('EKF2_OF_DELAY')
            if delay != 0:
                failure('EKF2_OF_DELAY is %.2f, but it should be zero', delay)
            rospy.loginfo('EKF2_OF_QMIN is %s, EKF2_OF_N_MIN is %.4f, EKF2_OF_N_MAX is %.4f, SENS_FLOW_MINHGT is %.3f, SENS_FLOW_MAXHGT is %.3f',
                get_param('EKF2_OF_QMIN'),
                get_param('EKF2_OF_N_MIN'),
                get_param('EKF2_OF_N_MAX'),
                get_param('SENS_FLOW_MINHGT'),
                get_param('SENS_FLOW_MAXHGT'))

    except rospy.ROSException:
        failure('no optical flow data (from Raspberry)')

@check('Rangefinder')
def check_rangefinder():
    # TODO: check FPS!
    rng = False
    try:
        rospy.wait_for_message('mavros/distance_sensor/rangefinder_sub', Range, timeout=4)
        rng = True
    except rospy.ROSException:
        failure('no rangefinder data from Raspberry')

    try:
        rospy.wait_for_message('mavros/distance_sensor/rangefinder', Range, timeout=4)
        rng = True
    except rospy.ROSException:
        failure('no rangefinder data from PX4')

    if not rng:
        return

    est = get_param('SYS_MC_EST_GROUP')
    if est == 1:
        fuse = get_param('LPE_FUSION')
        if not fuse & (1 << 5):
            rospy.loginfo('"pub agl as lpos down" in LPE_FUSION is disabled, NOT operating over flat surface')
        else:
            rospy.loginfo('"pub agl as lpos down" in LPE_FUSION is enabled, operating over flat surface')

    elif est == 2:
        hgt = get_param('EKF2_HGT_MODE')
        if hgt != 2:
            rospy.loginfo('EKF2_HGT_MODE != Range sensor, NOT operating over flat surface')
        else:
            rospy.loginfo('EKF2_HGT_MODE = Range sensor, operating over flat surface')
        aid = get_param('EKF2_RNG_AID')
        if aid != 1:
            rospy.loginfo('EKF2_RNG_AID != 1, range sensor aiding disabled')
        else:
            rospy.loginfo('EKF2_RNG_AID = 1, range sensor aiding enabled')

@check('Boot duration')
def check_boot_duration():
    proc = Popen('systemd-analyze', stdout=PIPE)
    proc.wait()
    output = proc.communicate()[0]
    r = re.compile(r'([\d\.]+)s$')
    duration = float(r.search(output).groups()[0])
    if duration > 15:
        failure('long Raspbian boot duration: %ss (systemd-analyze for analyzing)', duration)

@check('CPU usage')
def check_cpu_usage():
    WHITELIST = 'nodelet',
    CMD = "top -n 1 -b -i | tail -n +8 | awk '{ printf(\"%-8s\\t%-8s\\t%-8s\\n\", $1, $9, $12); }'"
    proc = Popen(CMD, stdout=PIPE, shell=True)
    proc.wait()
    output = proc.communicate()[0]
    processes = output.split('\n')
    for process in processes:
        if not process:
            continue
        pid, cpu, cmd = process.split('\t')

        if cmd.strip() not in WHITELIST and float(cpu) > 30:
            failure('high CPU usage (%s%%) detected: %s (PID %s)',
                    cpu.strip(), cmd.strip(), pid.strip())

def selfcheck():
    check_fcu()
    check_imu()
    check_local_position()
    check_velocity()
    check_global_position()
    check_camera('main_camera')
    check_aruco()
    check_simpleoffboard()
    check_optical_flow()
    check_vpe()
    check_rangefinder()
    check_cpu_usage()
    check_boot_duration()

if __name__ == '__main__':
    rospy.loginfo('Performing selfcheck...')
    selfcheck()


from __future__ import print_function
import argparse
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import time
import math
from sintel_io import depth_read
from dataloader import sintel_listflowfile as lt
# from dataloader import sintel_listflowfile_with_filter as lt
# from dataloader import sintel_listflowfile_with_filter_with_depth as lt
# from dataloader import sintel_listflowfile_without_filter_with_depth as lt
from dataloader import SintelFlowLoader as DA
from models import *

parser = argparse.ArgumentParser(description='PSMNet')
parser.add_argument('--maxdisp', type=int, default=192,
                    help='maxium disparity')
parser.add_argument('--model', default='stackhourglass',
                    help='select model')
parser.add_argument('--datapath', default='/media/yotamg/bd0eccc9-4cd5-414c-b764-c5a7890f9785/Yotam/Stereo',
                    help='datapath')
parser.add_argument('--left_imgs', default=None,
                    help='left img train dir name')
parser.add_argument('--right_imgs', default=None,
                    help='left img train dir name')
parser.add_argument('--disp_imgs', default=None,
                    help='left img train dir name')
parser.add_argument('--disp_R_imgs', default=None,
                    help='right img train dir name')
parser.add_argument('--occ_L_dir', default='occ_flatten',
                    help='occlusions left dir')
parser.add_argument('--oof_L_dir', default='oof_flatten',
                    help='out-of-frame left dir')
parser.add_argument('--epochs', type=int, default=200,
                    help='number of epochs to train')
parser.add_argument('--loadmodel', default=None,
# parser.add_argument('--loadmodel', default='./checkpoints/checkpoint_9.tar',
                    help='load model')
# parser.add_argument('--loadmodel', default='./checkpoints/checkpoint_filter_loss_2.6.tar',
# parser.add_argument('--loadmodel', default='./checkpoints/checkpoint_clean_from_scratch_loss_2.1.tar',
# parser.add_argument('--loadmodel', default='./pretrained_model_KITTI2015.tar',
parser.add_argument('--savemodel', default='./checkpoints/',
                    help='save model')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--clean', action='store_true', default=False,
                    help='random seed (default: 1)')
parser.add_argument('--cont', action='store_true', default=False,
                    help='random seed (default: 1)')
parser.add_argument('--dfd', action='store_true', default=False,
                    help='include dfd net')
parser.add_argument('--dfd_at_end', action='store_true', default=False,
                    help='include dfd net')
parser.add_argument('--right_head', action='store_true', default=False,
                    help='add right disparity head')
parser.add_argument('--pred_occlusion', action='store_true', default=False,
                    help='pred occlusion or right depth')
parser.add_argument('--pred_oof', action='store_true', default=False,
                    help='pred out of frame head')

# parser.add_argument('--gpu_id', type=int, default=0, metavar='S',
#                     help='gpu_id')
args = parser.parse_args()

args.cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)

all_left_img, all_right_img, all_left_disp, all_right_disp, all_left_occ, all_left_oof, test_left_img, test_right_img, test_left_disp, test_right_disp, test_left_occ, test_left_oof= lt.dataloader(args.datapath, args.left_imgs, args.right_imgs, args.disp_imgs,args.disp_R_imgs, occ_L_dir=args.occ_L_dir, oof_L_dir=args.oof_L_dir,clean=args.clean)

TrainImgLoader = torch.utils.data.DataLoader(
    # DA.myImageFloder(all_left_img, all_right_img, all_left_disp, True, dploader=DA.depth_loader),
    DA.myImageFloder(all_left_img, all_right_img, all_left_disp, all_right_disp, left_occ=all_left_occ, left_oof=all_left_oof, training=True, dploader=depth_read, cont=args.cont),
    batch_size=1, shuffle=True, num_workers=8, drop_last=False)

TestImgLoader = torch.utils.data.DataLoader(
    DA.myImageFloder(test_left_img, test_right_img, test_left_disp, test_right_disp, left_occ=test_left_occ, left_oof=test_left_oof, training=False, dploader=depth_read, cont=args.cont),
    batch_size=1, shuffle=False, num_workers=4, drop_last=False)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if args.model == 'stackhourglass':
    model = stackhourglass(args.maxdisp, device=device, dfd_net=args.dfd, dfd_at_end=args.dfd_at_end, right_head=args.right_head, pred_oof=args.pred_oof)
elif args.model == 'basic':
    model = basic(args.maxdisp)
else:
    print('no model')

if args.cuda:
    model = nn.DataParallel(model)
    model.to(device)

if args.loadmodel is not None:
    state_dict = torch.load(args.loadmodel)
    model.load_state_dict(state_dict['state_dict'], strict=False)

print('Number of model parameters: {}'.format(sum([p.data.nelement() for p in model.parameters()])))

optimizer = optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.999))
# optimizer = optim.Adam(model.parameters(), lr=0.1, betas=(0.9, 0.999))

def train(imgL, imgR, disp_L, disp_R,left_occ, left_oof):
    model.train()
    imgL = Variable(torch.FloatTensor(imgL))
    imgR = Variable(torch.FloatTensor(imgR))
    disp_L = Variable(torch.FloatTensor(disp_L))

    if args.cuda:
        imgL, imgR, disp_true, disp_R, left_occ, left_oof = imgL.to(device), imgR.to(device), disp_L.to(device), disp_R.to(device), left_occ.to(device), left_oof.to(device)

    # ---------
    mask = disp_true < args.maxdisp
    mask.detach_()
    mask_R = disp_R < args.maxdisp
    mask_R.detach_()
    # ----
    optimizer.zero_grad()

    if args.model == 'stackhourglass':
        if args.pred_oof:
            output1, output2, output3, output4, output5, output6,output7, output8, output9 = model(imgL, imgR)
        if args.right_head:
            output1, output2, output3, output4, output5, output6= model(imgL, imgR)
        else:
            output1, output2, output3 = model(imgL, imgR)
        output1 = torch.squeeze(output1, 1)
        output2 = torch.squeeze(output2, 1)
        output3 = torch.squeeze(output3, 1)
        if args.right_head:
            output4 = torch.squeeze(output4, 1)
            output5 = torch.squeeze(output5, 1)
            output6 = torch.squeeze(output6, 1)
        if args.pred_oof:
            output7 = torch.squeeze(output7, 1)
            output8 = torch.squeeze(output8, 1)
            output9 = torch.squeeze(output9, 1)
        # loss = F.smooth_l1_loss(output3[mask], disp_true[mask],size_average=True)
        loss = 0.5 * F.smooth_l1_loss(output1[mask], disp_true[mask], size_average=True) + 0.7 * F.smooth_l1_loss(
            output2[mask], disp_true[mask], size_average=True) + F.smooth_l1_loss(output3[mask], disp_true[mask],size_average=True)
        if args.right_head:
            if args.pred_occlusion:
                loss += 0.5 * F.smooth_l1_loss(output4[mask], left_occ[mask], size_average=True) + 0.7 * F.smooth_l1_loss(
                    output5[mask], left_occ[mask], size_average=True) + F.smooth_l1_loss(output6[mask], left_occ[mask],  size_average=True)
                if args.pred_oof:
                    loss += 0.5 * F.smooth_l1_loss(output7[mask], left_oof[mask], size_average=True) + 0.7 * F.smooth_l1_loss(
                        output8[mask], left_oof[mask], size_average=True) + F.smooth_l1_loss(output9[mask], left_oof[mask], size_average=True)
            else:
                loss += 0.5 * F.smooth_l1_loss(output4[mask_R], disp_R[mask_R], size_average=True) + 0.7 * F.smooth_l1_loss(
                    output5[mask_R], disp_R[mask_R], size_average=True) + F.smooth_l1_loss(output6[mask_R], disp_R[mask_R],size_average=True)

    elif args.model == 'basic':
        output = model(imgL, imgR)
        output = torch.squeeze(output, 1)
        loss = F.smooth_l1_loss(output[mask], disp_true[mask], size_average=True)

    loss.backward()
    optimizer.step()

    return loss.item()

def test(imgL, imgR, disp_true):
    model.eval()
    imgL = Variable(torch.FloatTensor(imgL))
    imgR = Variable(torch.FloatTensor(imgR))
    if args.cuda:
        imgL, imgR = imgL.to(device), imgR.to(device)

    # ---------
    mask = disp_true < 192
    # ----

    with torch.no_grad():
        output3 = model(imgL, imgR)

    output = torch.squeeze(output3.data.cpu(), 1)[:, 4:, :]

    if len(disp_true[mask]) == 0:
        loss = 0
    else:
        loss = torch.mean(torch.abs(output[mask] - disp_true[mask]))  # end-point-error

    return loss

def adjust_learning_rate(optimizer, epoch):
    lr = 0.001*(1/epoch)
    print(lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def main():
    start_full_time = time.time()
    for epoch in range(1, args.epochs + 1):
        print('This is %d-th epoch' % (epoch))
        total_train_loss = 0
        # adjust_learning_rate(optimizer, epoch)

        ## training ##
        for batch_idx, (imgL_crop, imgR_crop, disp_crop_L, disp_crop_R,left_occ, left_oof) in enumerate(TrainImgLoader):
            start_time = time.time()

            loss = train(imgL_crop, imgR_crop, disp_crop_L, disp_crop_R, left_occ, left_oof)
            print('Iter %d training loss = %.3f , time = %.2f' % (batch_idx, loss, time.time() - start_time))
            total_train_loss += loss
        print('epoch %d total training loss = %.3f' % (epoch, total_train_loss / len(TrainImgLoader)))

        # SAVE
        if not os.path.isdir(args.savemodel):
            os.makedirs(args.savemodel)
        savefilename = args.savemodel + '/checkpoint_' + str(epoch) + '.tar'
        torch.save({
            'epoch': epoch,
            'state_dict': model.state_dict(),
            'train_loss': total_train_loss / len(TrainImgLoader),
        }, savefilename)

    print('full training time = %.2f HR' % ((time.time() - start_full_time) / 3600))

    # ------------- TEST ------------------------------------------------------------
    # total_test_loss = 0
    # for batch_idx, (imgL, imgR, disp_L) in enumerate(TestImgLoader):
    #     test_loss = test(imgL, imgR, disp_L)
    #     print('Iter %d test loss = %.3f' % (batch_idx, test_loss))
    #     total_test_loss += test_loss
    #
    # print('total test loss = %.3f' % (total_test_loss / len(TestImgLoader)))
    # ----------------------------------------------------------------------------------
    # SAVE test information
    # savefilename = args.savemodel + 'testinformation.tar'
    # torch.save({
    #     'test_loss': total_test_loss / len(TestImgLoader),
    # }, savefilename)

if __name__ == '__main__':
    main()


import argparse
import os
import numpy as np
import cupy as cp
import chainer
from chainer import training
from chainer.training import extensions
from chainer.dataset import convert
from chainer import cuda, Variable

import dataset
import image
from nets.models import Generator, Discriminator, padding
from dataset import WaveDataset, PreEncodedDataset
from updater import Updater
from gla.gla_util import GLA

def init_gene(gpu):
    n = Generator()
    n.to_gpu()

    o = chainer.optimizers.Adam(5e-4, beta1=0.0, beta2=0.9)
    o.setup(n)
    o.add_hook(chainer.optimizer.GradientClipping(10))
    o.add_hook(chainer.optimizer.WeightDecay(1e-4))

    return n, o

def init_disc(gpu):
    n = Discriminator()
    n.to_gpu()

    o = chainer.optimizers.Adam(5e-4, beta1=0.0, beta2=0.9)
    o.setup(n)
    o.add_hook(chainer.optimizer.GradientClipping(10))
    o.add_hook(chainer.optimizer.WeightDecay(1e-4))

    return n, o

def init_dataset(path, length, is_test, batchsize):
    if path[-4:] == '.npy':
        ds = PreEncodedDataset(path, length, is_test)
    elif path[-4:] == '.wav':
        ds =  WaveDataset(path, length, is_test)
    else:
        raise Exception('file format is missing')

    return chainer.iterators.SerialIterator(ds, batchsize, shuffle=False)

def resume(args, trainer, generator_ab, generator_ba, discriminator_a, discriminator_b):
    if args.resume:
        chainer.serializers.load_npz(args.resume, trainer)
    if args.gene_ab:
        chainer.serializers.load_npz(args.gene_ab, generator_ab)
    if args.gene_ba:
        chainer.serializers.load_npz(args.gene_ba, generator_ba)
    if args.disc_a:
        chainer.serializers.load_npz(args.disc_a, discriminator_a)
    if args.disc_b:
        chainer.serializers.load_npz(args.disc_b, discriminator_b)
    if args.folder:
        folder = args.folder
        chainer.serializers.load_npz(folder+'/generator_ab.npz', generator_ab)
        chainer.serializers.load_npz(folder+'/generator_ba.npz', generator_ba)
        chainer.serializers.load_npz(folder+'/discriminator_a.npz', discriminator_a)
        chainer.serializers.load_npz(folder+'/discriminator_b.npz', discriminator_b)

def preview_convert(iterator_a, iterator_b, g_a, g_b, device, gla, dst):
    @chainer.training.make_extension()
    def make_preview(trainer):
        with chainer.using_config('train', False):
            with chainer.no_backprop_mode():
                x_a = iterator_a.next()
                x_a = convert.concat_examples(x_a, device)
                x_a = chainer.Variable(x_a)

                x_b = iterator_b.next()
                x_b = convert.concat_examples(x_b, device)
                x_b = chainer.Variable(x_b)

                x_ab = g_a(x_a)
                x_ba = g_b(x_b)

                x_bab = g_a(x_ba)
                x_aba = g_b(x_ab)

                preview_dir = '{}/preview'.format(dst)
                if not os.path.exists(preview_dir):
                    os.makedirs(preview_dir)
                image_dir = '{}/image'.format(dst)
                if not os.path.exists(image_dir):
                    os.makedirs(image_dir)

                names = ['a', 'ab', 'aba', 'b', 'ba', 'bab']
                images = [x_a, x_ab, x_aba, x_b, x_ba, x_bab]
                for n, i in zip(names, images):
                    i = cp.asnumpy(i.data)[:,:,padding:-padding,:].reshape(1, -1, 128)
                    image.save(image_dir+'/{}{}.jpg'.format(trainer.updater.epoch,n), i)
                    w = np.concatenate([gla.inverse(_i) for _i in dataset.reverse(i)])
                    dataset.save(preview_dir+'/{}{}.wav'.format(trainer.updater.epoch,n), 16000, w)

    return make_preview

def main():
    parser = argparse.ArgumentParser(description='Deep_VoiceChanger')
    parser.add_argument('--batchsize', '-b', type=int, default=32,
                        help='Number of images in each mini-batch')
    parser.add_argument('--iteration', '-i', type=int, default=100000,
                        help='Number of to train iteration')
    parser.add_argument('--gpu', '-g', type=int, default=0,
                        help='GPU ID')
    parser.add_argument('--out', '-o', default='results',
                        help='Directory to output the result')
    parser.add_argument('--resume', '-r', default='',
                        help='Resume the training from snapshot')
    parser.add_argument('--gene_ab', '-j', default='',
                        help='Resume generator a2b from file')
    parser.add_argument('--gene_ba', '-k', default='',
                        help='Resume generator b2a from file')
    parser.add_argument('--disc_a', '-m', default='',
                        help='Resume discriminator a from file')
    parser.add_argument('--disc_b', '-l', default='',
                        help='Resume discriminator b from file')
    parser.add_argument('--folder', '-f', default='',
                        help='Resume all model from foledr')
    parser.add_argument('--voice_a', '-v', default='../src/KizunaAI_long.wav',
                        help='Path of train wave file of voice a')
    parser.add_argument('--voice_b', '-w', default='../src/nekomasu_long.wav',
                        help='Path of train wave file of voice b')
    parser.add_argument('--test_a', '-s', default='../src/KizunaAI_short.wav',
                        help='Path of test wave file of voice a')
    parser.add_argument('--test_b', '-u', default='../src/nekomasu_short.wav',
                        help='Path of test wave file of voice b')
    args = parser.parse_args()

    chainer.cuda.set_max_workspace_size(256*1024*1024)
    chainer.config.type_check = False
    chainer.config.autotune = True

    if args.gpu < 0:
        print('sorry, but CPU is not recommended')
        quit()
    cp.cuda.Device(args.gpu).use()

    if args.test_a == '':
        args.test_a = args.voice_a
    if args.test_b == '':
        args.test_b = args.voice_b

    generator_ab,    opt_g_a = init_gene(args.gpu)
    generator_ba,    opt_g_b = init_gene(args.gpu)
    discriminator_a, opt_d_a = init_disc(args.gpu)
    discriminator_b, opt_d_b = init_disc(args.gpu)

    gla = GLA()

    train_iter_a = init_dataset(args.voice_a, 20000, False, args.batchsize)
    train_iter_b = init_dataset(args.voice_b, 20000, False, args.batchsize)
    test_iter_a  = init_dataset(args.test_a,  -1,    True,  16)
    test_iter_b  = init_dataset(args.test_b,  -1,    True,  16)

    updater = Updater(train_iter_a, train_iter_b, opt_g_a, opt_g_b, opt_d_a, opt_d_b, device=args.gpu)
    trainer = training.Trainer(updater, (args.iteration, 'iteration'), out=args.out)

    trainer.extend(extensions.snapshot(filename='snapshot.npz'), trigger=(10, 'epoch'))
    trainer.extend(extensions.snapshot_object(generator_ab,    'generator_ab.npz'),    trigger=(10, 'epoch'))
    trainer.extend(extensions.snapshot_object(generator_ba,    'generator_ba.npz'),    trigger=(10, 'epoch'))
    trainer.extend(extensions.snapshot_object(discriminator_a, 'discriminator_a.npz'), trigger=(10, 'epoch'))
    trainer.extend(extensions.snapshot_object(discriminator_b, 'discriminator_b.npz'), trigger=(10, 'epoch'))
    trainer.extend(extensions.LogReport())
    trainer.extend(extensions.PrintReport(['epoch', 'loss/g/recon', 'loss/g/ident', 'loss/g/gene', 'loss/d/disc', 'elapsed_time']))
    trainer.extend(extensions.ProgressBar(update_interval=5))
    trainer.extend(extensions.ExponentialShift("alpha", 0.1, optimizer=opt_g_a), trigger=(25000, 'iteration'))
    trainer.extend(extensions.ExponentialShift("alpha", 0.1, optimizer=opt_g_b), trigger=(25000, 'iteration'))
    trainer.extend(extensions.ExponentialShift("alpha", 0.1, optimizer=opt_d_a), trigger=(25000, 'iteration'))
    trainer.extend(extensions.ExponentialShift("alpha", 0.1, optimizer=opt_d_b), trigger=(25000, 'iteration'))
    trainer.extend(preview_convert(test_iter_a, test_iter_b, generator_ab, generator_ba, args.gpu, gla, args.out), trigger=(1, 'epoch'))

    resume(args, trainer, generator_ab, generator_ba, discriminator_a, discriminator_b)

    trainer.run()

if __name__ == '__main__':
    main()


#!/usr/bin/env python
from __future__ import print_function, division
import numpy as np
from rdkit import Chem
from rdkit import rdBase
from rdkit.Chem import AllChem
from rdkit import DataStructs
import rdkit.Chem.QED as QED
import scripts.sascorer as sascorer
import os
import pickle
import torch

from chemprop.train import predict
from chemprop.data import MoleculeDataset, MoleculeDataLoader
from chemprop.data.utils import get_data, get_data_from_smiles
from chemprop.utils import load_args, load_checkpoint, load_scalers

rdBase.DisableLog('rdApp.error')

class gsk3_model():
    """Scores based on an ECFP classifier for activity."""

    kwargs = ["clf_path"]
    clf_path = 'data/gsk3/gsk3.pkl'

    def __init__(self):
        with open(self.clf_path, "rb") as f:
            self.clf = pickle.load(f)

    def __call__(self, smiles_list):
        fps = []
        mask = []
        for i,smiles in enumerate(smiles_list):
            mol = Chem.MolFromSmiles(smiles)
            mask.append( int(mol is not None) )
            fp = gsk3_model.fingerprints_from_mol(mol) if mol else np.zeros((1, 2048))
            fps.append(fp)

        fps = np.concatenate(fps, axis=0)
        scores = self.clf.predict_proba(fps)[:, 1]
        scores = scores * np.array(mask)
        return np.float32(scores)

    @classmethod
    def fingerprints_from_mol(cls, mol):  # use ECFP4
        features_vec = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
        features = np.zeros((1,))
        DataStructs.ConvertToNumpyArray(features_vec, features)
        return features.reshape(1, -1)

class jnk3_model():
    """Scores based on an ECFP classifier for activity."""

    kwargs = ["clf_path"]
    clf_path = 'data/jnk3/jnk3.pkl'

    def __init__(self):
        with open(self.clf_path, "rb") as f:
            self.clf = pickle.load(f)

    def __call__(self, smiles_list):
        fps = []
        mask = []
        for i,smiles in enumerate(smiles_list):
            mol = Chem.MolFromSmiles(smiles)
            mask.append( int(mol is not None) )
            fp = jnk3_model.fingerprints_from_mol(mol) if mol else np.zeros((1, 2048))
            fps.append(fp)

        fps = np.concatenate(fps, axis=0)
        scores = self.clf.predict_proba(fps)[:, 1]
        scores = scores * np.array(mask)
        return np.float32(scores)

    @classmethod
    def fingerprints_from_mol(cls, mol):  # use ECFP4
        features_vec = AllChem.GetMorganFingerprintAsBitVect(mol, 2, nBits=2048)
        features = np.zeros((1,))
        DataStructs.ConvertToNumpyArray(features_vec, features)
        return features.reshape(1, -1)

class qed_func():

    def __call__(self, smiles_list):
        scores = []
        for smiles in smiles_list:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                scores.append(0)
            else:
                scores.append(QED.qed(mol))
        return np.float32(scores)

class sa_func():

    def __call__(self, smiles_list):
        scores = []
        for smiles in smiles_list:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                scores.append(100)
            else:
                scores.append(sascorer.calculateScore(mol))
        return np.float32(scores)

class chemprop_model():

    def __init__(self, checkpoint_dir, features_generator=None):
        self.features_generator = features_generator
        self.checkpoints, self.scalers, self.features_scalers = [], [], []
        for root, _, files in os.walk(checkpoint_dir):
            for fname in files:
                if fname.endswith('.pt'):
                    fname = os.path.join(root, fname)

                    scaler, features_scaler, _, _ = load_scalers(fname)
                    self.scalers.append(scaler)
                    self.features_scalers.append(features_scaler)

                    model = load_checkpoint(fname, device=torch.device('cpu'))
                    self.checkpoints.append(model)

    def __call__(self, smiles, batch_size=500):
        test_data = get_data_from_smiles(
            smiles=[[s] for s in smiles],
            skip_invalid_smiles=False,
            features_generator=self.features_generator
        )
        valid_indices = [i for i in range(len(test_data)) if test_data[i].mol[0] is not None]
        full_data = test_data
        test_data = MoleculeDataset([test_data[i] for i in valid_indices])
        test_data_loader = MoleculeDataLoader(dataset=test_data, batch_size=batch_size)

        sum_preds = np.zeros((len(test_data), 1))
        for model, scaler, features_scaler in zip(self.checkpoints, self.scalers, self.features_scalers):
            test_data.reset_features_and_targets()
            if features_scaler is not None:
                test_data.normalize_features(features_scaler)

            model_preds = predict(
                model=model,
                data_loader=test_data_loader,
                scaler=scaler
            )
            sum_preds += np.array(model_preds)

        # Ensemble predictions
        avg_preds = sum_preds / len(self.checkpoints)
        avg_preds = avg_preds.squeeze(-1).tolist()

        # Put zero for invalid smiles
        full_preds = [0.0] * len(full_data)
        for i, si in enumerate(valid_indices):
            full_preds[si] = avg_preds[i]

        return np.array(full_preds, dtype=np.float32)

def get_scoring_function(prop_name, features_generator=None):
    """Function that initializes and returns a scoring function by name"""
    if prop_name == 'jnk3':
        return jnk3_model()
    elif prop_name == 'gsk3':
        return gsk3_model()
    elif prop_name == 'qed':
        return qed_func()
    elif prop_name == 'sa':
        return sa_func()
    else:
        return chemprop_model(prop_name, features_generator)

if __name__ == "__main__":
    import sys
    from argparse import ArgumentParser

    parser = ArgumentParser()
    parser.add_argument('--prop', required=True)
    parser.add_argument('--features_generator', default=None, nargs='*')

    args = parser.parse_args()
    funcs = [get_scoring_function(prop, args.features_generator) for prop in args.prop.split(',')]

    data = [line.split()[:2] for line in sys.stdin]
    all_x, all_y = zip(*data)
    props = [func(all_y) for func in funcs]

    col_list = [all_x, all_y] + props
    for tup in zip(*col_list):
        print(*tup)


#!/usr/bin/python3

import os
import json
import logging
import argparse

from tqdm import tqdm
import numpy as np
import hnswlib

logger = logging.getLogger()

def str2bool(v):
    return v.lower() in ('yes', 'true', 't', '1', 'y')

def create_knn_index(args):
    # build knn index from document files in numpy format
    logger.info('Starting to load encoded numpy file')
    max_elements = args.max_elements

    p = hnswlib.Index(space=args.similarity, dim=args.dimension)
    p.init_index(max_elements=max_elements, ef_construction=args.ef_construction, M=args.M)  # parameter tuning
    idx = 0
    docid2indexid = {}

    for i in range(0, args.number_of_doc_files):
        data = np.load(args.passage_file_format.format(i))
        indices = np.load(args.indices_file_format.format(i))
        print(indices.shape)
        current_idxs = np.empty(len(indices))
        for idy, docid in enumerate(indices):
            current_idxs[idy] = idx
            docid2indexid[docid] = idx
            idx += 1
        logger.info('Starting adding current chunk of docs to knn index...')
        p.add_items(data, current_idxs)
        logger.info(f'Indexed {idx} / {max_elements} passages!')

    logger.info(f'Finished creating index added {idx} chunks, starting saving index and docid2indexid file')
    index_name = args.out_dir + f'msmarco_knn_index_M_{args.M}_efc_{args.ef_construction}.bin'
    p.save_index(index_name)
    with open(args.mapping_file, 'w+') as f:
        json.dump(docid2indexid, f)
    logger.info('Finished!')

    if args.test:
        data = np.load(args.passage_file_format.format(0))
        indices = np.load(args.indices_file_format.format(0))
        labels, distances = p.knn_query(data, k=1)
        logger.info("Recall for dataset: ", np.mean(labels.reshape(labels.shape[0]) == indices))

# convert documents into jsonl format expected by anserini to build inverted index
def convert_tsv_to_json(args):
    number_of_chunks = len([name for name in os.listdir(args.embedding_dir)])

    logger.info('Starting loading passage chunks and writing to jsonl...')
    fout = args.out_dir + 'full_msmarco_passage_collection_150_pyseriniformat.jsonl'
    with open(fout, 'w', encoding='utf8') as fout:
        for chunk_id in tqdm(range(0, number_of_chunks)):
            fin = args.embedding_dir + str(chunk_id) + '_passage_collection_150.tsv'
            with open(fin, 'r') as f:
                j = 1
                for line in f:
                    #hotfix

                    if chunk_id == 0 and 1308 <= j <= 1310:
                        j += 1
                        continue
                    j += 1
                    #hotfix end
                    split = line.split('\t')
                    pid = split[0]
                    passage = split[1].replace('"', '').replace("\\", "/").strip('\n')
                    if not isinstance(passage, str):
                        logger.info(f"pid {pid} just got skipped with passage:\n {passage}")
                        continue
                    fout.write(f'{{"id": "{str(pid)}", "contents": "{passage}"}}\n')
    logger.info('Conversion done!')

# extend the built index with remaining parts of documents
def extend_knn_index(args):
    max_elements = args.max_elements
    index_name = args.out_dir + f'msmarco_firstP_and_remainP_512_knn_M_{args.M}_efc_{args.ef_construction}.bin'
    logger.info(f'Start loading existing index from {index_name}...')
    p = hnswlib.Index(space=args.similarity, dim=args.dimension)
    p.load_index(index_name, max_elements=max_elements)
    idx = p.get_current_count()
    logger.info(f"currently the index contains {idx} elements")
    with open(args.mapping_file, "r") as f:
        docid2indexid = json.load(f)

    for i in range(0, args.number_of_doc_files):
        data = np.load(args.passage_file_format.format(i))
        indices = np.load(args.indices_file_format.format(i))
        current_idxs = np.empty(len(indices))
        for idy, docid in enumerate(indices):
            current_idxs[idy] = idx
            if docid in docid2indexid:
                docid2indexid[docid] = (docid2indexid[docid], idx)
            else:
                logger.info(f"docid {docid} was not found in mapping file...")
                docid2indexid[docid] = idx
            idx += 1
        logger.info('Starting adding current chunk of docs to knn index...')
        p.add_items(data, current_idxs)
        logger.info(f'Indexed {idx} / {max_elements} passages!')

    logger.info(f'Finished creating index added {idx} chunks, starting saving index and docid2indexid file')
    p.save_index(index_name)
    with open(args.mapping_file, 'w+') as f:
        json.dump(docid2indexid, f)
    logger.info('Finished!')

# truncate documents to 512 tokens and write in jsonl format expected by anserini
def truncate_docs(args):
    with open(args.doc_file, "r") as f, open(args.output_file, "w") as out:
        for idx, line in enumerate(f):
            if idx == 10000:
                print(f"we did {idx}")
            did, _, _, content = line.split('\t')
            tokens = content.split(" ")
            if len(tokens) > 512:
                tokens = tokens[:512]
            contents = " ".join(tokens)
            json_object = {"id": did, "contents": contents}
            out.write(json.dumps(json_object) + "\n")

if __name__ == '__main__':
    #start indexing on hadoop/nvidia:
    # python3 ./scripts/encode/build_index.py -index_type knn -out_dir ./data/indexes/ -ef_construction 100 -M 84 -similarity ip
    parser = argparse.ArgumentParser()

    parser.register('type', 'bool', str2bool)

    # knn index settings
    parser.add_argument('-similarity', type=str, default='ip', choices=['cosine', 'l2', 'ip'],
                        help='similarity score to use when knn index is chosen')
    parser.add_argument('-dimension', type=int, default=768,
                        help='dimension of the embeddings for knn index')
    parser.add_argument('-ef_construction', type=int, default=400,
                        help='hnswlib parameter, the size of the dynamic list for the nearest neighbors, higher ef'
                             ' leads to higher accuracy but slower search/construction time ')
    parser.add_argument('-M', type=int, default=64,
                        help='hnswlib parameter, the number of bi-directional links created for every new element '
                             'during construction. Range: 0-100. For embeddings 48-64 is reasonable')
    parser.add_argument('-max_elements', type=int, default=3213835)

    # run options
    parser.add_argument('-test', type='bool', default=False,
                        help='if true testing recall for knn index with querying dataset and receive top 1')
    parser.add_argument('-convert_tsv_to_json', type='bool', default=False,
                        help='convert chunks in tsv files in folder to .json files for indexing')
    parser.add_argument('-extend_index', type='bool', default=False)
    parser.add_argument('-truncate_bm25', type='bool', default=False)

    # data settings
    parser.add_argument('-embedding_dir', type=str,
                        help='path to encoded passages, should be in chunks as dicts in .json files with pid:passage')
    parser.add_argument('-out_dir', type=str, default='./data/indexes/',
                        help='output directory for the index')
    parser.add_argument('-passage_file_format', type=str, default='./data/embeddings/marco_doc_embeddings_{}.npy',
                        help='path to the passage encoding file')
    parser.add_argument('-indices_file_format', type=str,
                        default='./data/embeddings/marco_doc_embeddings_indices_{}.npy',
                        help='path to the indices for the passages')
    parser.add_argument('-number_of_doc_files', type=int, default=13)
    parser.add_argument('-mapping_file', type=str, default='./data/indexes/mapping_docid2indexid.json')
    parser.add_argument('-truncate_limit', type=int, default=512)
    parser.add_argument('-doc_file', type=str, default="./data/msmarco-docs.tsv")
    parser.add_argument('-output_file', type=str, default="./data/msmarco-docs-truncated-512.jsonl")

    logger.setLevel(logging.INFO)
    fmt = logging.Formatter('%(asctime)s: [ %(message)s ]',
                            '%m/%d/%Y %I:%M:%S %p')
    console = logging.StreamHandler()
    console.setFormatter(fmt)
    logger.addHandler(console)

    args = parser.parse_args()

    if args.convert_tsv_to_json:
        convert_tsv_to_json(args)
    elif args.extend_index:
        extend_knn_index(args)
    elif args.truncate_bm25:
        truncate_docs(args)
    else:
        create_knn_index(args)


import time, sys, cv2, json
from datetime import date
from datetime import datetime
from detection import detection
import numpy as np
import scipy.spatial.distance as scipydist
from munkres import Munkres
from tracks import track
from simulation import SimulationEngine

#initalize tracking
strikelimit = 6
threshold = 50
detections = []
delay = 1
lostTracks = []
simulation = SimulationEngine.gameEngine()

switch = 'Tracking\n0 : OFF \n1 : ON'

'''Write data into json file

Parameters
----------
filePathName: path and file Name
data: data to write

'''
def toJson(filePathName, data):

    with open(filePathName, 'w') as fp:
        json.dump(data, fp)

'''
Save trackdata

Parameters
----------
input: tracks with data to save

'''

def saveTracks(input):
    print("saving tracks...")
    now = datetime.now()
    filename = str(date.today())+ "_" + str(now.strftime("%H-%M"))
    filePathName = './' + './output' + '/' + filename + '.json'
    d = {}
    da = {}
    data = {}

    for x in range(0, len(input)-1):
        path = input[x]
        d['fps'] = path.fps
        for point in range(0, len(path.track)-1):

            vel = path.velocity[point]
            acc = path.acceleration[point]
            dirVec= path.directionVect[point]
            pos = path.track[point]
            pred = path.track[point]

            d['velocityx'] =vel[0]
            d['velocityy'] = vel[1]
            d['accelerationx'] = acc[0]
            d['accelerationx'] = acc[1]
            d['directionVectorx'] = dirVec[0]
            d['directionVectory'] = dirVec[1]
            d['datetime'] = str(path.datetime[point])
            d['position x'] =pos[0]
            d['position y'] = pos[1]
            d['predicted x'] = pred[0]
            d['predicted y'] = pred[1]
            da[point] = d
            d={}

        data[x]=da

    toJson(filePathName, data)

'''does nothing'''
def nothing(x):
    pass

'''
Assigns two point-arrays, based on euclidean distance and Munkres-algorithm.

Parameters
----------
rowMat: array 1
columnMat: array 2

Returns
----------
assignment: array of assignment indieces
distance: distance matrix
'''

def assign(rowMat, columnMat):

    distance = scipydist.cdist(rowMat, columnMat, 'euclidean') #predicted row, existing column
    dist = distance.copy()
    m = Munkres()
    return m.compute(dist), distance

'''
Swap row and column in Array of 2D-tupels.

Parameters
----------
mat: Input-array of Tupels

Returns
----------
swapped tupel-array
'''

def swap(mat):
    buf = []
    for element in mat:
        buf.append((element[1], element[0]))

    return buf

'''
This function fixes rectangular shapes and assigs two arrays.

Parameters
----------
row: array 1
column: array2

Returns
----------
assignment: 2D-Array with indices to assign arrays (array 1, array 2)
distances: 2D-Array filled with distances, of array 1 (row) and array 2 (column)
'''

def rectAssing(row, column):
    #mat1:z.b. last, mat2, existing

    if len(row) > len(column):
        assignment, distance = assign(column, row)
        assignment = swap(assignment)
        distance = np.transpose(distance)
        # sort
        assignment.sort(key=lambda x: x[0])

    else:
        assignment, distance = assign(row, column)

    return assignment, distance

'''
This function tries to find the right tracks, predict the next position and correct false estimations.
It also creates new tracks, if there are new detections.
Parameters
----------
newDetected: 2D-Array with new detections (x, y)
assignment: 2D-Array with indieces to assigned points (old detection, new detection)
distance: 2D-Array with distances between assigned points
'''

def correctPosition(newDetected = None, assignment= None, distance = None):

    cols = []

    if distance is not None:

        # row= prediction in previous detections, column = new detections

        for row, column in assignment:
            cols.append(column)

            if distance[row, column] > threshold:

                #correct false estimation and predict new postion

                correction = scipydist.pdist([detections[row].pos, newDetected[column]], 'euclidean')
                if threshold > correction:
                    distance[row, column] = correction

                    pos = newDetected[column]
                    detections[row].predict(pos)
                    detections[row].strk(False)

                #correct lost tracks, and predict new position
                #elif strikelimit >= detections[row].strike:
                else:
                    o = detections[row]
                    pred = o.pred
                    o.predict(pred)
                    o.strk()
                    cv2.circle(image, o.pos, 1, (0,0,200),20, 5)

                # one strike more

            #predict new postion
            else:

                pos = newDetected[column]
                detections[row].predict(pos)
                detections[row].strk(False)

    # appending all new
    #invert column list
    mask = np.ones(len(newDetected), np.bool)
    mask[cols] = 0

    for x in range(0, len(newDetected)):
        if mask[x]:
            mp = newDetected[x]

            detect = track(mp, len(detections), dt, fps)
            detect.predict(mp)
            detections.append(detect)

'''
This function tries to identify detections and concatenate old and new in tracks.
----------
xy: this is a 2D-Array of detected points
frame: 2D-Array of an image

Returns
----------
nothing yet

'''
def buildTracks(xy, frame):

    oldprediction = []

    #first detection
    if len(detections) == 0:
        correctPosition(xy)

    #next detections
    else:

        delete = []
        #check for lost tracks
        for x in range(0, len(detections)):
            obj = detections[x]

            if obj.strike >= strikelimit:
                delete.append(x)
            else:
                oldprediction.append(obj.pred)

        #delete lost tracks
        for delete in sorted(delete, reverse=True):
            lostTracks.append(detections[delete])
            del detections[delete]

        if xy is None:
            return

        #assign and correct new detections
        if len(detections)!=0 and len(xy)!= 0:
            assignment, distance = rectAssing(oldprediction, xy)
            correctPosition(xy, assignment, distance)

'''
This function is used to set change process-setting via userinput.

Parameters
----------
blobs: is a detection-object, we need to update the detection-setting

Returns
----------
Boolean: decision if tracking should be acitvated or not

'''
def update(blobs):
    global bt, dk, strikelimit, threshold, dt
    #mk = cv2.getTrackbarPos("Morph Kernel", "Settings")
    bt = cv2.getTrackbarPos("Bin Threshold (%)", "Settings")
    dk = cv2.getTrackbarPos("Distance Kernel", "Settings")
    strikelimit = cv2.getTrackbarPos("Reject track at:", "Settings")
    threshold = cv2.getTrackbarPos("Assignm. Distance:", "Settings")
    blobs.settings(3, bt/100, dk)

    #simulation
    asize = cv2.getTrackbarPos("number of agents", "Settings")
    simulation.agentNumbers(asize)

    return cv2.getTrackbarPos(switch, "Settings")

'''
This function is used to set initialize the trackbar for userinputs.

Parameters
----------
blobs: is a detection-object, we need to set start values
'''

def trackbar(blobs):

    cv2.namedWindow('Settings')

    wnd = 'Kernel'

    # create trackbar for simulation
    cv2.createTrackbar("number of agents", "Settings", 1, 200, nothing)

    #create Trackbar for object detection
    #cv2.createTrackbar("Morph Kernel", "Settings", 1, 20, nothing)
    cv2.createTrackbar("Bin Threshold (%)", "Settings", 2, 100, nothing)
    cv2.createTrackbar("Distance Kernel", "Settings", 1, 20, nothing)

    #create trackbar for trackbuilding
    cv2.createTrackbar(switch, "Settings", 0, 1, nothing)
    cv2.createTrackbar("Reject track at:", "Settings", 1, 20, nothing)
    cv2.createTrackbar("Assignm. Distance:", "Settings",10,100, nothing)

    #set Values
    cv2.setTrackbarPos("number of agents", "Settings", simulation.getAgentsize())

   # cv2.setTrackbarPos("Morph Kernel", "Settings", blobs.kerSize)
    cv2.setTrackbarPos("Bin Threshold (%)", "Settings", int(blobs.binaryThresh*100))
    cv2.setTrackbarPos("Distance Kernel", "Settings", blobs.distKernel)

    cv2.setTrackbarPos("Reject track at:", "Settings", strikelimit)
    cv2.setTrackbarPos("Assignm. Distance:", "Settings",threshold)
    cv2.setTrackbarPos(switch, "Settings", 1)

    #set min values
    cv2.setTrackbarMin("number of agents", "Settings", 1)

    #cv2.setTrackbarMin("Morph Kernel", "Settings", 1)
    cv2.setTrackbarMin("Bin Threshold (%)", "Settings", 1)
    cv2.setTrackbarMin("Distance Kernel", "Settings", 1)

    cv2.setTrackbarMin("Reject track at:", "Settings", 1)
    cv2.setTrackbarMin("Assignm. Distance:", "Settings", 10)

#initalize
blobs = detection()
w = h = 5
trackbar(blobs)
framebuffer = []
cap = cv2.VideoCapture("Tomaten2.mp4")
fps = cap.get(cv2.CAP_PROP_FPS)
#dt = fps*0.25/fps#fps*0.5/fps
#dt = 1
pSpeed = 1/fps
dt = pSpeed * 2
#dt*2

ret, frame = cap.read()
frNr = 0
image = frame
objects = []

#face_cascade = cv2.CascadeClassifier("haarcascade_frontalface_default.xml")

while(1):

    simulation.rederLoop()
    img = simulation.im

    #ret, img = cap.read()
    #img = cv2.blur(img,(10,10))
    framebuffer.append(img)

    trackswitch = update(blobs)

    if frNr%delay==0:
        pass
        image = framebuffer.pop(0)

    objects, dist = blobs.detectObjects(image)

    #
    # #faceversion
    # dist = image
    # gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # faces = face_cascade.detectMultiScale(gray, 1.3, 3)
    #
    # for (x, y, w, h) in faces:
    #     print(x, y)
    #     objects.append((int((x+w/2)), int((y+h/2))))
    #     cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 0), 2)
    # #####

    if trackswitch:
        buildTracks(objects, image)

    # draw stored position, predicted position and calculated velocity
    for det in detections:
        cv2.circle(image, det.pos, 5, det.color)
        cv2.circle(image, det.pred, 10, (200,200,0))
        if len(det.veloVec)>0:
            v = det.veloVec[len(det.veloVec)-2]
            cv2.arrowedLine(image, det.pos, v, (150,20,20))
        tracks = det.track
        s=0
        # if len(tracks)>7:
        #     s = len(tracks)-7

        # draw tracks
        for x in range(s, len(tracks) - 1):
            pos = tracks[x]
            pred = tracks[x + 1]
            x = int(pos[0])
            y = int(pos[1])
            px = int(pred[0])
            py = int(pred[1])
            color = det.color
            cv2.line(image, (x, y), (px, py), color)
    # draw detected objects
    if objects is not None:
        for obj in objects:
            cv2.circle(image, obj, 5, (0,0,200))

    # draw detected cores
    if dist is not None:
        pass
        cv2.imshow('Distance transformed Image', dist)

    # show image and draws
    if image is not None:
        cv2.imshow("Detected", image)
    else:
        break

    time.sleep(pSpeed)

    k = cv2.waitKey(1) & 0xFF

    if k == ord('q'):
        break

    elif k == 27:
        break
    frNr += 1
input = lostTracks
for det in detections:
    input.append(det)
saveTracks(input)
cv2.destroyAllWindows()


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
import time
import cv2
from edgetpu.detection.engine import DetectionEngine

CONFIDENCE_TH = 0.6                                                     # 顔検出閾値
MODEL_PATH = "./ssd_mobilenet_v2_face_quant_postprocess_edgetpu.tflite" # モデルパス
IMAGE_PATH = "./P_20201023_142737_BF.jpg"                               # 入力画像パス

# 画を320x320サイズに縮小・拡大（長方形の場合は余白を黒地にして正方形に変換）
def resize_320x320(image):
    size = 320
    interpolation = cv2.INTER_AREA
    (h, w) = image.shape[:2]

    #入力画像が正方形ならそのままresize
    if h == w:
        return cv2.resize(image, (size, size), interpolation)

    #入力画像が長方形ならば余白を入れて正方形に変換してからresize
    mask_size = h if h > w else w
    channel = None if len(image.shape) < 3 else image.shape[2]
    if channel is None:
        mask = np.zeros((mask_size, mask_size), dtype=image.dtype)
        mask[:h, :w] = image[:h, :w]
    else:
        mask = np.zeros((mask_size, mask_size, channel), dtype=image.dtype)
        mask[:h, :w, :] = image[:h, :w, :]

    return cv2.resize(mask, (size, size), interpolation)

if __name__ == '__main__':
    # 検出エンジン初期化
    start = time.time()
    detection_engine = DetectionEngine(MODEL_PATH)
    elapsed_time = time.time() - start
    print ("検出エンジン初期化：{:.6f}".format(elapsed_time) + "[sec]")

    # 入力画像読込み
    start = time.time()
    image = cv2.imread(IMAGE_PATH)
    elapsed_time = time.time() - start
    print ("入力画像読込み：{:.6f}".format(elapsed_time) + "[sec]")

    # 入力画像リサイズ（320x320へ変換）
    start = time.time()
    resized_image = resize_320x320(image)
    elapsed_time = time.time() - start
    print ("入力画像リサイズ：{:.6f}".format(elapsed_time) + "[sec]")

    # TPUによる顔検出
    start = time.time()
    detections = detection_engine.detect_with_input_tensor(
                input_tensor=resized_image.reshape(-1),
                threshold=CONFIDENCE_TH,
                top_k=10
            )
    (h, w) = resized_image.shape[:2]
    for detection in detections:
        box = (detection.bounding_box.flatten().tolist()) * np.array([w, h, w, h])
        (face_left, face_top, face_right, face_bottom) = box.astype('int')
        cv2.rectangle(resized_image, (face_left, face_top),(face_right, face_bottom),(255,0,0),2)
    elapsed_time = time.time() - start
    print ("TPUによる顔検出：{:.6f}".format(elapsed_time) + "[sec]")

    # 検出結果描画
    start = time.time()
    cv2.imshow("Face Detection", resized_image)
    elapsed_time = time.time() - start
    print ("検出結果描画：{:.6f}".format(elapsed_time) + "[sec]")

    # 後処理
    cv2.waitKey(0)
    cv2.destroyAllWindows()