# *****************************************************************************
#  Copyright (c) 2018, NVIDIA CORPORATION.  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.
#      * Neither the name of the NVIDIA CORPORATION nor the
#        names of its contributors may be used to endorse or promote products
#        derived from this software without specific prior written permission.
#
#  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 NVIDIA CORPORATION 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.
#
# *****************************************************************************

import numpy as np
from scipy.io.wavfile import read
import math
import torch
import os

def cosine_decay(init_val, final_val, step, decay_steps):
    alpha = final_val / init_val
    cosine_decay = 0.5 * (1 + math.cos(math.pi * step / decay_steps))
    decayed = (1 - alpha) * cosine_decay + alpha
    return init_val * decayed

def get_mask_from_lengths(lengths):
    max_len = torch.max(lengths).item()
    ids = torch.arange(0, max_len, out=torch.cuda.IntTensor(max_len))
    mask = ids < lengths.unsqueeze(1)
    return mask

def load_wav_to_torch(full_path):
    sampling_rate, data = read(full_path)
    return torch.FloatTensor(data.astype(np.float32)), sampling_rate

def load_metadata(dirname, filename='train.txt', split="|"):
    with open(os.path.join(dirname, filename)) as f:
        def split_line(line):
            parts = line.strip().split(split)
            mel_path = os.path.join(dirname, 'mels', parts[0])
            text = parts[-1]
            return mel_path, text
        return [split_line(line) for line in f.readlines()]

def to_gpu(x):
    x = x.contiguous()

    if torch.cuda.is_available():
        x = x.cuda(non_blocking=True)
    return torch.autograd.Variable(x)


import numpy as np
from rubin_sim.scheduler.surveys import BaseSurvey
import copy
import rubin_sim.scheduler.basis_functions as basis_functions
from rubin_sim.scheduler.utils import empty_observation
from rubin_sim.scheduler import features
import logging
import random
from rubin_sim.utils import ddf_locations

__all__ = ['Deep_drilling_survey', 'generate_dd_surveys', 'dd_bfs']

log = logging.getLogger(__name__)

class Deep_drilling_survey(BaseSurvey):
    """A survey class for running deep drilling fields.

    Parameters
    ----------
    basis_functions : list of rubin_sim.scheduler.basis_function objects
        These should be feasibility basis functions.
    RA : float
        The RA of the field (degrees)
    dec : float
        The dec of the field to observe (degrees)
    sequence : list of observation objects or str (rgizy)
        The sequence of observations to take. Can be a string of list of obs objects.
    nvis : list of ints
        The number of visits in each filter. Should be same length as sequence.
    survey_name : str (DD)
        The name to give this survey so it can be tracked
    reward_value : float (101.)
        The reward value to report if it is able to start (unitless).
    readtime : float (2.)
        Readout time for computing approximate time of observing the sequence. (seconds)
    flush_pad : float (30.)
        How long to hold observations in the queue after they were expected to be completed (minutes).
    """

    def __init__(self, basis_functions, RA, dec, sequence='rgizy',
                 nvis=[20, 10, 20, 26, 20],
                 exptime=30., u_exptime=30., nexp=2, ignore_obs=None, survey_name='DD',
                 reward_value=None, readtime=2., filter_change_time=120.,
                 nside=None, flush_pad=30., seed=42, detailers=None):
        super(Deep_drilling_survey, self).__init__(nside=nside, basis_functions=basis_functions,
                                                   detailers=detailers, ignore_obs=ignore_obs)
        random.seed(a=seed)

        self.ra = np.radians(RA)
        self.ra_hours = RA/360.*24.
        self.dec = np.radians(dec)
        self.survey_name = survey_name
        self.reward_value = reward_value
        self.flush_pad = flush_pad/60./24.  # To days
        self.filter_sequence = []
        if type(sequence) == str:
            self.observations = []
            for num, filtername in zip(nvis, sequence):
                for j in range(num):
                    obs = empty_observation()
                    obs['filter'] = filtername
                    if filtername == 'u':
                        obs['exptime'] = u_exptime
                    else:
                        obs['exptime'] = exptime
                    obs['RA'] = self.ra
                    obs['dec'] = self.dec
                    obs['nexp'] = nexp
                    obs['note'] = survey_name
                    self.observations.append(obs)
        else:
            self.observations = sequence

        # Let's just make this an array for ease of use
        self.observations = np.concatenate(self.observations)
        order = np.argsort(self.observations['filter'])
        self.observations = self.observations[order]

        n_filter_change = np.size(np.unique(self.observations['filter']))

        # Make an estimate of how long a seqeunce will take. Assumes no major rotational or spatial
        # dithering slowing things down.
        self.approx_time = np.sum(self.observations['exptime']+readtime*self.observations['nexp'])/3600./24. \
                           + filter_change_time*n_filter_change/3600./24.  # to days

        if self.reward_value is None:
            self.extra_features['Ntot'] = features.N_obs_survey()
            self.extra_features['N_survey'] = features.N_obs_survey(note=self.survey_name)

    def check_continue(self, observation, conditions):
        # feasibility basis functions?
        '''
        This method enables external calls to check if a given observations that belongs to this survey is
        feasible or not. This is called once a sequence has started to make sure it can continue.

        XXX--TODO:  Need to decide if we want to develope check_continue, or instead hold the
        sequence in the survey, and be able to check it that way.
        '''

        result = True

        return result

    def calc_reward_function(self, conditions):
        result = -np.inf
        if self._check_feasibility(conditions):
            if self.reward_value is not None:
                result = self.reward_value
            else:
                # XXX This might backfire if we want to have DDFs with different fractions of the
                # survey time. Then might need to define a goal fraction, and have the reward be the
                # number of observations behind that target fraction.
                result = self.extra_features['Ntot'].feature / (self.extra_features['N_survey'].feature+1)
        return result

    def generate_observations_rough(self, conditions):
        result = []
        if self._check_feasibility(conditions):
            result = copy.deepcopy(self.observations)

            # Set the flush_by
            result['flush_by_mjd'] = conditions.mjd + self.approx_time + self.flush_pad

            # remove filters that are not mounted
            mask = np.isin(result['filter'], conditions.mounted_filters)
            result = result[mask]
            # Put current loaded filter first
            ind1 = np.where(result['filter'] == conditions.current_filter)[0]
            ind2 = np.where(result['filter'] != conditions.current_filter)[0]
            result = result[ind1.tolist() + (ind2.tolist())]

            # convert to list of array. Arglebargle, don't understand why I need a reshape there
            final_result = [row.reshape(1,) for row in result]
            result = final_result

        return result

def dd_bfs(RA, dec, survey_name, ha_limits, frac_total=0.0185/2., aggressive_frac=0.011/2.,
           delays=[0., 0.5, 1.5], time_needed=62.):
    """
    Convienence function to generate all the feasibility basis functions
    """
    sun_alt_limit = -18.
    fractions = [0.00, aggressive_frac, frac_total]
    bfs = []
    bfs.append(basis_functions.Not_twilight_basis_function(sun_alt_limit=sun_alt_limit))
    bfs.append(basis_functions.Time_to_twilight_basis_function(time_needed=time_needed))
    bfs.append(basis_functions.Hour_Angle_limit_basis_function(RA=RA, ha_limits=ha_limits))
    bfs.append(basis_functions.Moon_down_basis_function())
    bfs.append(basis_functions.Fraction_of_obs_basis_function(frac_total=frac_total, survey_name=survey_name))
    bfs.append(basis_functions.Look_ahead_ddf_basis_function(frac_total, aggressive_frac,
                                                             sun_alt_limit=sun_alt_limit, time_needed=time_needed,
                                                             RA=RA, survey_name=survey_name,
                                                             ha_limits=ha_limits))
    bfs.append(basis_functions.Soft_delay_basis_function(fractions=fractions, delays=delays,
                                                         survey_name=survey_name))
    bfs.append(basis_functions.Time_to_scheduled_basis_function(time_needed=time_needed))

    return bfs

def generate_dd_surveys(nside=None, nexp=2, detailers=None, euclid_detailers=None, reward_value=100,
                        frac_total=0.0185/2., aggressive_frac=0.011/2., exptime=30, u_exptime=30,
                        nvis_master=[8, 20, 10, 20, 26, 20], delays=[0., 0.5, 1.5]):
    """Utility to return a list of standard deep drilling field surveys.

    XXX-Someone double check that I got the coordinates right!

    """

    if euclid_detailers is None:
        euclid_detailers = detailers

    surveys = []

    locations = ddf_locations()

    # ELAIS S1
    survey_name = 'DD:ELAISS1'
    RA = locations['ELAISS1'][0]
    dec = locations['ELAISS1'][1]
    ha_limits = ([0., 1.5], [21.5, 24.])
    bfs = dd_bfs(RA, dec, survey_name, ha_limits, frac_total=frac_total, aggressive_frac=aggressive_frac, delays=delays)
    surveys.append(Deep_drilling_survey(bfs, RA, dec, sequence='urgizy',
                                        nvis=nvis_master, exptime=exptime, u_exptime=u_exptime,
                                        survey_name=survey_name, reward_value=reward_value,
                                        nside=nside, nexp=nexp, detailers=detailers))

    # XMM-LSS
    survey_name = 'DD:XMM-LSS'
    RA = locations['XMM_LSS'][0]
    dec = locations['XMM_LSS'][1]
    ha_limits = ([0., 1.5], [21.5, 24.])
    bfs = dd_bfs(RA, dec, survey_name, ha_limits, frac_total=frac_total, aggressive_frac=aggressive_frac, delays=delays)

    surveys.append(Deep_drilling_survey(bfs, RA, dec, sequence='urgizy', exptime=exptime, u_exptime=u_exptime,
                                        nvis=nvis_master, survey_name=survey_name, reward_value=reward_value,
                                        nside=nside, nexp=nexp, detailers=detailers))

    # Extended Chandra Deep Field South
    survey_name = 'DD:ECDFS'
    RA = locations['ECDFS'][0]
    dec = locations['ECDFS'][1]
    ha_limits = [[0.5, 3.0], [20., 22.5]]
    bfs = dd_bfs(RA, dec, survey_name, ha_limits, frac_total=frac_total, aggressive_frac=aggressive_frac, delays=delays)
    surveys.append(Deep_drilling_survey(bfs, RA, dec, sequence='urgizy',
                                        nvis=nvis_master, exptime=exptime, u_exptime=u_exptime,
                                        survey_name=survey_name, reward_value=reward_value, nside=nside,
                                        nexp=nexp, detailers=detailers))

    # COSMOS
    survey_name = 'DD:COSMOS'
    RA = locations['COSMOS'][0]
    dec = locations['COSMOS'][1]
    ha_limits = ([0., 2.5], [21.5, 24.])
    bfs = dd_bfs(RA, dec, survey_name, ha_limits, frac_total=frac_total, aggressive_frac=aggressive_frac, delays=delays)
    surveys.append(Deep_drilling_survey(bfs, RA, dec, sequence='urgizy',
                                        nvis=nvis_master, exptime=exptime, u_exptime=u_exptime,
                                        survey_name=survey_name, reward_value=reward_value, nside=nside,
                                        nexp=nexp, detailers=detailers))

    # Euclid Fields
    # I can use the sequence kwarg to do two positions per sequence
    filters = 'urgizy'
    nviss = nvis_master
    survey_name = 'DD:EDFS'
    # Note the sequences need to be in radians since they are using observation objects directly
    # Coords from jc.cuillandre@cea.fr Oct 15, 2020
    RAs = np.radians([locations['EDFS_a'][0], locations['EDFS_b'][0]])
    decs = np.radians([locations['EDFS_a'][1], locations['EDFS_b'][1]])
    suffixes = [', a', ', b']
    sequence = []

    for filtername, nvis in zip(filters, nviss):
        for ra, dec, suffix in zip(RAs, decs, suffixes):
            for num in range(nvis):
                obs = empty_observation()
                obs['filter'] = filtername
                if filtername == 'u':
                    obs['exptime'] = u_exptime
                else:
                    obs['exptime'] = exptime
                obs['RA'] = ra
                obs['dec'] = dec
                obs['nexp'] = nexp
                obs['note'] = survey_name + suffix
                sequence.append(obs)

    ha_limits = ([0., 1.5], [22.5, 24.])
    # And back to degrees for the basis function. Need to bump up the time needed since it's a double field.
    bfs = dd_bfs(np.degrees(RAs[0]), np.degrees(decs[0]), survey_name, ha_limits,
                 frac_total=frac_total, aggressive_frac=aggressive_frac, delays=delays, time_needed=120.)
    surveys.append(Deep_drilling_survey(bfs, np.degrees(RAs), np.degrees(decs), sequence=sequence,
                                        survey_name=survey_name, reward_value=reward_value, nside=nside,
                                        nexp=nexp, detailers=euclid_detailers))

    return surveys


"""
Simulated environment wrapper
"""

import numpy as np
from base_world import BaseWorld
from builtins import range

class SimWorld(BaseWorld):

    @property
    def image(self):
        """return object position relative to the gripper"""
        return np.array(self.interface.bullet.capture_image())

    def start_world(self):
        self.setup_scene()
        self.reset_custom()
        self.interface.start()

    def reset_world(self):
        self.interface.reset()
        self.start_world()

    def setup_scene(self):
        """
        Setup the task scene given the scene specs
        """
        for tobj in self.task_objects:
            tobj.reset()

        workspace = self.scene_specs['workspace']
        objects = self.scene_specs['objects']
        # divide workspace into xy grid
        xmin, xmax = workspace['xlim']
        ymin, ymax = workspace['ylim']
        grid_size = workspace['size']

        if 'pos_eps' in workspace:
            pos_eps = workspace['pos_eps']
        else:
            pos_eps = 0

        grid_x, grid_y = np.mgrid[xmin:xmax:grid_size, ymin:ymax:grid_size]
        grid = np.vstack((grid_x.flatten(), grid_y.flatten())).T
        grid = np.random.permutation(grid)
        eps = (np.random.random(grid.shape) * 2 - 1) * pos_eps
        # [-pos_eps, pos_eps]
        grid += eps

        def get_spec(ospec):
            s = {'name': ospec['name'],
                 'model': {'filename': ospec['filename']},
                 'fixed': ospec['fixed'],
                 'boundary': ospec['boundary'],
                 'scale': ospec['scale'],
                 'pose': {'rpy': None, 'xyz': None}}
            return s

        def find_valid_pos(s, z, grid):
            for x, y in grid:
                tpos = np.array([x, y, z])
                if self.no_collision(tpos, s['boundary']):
                    return tpos
            raise ValueError('cannot find valid position')

        # resolve fixed first
        for ospec in objects:
            if ospec['pose']['type'] == 'preset':
                s = get_spec(ospec)
                s['pose']['rpy'] = ospec['pose']['rpy']
                s['pose']['xyz'] = ospec['pose']['xyz']
                self.add_instance(s)

        # resolve random
        for ospec in objects:
            if ospec['pose']['type'] == 'random':
                s = get_spec(ospec)
                s['pose']['rpy'] = ospec['pose']['rpy']
                z = ospec['pose']['z']
                s['pose']['xyz'] = find_valid_pos(s, z, grid)
                self.add_instance(s)
            elif ospec['pose']['type'] == 'random_repeat':
                repeat_range = ospec['pose']['n_repeat']
                z = ospec['pose']['z']
                n_repeat = np.random.choice(
                    np.arange(*repeat_range), size=1)[0]
                for i in range(n_repeat):
                    s = get_spec(ospec)
                    s['pose']['xyz'] = find_valid_pos(s, z, grid)
                    s['pose']['rpy'] = ospec['pose']['rpy']
                    self.add_instance(s)
            elif ospec['pose']['type'] == 'repeat':
                repeat_range = ospec['pose']['n_repeat']
                n_repeat = np.random.choice(
                    np.arange(*repeat_range), size=1)[0]
                for i in range(n_repeat):
                    s = get_spec(ospec)
                    s['pose']['xyz'] = ospec['pose']['xyz']
                    s['pose']['rpy'] = ospec['pose']['rpy']
                    self.add_instance(s)

        # resolve dependencies
        for ospec in objects:
            if ospec['pose']['type'] == 'on_top':
                s = get_spec(ospec)
                s['pose']['rpy'] = ospec['pose']['rpy']
                s['pose']['xyz'] = self.pos_on_top_of(s['boundary'],
                                                      ospec['pose']['target'])
                self.add_instance(s)

    def no_collision(self, tpos, boundary):
        """
        check if tpos has any collision with other object
        tpos:
            tpos: target position
            boundary: object boundary of the target position
        """
        for o in self.all_object_instances:
            if not self.interface.no_collision_with(tpos, boundary, o):
                return False
        return True


from __future__ import print_function, division
from PyAstronomy.pyaC import pyaPermanent as pp
from PyAstronomy.pyaC import pyaErrors as PE
import os
from .. import readFitsSpec
import numpy as np
import astropy.io.fits as pyfits

class SpectralLib:

    def refreshDescription(self):
        """
        (Re-)Download file description.txt

        This file contains an identifier and a brief description of the available spectral libraries.
        """
        self.pfs.downloadToFile(self.baseurl+"/description.txt", os.path.join(self.dpa, "description.txt"), True, self.dlv)

    def getREADME(self):
        """
        Get content of README file

        Returns
        -------
        README : list of strings
            Content of README file
        """
        rurl = self.baseurl + self.lib + "/README"
        rfn = os.path.join("pyasl", "resBased", "spectrallib", self.lib, "README")
        if not self.pfs.fileExists(rfn):
            self.pfs.downloadToFile(rurl, rfn, True, self.dlv)
        f = self.pfs.requestFile(rfn, 'rt')
        return f.readlines()

    def _readDescription(self):
        """
        Read description
        """
        self.description = {}
        ll = self.pfs.requestFile(os.path.join(self.dpa, "description.txt"), 'rt').readlines()
        for l in ll:
            l = l.rstrip("\n")
            s = l.split(":")
            if len(s) == 2:
                self.description[s[0].strip()] = s[1].strip()

    def listDescription(self):
        """
        List descriptions of available libraries
        """
        print("-"*80)
        print("--- Description of libraries")
        print("-"*80)
        for k, v in self.description.items():
            print("Key         : ", k)
            print("Description : ", v)
        print("-"*80)

    def refreshInventory(self):
        """
        Re-download inventory.txt file
        """
        self.pfs.downloadToFile(self.invurl, self.invfn, True, self.dlv)

    def _initLib(self, lib):
        """
        Initialize spectral library
        """
        # Get inventory
        self.invfn = os.path.join("pyasl", "resBased", "spectrallib", lib, "inventory.txt")
        self.invurl = self.baseurl + lib + "/inventory.txt"
        if not self.pfs.fileExists(self.invfn):
            self.pfs.downloadToFile(self.invurl, self.invfn, True, self.dlv)
        # Read inventory
        f = self.pfs.requestFile(self.invfn, 'rt')
        ll = [l.rstrip("\n") for l in f.readlines()]
        self.inv = {}
        for l in ll:
            if l.startswith("#"):
                continue
            s = l.split()
            code = []
            # First is number, last is filename
            for i in range(1,len(s)-1):
                try:
                    x = float(s[i])
                except ValueError:
                    x = s[i]
                code.append(x)
            self.inv[tuple(code)] = s[-1]

    def listInventory(self):
        """
        Print content of inventory

        The inventory is stored in the attribute `inv` (dictionary)
        """
        print("-"*80)
        print("--- Inventory")
        print("-"*80)
        for k,v in self.inv.items():
            print(k, v)
        print("-"*80)

    def read1dFitsSpec(self, fn):
        """
        Read 1d spectral model

        Returns
        -------
        w, f : arrays
            Wavelength [A] and flux
        """
        return readFitsSpec.read1dFitsSpec(fn)

    def readMuFits(self, fn):
        """
        Read specific intensity file

        Returns
        -------

        """
        hl = pyfits.open(fn)
        hkeys = hl[1].header
        w = ((np.arange(hkeys["NAXIS2"]) + 1.0) - hkeys["CRPIX1"]) * hkeys["CDELT1"] + hkeys["CRVAL1"]
        mus = np.array(hl["MUS"].data)
        specs = np.array(hl["I(mu)"].data)
        return w, mus, specs

    def requestModel(self, teff, logg, mh=0.0, nex="diskint"):
        """
        Get filename of model spectrum in the library

        Model files are downloaded on demand.

        Parameters
        ----------
        teff : float
            Effective temperature [K]
        logg : float
            Surface gravity
        mh : float
            Metallicity
        nex : string
            Name extension (e.g., diskint or muspecs)

        Returns
        -------
        Spectrum : numpy npz file
            The spectrum in the form of the context of a numpy npz file
        """
        if self.lib == "A":
            p = (float(teff), float(logg), float(mh), nex)
        if not p in self.inv:
            raise(PE.PyAValError("The key '" + str(p) + "' is not in the inventory.", \
                                 where="requestSpectrum", \
                                 solution="Use 'listInventory' to see available combinations."))
        fn = os.path.join(self.dpa, self.lib, self.inv[p])
        if not self.pfs.fileExists(fn):
            surl = self.baseurl + self.lib + "/" + self.inv[p]
            self.pfs.downloadToFile(surl, fn, True, self.dlv)
        fp = self.pfs.getFullPath(fn)
        return fp

    def __init__(self, lib="A", refreshDescr=False, dlv=True):
        """
        Access spectral library/ies

        Parameters
        ----------
        lib : string, optional
            Library key
        refreshDescr : boolean, optional
            Re-download library description file (default is False)
        dlv : boolean, optional
            Verbose download, default is True
        """

        # Download verbose
        self.dlv = dlv
        self.baseurl = "ftp://ftp.hs.uni-hamburg.de/pub/outgoing/czesla/spectrallib/"

        self.pfs = pp.pyaFS.PyAFS()
        # Data path
        self.dpa = os.path.join("pyasl", "resBased", "spectrallib")
        self.lib = lib

        self.fndescr = os.path.join(self.dpa, "description.txt")
        if (not self.pfs.fileExists(self.fndescr)) or refreshDescr:
            # Get description
            self.refreshDescription()
        self._readDescription()

        if not self.lib in self.description:
            raise(PE.PyAValError("No such library key: " + str(self.lib), \
                                 where="SpectralLib", \
                                 solution="Chose one of: " + ','.join(self.description.keys())))

        self._initLib(self.lib)


# coding: utf-8

import numpy as np
import config as cfg
import cv2
import os
import tensorflow as tf
from model.head.yolov3 import YOLOV3
try:
    import xml.etree.cElementTree as ET
except ImportError:
    import xml.etree.ElementTree as ET
import argparse
from utils import tools
from eval.evaluator import Evaluator

class Yolo_test(Evaluator):
    def __init__(self, test_weight):
        log_dir = os.path.join(cfg.LOG_DIR, 'test')
        test_weight_path = os.path.join(cfg.WEIGHTS_DIR, test_weight)

        with tf.name_scope('input'):
            input_data = tf.placeholder(dtype=tf.float32, name='input_data')
            training = tf.placeholder(dtype=tf.bool, name='training')
        _, _, _, pred_sbbox, pred_mbbox, pred_lbbox = YOLOV3(training).build_nework(input_data)
        with tf.name_scope('summary'):
            tf.summary.FileWriter(log_dir).add_graph(tf.get_default_graph())
        self.__sess = tf.Session()
        net_vars = tf.get_collection('YoloV3')
        saver = tf.train.Saver(net_vars)
        saver.restore(self.__sess, test_weight_path)
        super(Yolo_test, self).__init__(self.__sess, input_data, training, pred_sbbox, pred_mbbox, pred_lbbox)
        print("input_data.name=", input_data.name)
        print("pred_sbbox=", pred_sbbox.name)
        print("pred_mbbox=", pred_mbbox.name)
        print("pred_lbbox=", pred_lbbox.name)

    def detect_image(self, image):
        original_image = np.copy(image)
        bboxes = self.get_bbox(image)
        image = tools.draw_bbox(original_image, bboxes, self._classes)
        self.__sess.close()
        return image

    def test(self, year=2007, multi_test=False, flip_test=False):
        APs, ave_times = self.APs_voc(year, multi_test, flip_test)
        APs_file = os.path.join(self._project_path, 'eval', 'APs.txt')
        with file(APs_file, 'w') as f:
            for cls in APs:
                AP_mess = 'AP for %s = %.4f\n' % (cls, APs[cls])
                print(AP_mess.strip())
                f.write(AP_mess)
            mAP = np.mean([APs[cls] for cls in APs])
            mAP_mess = 'mAP = %.4f\n' % mAP
            print(mAP_mess.strip())
            f.write(mAP_mess)
        for key in ave_times:
            print('Average time for %s :\t%.2f ms' % (key, ave_times[key]))
        self.__sess.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Test model')
    parser.add_argument('--test_weight', help='name of test weights file', default='', type=str)
    parser.add_argument('--gpu', help='select a gpu for test', default='0', type=str)
    parser.add_argument('-mt', help='multi scale test', dest='mt', action='store_true', default=False)
    parser.add_argument('-ft', help='flip test', dest='ft', action='store_true', default=False)
    parser.add_argument('-t07', help='test voc 2007', dest='t07', action='store_true', default=False)
    parser.add_argument('-t12', help='test voc 2012', dest='t12', action='store_true', default=False)
    args = parser.parse_args()

    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

    T = Yolo_test(args.test_weight)
    if args.t07:
        T.test(2007, args.mt, args.ft)
    elif args.t12:
        T.test(2012, args.mt, args.ft)
    else:
        #test_set_path = os.path.join(cfg.DATASET_PATH, '%d_test' % 2007)
        test_set_path = '/media/xxxx/Windows/datasets/VOC/test/VOCdevkit/VOC2007'
        #img_inds_file = os.path.join(test_set_path, 'ImageSets', 'Main', 'test.txt')
        img_inds_file = '/media/xxxx/Windows/datasets/VOC/test/VOCdevkit/VOC2007/ImageSets/Main/test.txt'
        with file(img_inds_file, 'r') as f:
            txt = f.readlines()
            image_inds = [line.strip() for line in txt]
        image_ind = np.random.choice(image_inds)
        image_path = os.path.join(test_set_path, 'JPEGImages', image_ind + '.jpg')
        image = cv2.imread(image_path)
        image = T.detect_image(image)
        cv2.imwrite('detect_result.jpg', image)


import numpy as np
import torch
import torch.utils.data

class ig_dataset(torch.utils.data.Dataset):
    def __init__(self, args):
        self.data_ids = [i for i in range(1, 117)] + [i for i in range(118, 149)] + [i for i in range(150, 156)]
        print(self.data_ids)

        self.args = args

        self.data_buff_states = {}
        self.data_buff_actions = {}

        self.video_idx_list = []
        self.frame_idx_list = []

        self.action_mean = []
        self.action_max = []
        self.action_min = []
        self.steps_size = []

        self.seq_length = 10
        self.state_size = 15

        self.dataset_name = self.args.gail_experts_dir

        for item in self.data_ids:
            self.data_buff_states[item] = np.load('{0}/{1}_states.npy'.format(self.dataset_name, item))
            self.data_buff_actions[item] = np.load('{0}/{1}_actions.npy'.format(self.dataset_name, item))

            self.action_mean.append(np.mean(self.data_buff_actions[item], axis=0))
            self.action_max.append(np.max(self.data_buff_actions[item], axis=0))
            self.action_min.append(np.min(self.data_buff_actions[item], axis=0))
            self.steps_size.append(len(self.data_buff_states[item]))

            self.video_idx_list += [item for i in range(1, len(self.data_buff_states[item]) - self.seq_length - 1)]
            self.frame_idx_list += [i for i in range(1, len(self.data_buff_states[item]) - self.seq_length - 1)]

        self.length = len(self.video_idx_list)

    def __getitem__(self, index):

        video_idx = self.video_idx_list[index]
        frame_idx = self.frame_idx_list[index]

        if frame_idx >= (self.seq_length - 1):
            inputs_raw = np.array(self.data_buff_states[video_idx][(frame_idx + 1 - self.seq_length):(frame_idx + 1)])
            base_1 = inputs_raw[:, :3]
            base_2 = inputs_raw[:, 3:6]
            pos_1 = inputs_raw[:, 6:9]
            pos_2 = inputs_raw[:, 9:12]
            cut_pos_1 = inputs_raw[:, 12:15]
            cut_ori_1 = inputs_raw[:, 15:18]
            cut_pos_2 = inputs_raw[:, 18:21]
            cut_ori_2 = inputs_raw[:, 21:24]
            obj_pos = inputs_raw[:, 24:27]
            obj_ori = inputs_raw[:, 27:30]

            inputs = np.concatenate((pos_1 - base_1, pos_2 - base_2, obj_pos - pos_1, obj_pos - pos_2, obj_ori), axis=1)

        else:
            inputs_raw = np.array(self.data_buff_states[video_idx][:(frame_idx + 1)])
            base_1 = inputs_raw[:, :3]
            base_2 = inputs_raw[:, 3:6]
            pos_1 = inputs_raw[:, 6:9]
            pos_2 = inputs_raw[:, 9:12]
            cut_pos_1 = inputs_raw[:, 12:15]
            cut_ori_1 = inputs_raw[:, 15:18]
            cut_pos_2 = inputs_raw[:, 18:21]
            cut_ori_2 = inputs_raw[:, 21:24]
            obj_pos = inputs_raw[:, 24:27]
            obj_ori = inputs_raw[:, 27:30]

            inputs = np.concatenate((pos_1 - base_1, pos_2 - base_2, obj_pos - pos_1, obj_pos - pos_2, obj_ori), axis=1)

            padding = np.array([[0.0 for i in range(self.state_size)] for _ in range(self.seq_length - len(inputs))])
            inputs = np.concatenate(([padding, inputs]), axis=0)

        action_gt = torch.FloatTensor(np.array(self.data_buff_actions[video_idx][frame_idx])).view(-1)
        action_gt[:3] *= 1000.0
        action_gt[7:10] *= 1000.0

        inputs = torch.FloatTensor(inputs)
        inputs = inputs.view(-1)

        vid_id = torch.FloatTensor(np.array([video_idx]))

        return inputs, action_gt, vid_id

    def __len__(self):
        return self.length


from absl import app, flags, logging
from absl.flags import FLAGS
import numpy as np
from yolov3_tf2.models import YoloV3, YoloV3Tiny
from yolov3_tf2.utils import load_darknet_weights
import tensorflow as tf

flags.DEFINE_string('weights', './data/yolov3.weights', 'path to weights file')
flags.DEFINE_string('output', './checkpoints/yolov3.tf', 'path to output')
flags.DEFINE_boolean('tiny', False, 'yolov3 or yolov3-tiny')
flags.DEFINE_integer('num_classes', 80, 'number of classes in the model')

def main(_argv):
    # physical_devices = tf.config.experimental.list_physical_devices('GPU')
    # if len(physical_devices) > 0:
    #     tf.config.experimental.set_memory_growth(physical_devices[0], True)

    if FLAGS.tiny:
        yolo = YoloV3Tiny(classes=FLAGS.num_classes)
    else:
        yolo = YoloV3(classes=FLAGS.num_classes)
    yolo.summary()
    logging.info('model created')

    load_darknet_weights(yolo, FLAGS.weights, FLAGS.tiny)
    logging.info('weights loaded')

    img = np.random.random((1, 320, 320, 3)).astype(np.float32)
    output = yolo(img)
    logging.info('sanity check passed')

    yolo.save_weights(FLAGS.output)
    logging.info('weights saved')

if __name__ == '__main__':
    try:
        app.run(main)
    except SystemExit:
        pass


import time
import warnings
from abc import abstractmethod
from itertools import product

import numpy as np
from pyspark.sql import SparkSession
from sklearn.model_selection._split import check_cv
from sklearn.model_selection._validation import _fit_and_score
from sklearn.base import is_classifier, clone
from sklearn.metrics.scorer import _check_multimetric_scoring
from sklearn.model_selection._search import BaseSearchCV, _check_param_grid, ParameterGrid, ParameterSampler
from sklearn.utils._joblib import Parallel, delayed
from sklearn.utils.validation import indexable

class SparkBaseSearchCV(BaseSearchCV):
    """Abstract base class for spark accelerated hyper parameter search with
    cross-validation.
    """

    @abstractmethod
    def __init__(self, spark, estimator, scoring=None, n_jobs=None, iid='warn',
                 refit=True, cv='warn', verbose=0, pre_dispatch='2*n_jobs',
                 error_score='raise-deprecating', return_train_score=True):

        super(SparkBaseSearchCV, self).__init__(
            estimator, scoring=scoring, n_jobs=n_jobs, iid=iid,
            refit=refit, cv=cv, verbose=verbose, pre_dispatch=pre_dispatch,
            error_score=error_score, return_train_score=return_train_score)

        if isinstance(spark, SparkSession):
            self.spark = spark
        elif spark:
            self.spark = SparkSession.builder.getOrCreate()
        else:
            self.spark = None

    def _run_search(self, evaluate_candidates):
        """Repeatedly calls `evaluate_candidates` to conduct a search.

        This method, implemented in sub-classes, makes it possible to
        customize the the scheduling of evaluations: GridSearchCV and
        RandomizedSearchCV schedule evaluations for their whole parameter
        search space at once but other more sequential approaches are also
        possible: for instance is possible to iteratively schedule evaluations
        for new regions of the parameter search space based on previously
        collected evaluation results. This makes it possible to implement
        Bayesian optimization or more generally sequential model-based
        optimization by deriving from the BaseSearchCV abstract base class.

        Parameters
        ----------
        evaluate_candidates : callable
            This callback accepts a list of candidates, where each candidate is
            a dict of parameter settings. It returns a dict of all results so
            far, formatted like ``cv_results_``.

        Examples
        --------

        ::

            def _run_search(self, evaluate_candidates):
                'Try C=0.1 only if C=1 is better than C=10'
                all_results = evaluate_candidates([{'C': 1}, {'C': 10}])
                score = all_results['mean_test_score']
                if score[0] < score[1]:
                    evaluate_candidates([{'C': 0.1}])
        """
        raise NotImplementedError("_run_search not implemented.")

    def fit(self, X, y=None, groups=None, **fit_params):
        """Run fit with all sets of parameters.

        Parameters
        ----------

        X : array-like, shape = [n_samples, n_features]
            Training vector, where n_samples is the number of samples and
            n_features is the number of features.

        y : array-like, shape = [n_samples] or [n_samples, n_output], optional
            Target relative to X for classification or regression;
            None for unsupervised learning.

        groups : array-like, with shape (n_samples,), optional
            Group labels for the samples used while splitting the dataset into
            train/test set. Only used in conjunction with a "Group" `cv`
            instance (e.g., `GroupKFold`).

        **fit_params : dict of string -> object
            Parameters passed to the ``fit`` method of the estimator
        """
        estimator = self.estimator
        cv = check_cv(self.cv, y, classifier=is_classifier(estimator))

        scorers, self.multimetric_ = _check_multimetric_scoring(
            self.estimator, scoring=self.scoring)

        if self.multimetric_:
            if self.refit is not False and (
                    not isinstance(self.refit, str) or
                    # This will work for both dict / list (tuple)
                    self.refit not in scorers) and not callable(self.refit):
                raise ValueError("For multi-metric scoring, the parameter "
                                 "refit must be set to a scorer key or a "
                                 "callable to refit an estimator with the "
                                 "best parameter setting on the whole "
                                 "data and make the best_* attributes "
                                 "available for that metric. If this is "
                                 "not needed, refit should be set to "
                                 "False explicitly. %r was passed."
                                 % self.refit)
            else:
                refit_metric = self.refit
        else:
            refit_metric = 'score'

        X, y, groups = indexable(X, y, groups)
        n_splits = cv.get_n_splits(X, y, groups)

        base_estimator = clone(self.estimator)

        fit_and_score_kwargs = dict(scorer=scorers,
                                    fit_params=fit_params,
                                    return_train_score=self.return_train_score,
                                    return_n_test_samples=True,
                                    return_times=True,
                                    return_parameters=False,
                                    error_score=self.error_score,
                                    verbose=self.verbose)

        if self.spark is None:
            fitting_function = self._run_sklearn_fit
        else:
            fitting_function = self._run_skspark_fit

        results = fitting_function(
            base_estimator=base_estimator, X=X, y=y, scorers=scorers, cv=cv,
            groups=groups, n_splits=n_splits,
            fit_and_score_kwargs=fit_and_score_kwargs
        )

        # For multi-metric evaluation, store the best_index_, best_params_ and
        # best_score_ iff refit is one of the scorer names
        # In single metric evaluation, refit_metric is "score"
        if self.refit or not self.multimetric_:
            # If callable, refit is expected to return the index of the best
            # parameter set.
            if callable(self.refit):
                self.best_index_ = self.refit(results)
                if not isinstance(self.best_index_, (int, np.integer)):
                    raise TypeError('best_index_ returned is not an integer')
                if (self.best_index_ < 0 or
                   self.best_index_ >= len(results["params"])):
                    raise IndexError('best_index_ index out of range')
            else:
                self.best_index_ = results["rank_test_%s"
                                           % refit_metric].argmin()
                self.best_score_ = results["mean_test_%s" % refit_metric][
                                           self.best_index_]
            self.best_params_ = results["params"][self.best_index_]

        if self.refit:
            self.best_estimator_ = clone(base_estimator).set_params(
                **self.best_params_)
            refit_start_time = time.time()
            if y is not None:
                self.best_estimator_.fit(X, y, **fit_params)
            else:
                self.best_estimator_.fit(X, **fit_params)
            refit_end_time = time.time()
            self.refit_time_ = refit_end_time - refit_start_time

        # Store the only scorer not as a dict for single metric evaluation
        self.scorer_ = scorers if self.multimetric_ else scorers['score']

        self.cv_results_ = results
        self.n_splits_ = n_splits

        return self

    def _run_sklearn_fit(self, base_estimator, X, y, scorers, cv, groups,
                         n_splits, fit_and_score_kwargs):

        parallel = Parallel(n_jobs=self.n_jobs, verbose=self.verbose,
                            pre_dispatch=self.pre_dispatch)
        results = {}

        with parallel:
            all_candidate_params = []
            all_out = []

            def evaluate_candidates(candidate_params):
                candidate_params = list(candidate_params)
                n_candidates = len(candidate_params)

                if self.verbose > 0:
                    print("Fitting {0} folds for each of {1} candidates,"
                          " totalling {2} fits".format(
                              n_splits, n_candidates, n_candidates * n_splits))

                out = parallel(delayed(_fit_and_score)(clone(base_estimator),
                                                       X, y,
                                                       train=train, test=test,
                                                       parameters=parameters,
                                                       **fit_and_score_kwargs)
                               for parameters, (train, test)
                               in product(candidate_params,
                                          cv.split(X, y, groups)))

                self._validate_out(out, n_candidates, n_splits)

                all_candidate_params.extend(candidate_params)
                all_out.extend(out)

                nonlocal results
                results = self._format_results(
                    all_candidate_params, scorers, n_splits, all_out)
                return results

            self._run_search(evaluate_candidates)
        return results

    def _run_skspark_fit(self, base_estimator, X, y, scorers, cv, groups,
                         n_splits, fit_and_score_kwargs):
        all_candidate_params = []
        all_out = []
        results = {}

        def evaluate_candidates(candidate_params):
            candidate_params = list(candidate_params)
            n_candidates = len(candidate_params)

            param_grid = [(parameters, train, test)
                          for parameters, (train, test)
                          in product(candidate_params, cv.split(X, y, groups))]

            # empty grids fail during execution, so need to apply check here
            self._validate_out(param_grid, n_candidates, n_splits)

            # Because the original python code expects a certain order for the
            # elements, we need to respect it.
            indexed_param_grid = list(zip(range(len(param_grid)), param_grid))
            par_param_grid = self.spark.sparkContext.parallelize(
                indexed_param_grid, len(indexed_param_grid))

            X_bc = self.spark.sparkContext.broadcast(X)
            y_bc = self.spark.sparkContext.broadcast(y)

            def spark_task(tup):
                (index, (parameters, train, test)) = tup
                local_estimator = clone(base_estimator)
                local_X = X_bc.value
                local_y = y_bc.value

                error = None
                with warnings.catch_warnings(record=True) as warns:
                    try:
                        res = _fit_and_score(
                            local_estimator, X=local_X, y=local_y, train=train,
                            test=test, parameters=parameters,
                            **fit_and_score_kwargs)
                    except Exception as e:
                        res = None
                        error = e

                return index, {"results": res, "errors": error, "warnings": warns}

            indexed_out0 = dict(par_param_grid.map(spark_task).collect())

            # process errors and warnings
            for i in range(len(param_grid)):
                error_to_be_raised = indexed_out0[i]["errors"]
                if error_to_be_raised is not None:
                    raise error_to_be_raised

                for warn in indexed_out0[i]["warnings"]:
                    warnings.warn(warn.message, warn.category)

            out = [indexed_out0[i]["results"] for i in range(len(param_grid))]

            X_bc.unpersist()
            y_bc.unpersist()

            self._validate_out(out, n_candidates, n_splits)

            all_out.extend(out)
            all_candidate_params.extend(candidate_params)

            nonlocal results
            results = self._format_results(
                all_candidate_params, scorers, n_splits, all_out)
            return results

        self._run_search(evaluate_candidates)
        return results

    @staticmethod
    def _validate_out(out, n_candidates, n_splits):
        if len(out) < 1:
            raise ValueError('No fits were performed. '
                             'Was the CV iterator empty? '
                             'Were there no candidates?')
        elif len(out) != n_candidates * n_splits:
            raise ValueError('cv.split and cv.get_n_splits returned '
                             'inconsistent results. Expected {} '
                             'splits, got {}'
                             .format(n_splits,
                                     len(out) // n_candidates))

    def __getstate__(self):
        """To not try to pickle the non-serializable SparkSession"""
        attributes = dict(self.__dict__)
        if 'spark' in attributes.keys():
            del attributes['spark']
        return attributes

class GridSearchCV(SparkBaseSearchCV):
    """Exhaustive search over specified parameter values for an estimator.

    Important members are fit, predict.

    GridSearchCV implements a "fit" and a "score" method.
    It also implements "predict", "predict_proba", "decision_function",
    "transform" and "inverse_transform" if they are implemented in the
    estimator used.

    The parameters of the estimator used to apply these methods are optimized
    by cross-validated grid-search over a parameter grid.

    Read more in the :ref:`User Guide <grid_search>`.

    Parameters
    ----------
    estimator : estimator object.
        This is assumed to implement the scikit-learn estimator interface.
        Either estimator needs to provide a ``score`` function,
        or ``scoring`` must be passed.

    param_grid : dict or list of dictionaries
        Dictionary with parameters names (string) as keys and lists of
        parameter settings to try as values, or a list of such
        dictionaries, in which case the grids spanned by each dictionary
        in the list are explored. This enables searching over any sequence
        of parameter settings.

    scoring : string, callable, list/tuple, dict or None, default: None
        A single string (see :ref:`scoring_parameter`) or a callable
        (see :ref:`scoring`) to evaluate the predictions on the test set.

        For evaluating multiple metrics, either give a list of (unique) strings
        or a dict with names as keys and callables as values.

        NOTE that when using custom scorers, each scorer should return a single
        value. Metric functions returning a list/array of values can be wrapped
        into multiple scorers that return one value each.

        See :ref:`multimetric_grid_search` for an example.

        If None, the estimator's score method is used.

    n_jobs : int or None, optional (default=None)
        Number of jobs to run in parallel.
        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
        for more details.

    pre_dispatch : int, or string, optional
        Controls the number of jobs that get dispatched during parallel
        execution. Reducing this number can be useful to avoid an
        explosion of memory consumption when more jobs get dispatched
        than CPUs can process. This parameter can be:

            - None, in which case all the jobs are immediately
              created and spawned. Use this for lightweight and
              fast-running jobs, to avoid delays due to on-demand
              spawning of the jobs

            - An int, giving the exact number of total jobs that are
              spawned

            - A string, giving an expression as a function of n_jobs,
              as in '2*n_jobs'

    iid : boolean, default='warn'
        If True, return the average score across folds, weighted by the number
        of samples in each test set. In this case, the data is assumed to be
        identically distributed across the folds, and the loss minimized is
        the total loss per sample, and not the mean loss across the folds. If
        False, return the average score across folds. Default is True, but
        will change to False in version 0.22, to correspond to the standard
        definition of cross-validation.

        .. versionchanged:: 0.20
            Parameter ``iid`` will change from True to False by default in
            version 0.22, and will be removed in 0.24.

    cv : int, cross-validation generator or an iterable, optional
        Determines the cross-validation splitting strategy.
        Possible inputs for cv are:

        - None, to use the default 3-fold cross validation,
        - integer, to specify the number of folds in a `(Stratified)KFold`,
        - :term:`CV splitter`,
        - An iterable yielding (train, test) splits as arrays of indices.

        For integer/None inputs, if the estimator is a classifier and ``y`` is
        either binary or multiclass, :class:`StratifiedKFold` is used. In all
        other cases, :class:`KFold` is used.

        Refer :ref:`User Guide <cross_validation>` for the various
        cross-validation strategies that can be used here.

        .. versionchanged:: 0.20
            ``cv`` default value if None will change from 3-fold to 5-fold
            in v0.22.

    refit : boolean, string, or callable, default=True
        Refit an estimator using the best found parameters on the whole
        dataset.

        For multiple metric evaluation, this needs to be a string denoting the
        scorer that would be used to find the best parameters for refitting
        the estimator at the end.

        Where there are considerations other than maximum score in
        choosing a best estimator, ``refit`` can be set to a function which
        returns the selected ``best_index_`` given ``cv_results_``.

        The refitted estimator is made available at the ``best_estimator_``
        attribute and permits using ``predict`` directly on this
        ``GridSearchCV`` instance.

        Also for multiple metric evaluation, the attributes ``best_index_``,
        ``best_score_`` and ``best_params_`` will only be available if
        ``refit`` is set and all of them will be determined w.r.t this specific
        scorer. ``best_score_`` is not returned if refit is callable.

        See ``scoring`` parameter to know more about multiple metric
        evaluation.

        .. versionchanged:: 0.20
            Support for callable added.

    verbose : integer
        Controls the verbosity: the higher, the more messages.

    error_score : 'raise' or numeric
        Value to assign to the score if an error occurs in estimator fitting.
        If set to 'raise', the error is raised. If a numeric value is given,
        FitFailedWarning is raised. This parameter does not affect the refit
        step, which will always raise the error. Default is 'raise' but from
        version 0.22 it will change to np.nan.

    return_train_score : boolean, default=False
        If ``False``, the ``cv_results_`` attribute will not include training
        scores.
        Computing training scores is used to get insights on how different
        parameter settings impact the overfitting/underfitting trade-off.
        However computing the scores on the training set can be computationally
        expensive and is not strictly required to select the parameters that
        yield the best generalization performance.

    spark : SparkSession or bool
        If True use Spark (getting the exising spark session or creating one
        if there isn't one running). If False doesn't use Spark acceleration.
        Alternatively pass a SparkSession object.

    Examples
    --------
    >>> from sklearn import svm, datasets
    >>> from skspark.model_selection import GridSearchCV
    >>> iris = datasets.load_iris()
    >>> parameters = {'kernel':('linear', 'rbf'), 'C':[1, 10]}
    >>> svc = svm.SVC(gamma="scale")
    >>> clf = GridSearchCV(svc, parameters, cv=5, spark=True)
    >>> clf.fit(iris.data, iris.target)
    ...                             # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
    GridSearchCV(cv=5, error_score=...,
           estimator=SVC(C=1.0, cache_size=..., class_weight=..., coef0=...,
                         decision_function_shape='ovr', degree=..., gamma=...,
                         kernel='rbf', max_iter=-1, probability=False,
                         random_state=None, shrinking=True, tol=...,
                         verbose=False),
           iid=..., n_jobs=None,
           param_grid=..., pre_dispatch=..., refit=..., return_train_score=...,
           scoring=..., verbose=...)
    >>> sorted(clf.cv_results_.keys())
    ...                             # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
    ['mean_fit_time', 'mean_score_time', 'mean_test_score',...
     'param_C', 'param_kernel', 'params',...
     'rank_test_score', 'split0_test_score',...
     'split2_test_score', ...
     'std_fit_time', 'std_score_time', 'std_test_score']

    Attributes
    ----------
    cv_results_ : dict of numpy (masked) ndarrays
        A dict with keys as column headers and values as columns, that can be
        imported into a pandas ``DataFrame``.

        For instance the below given table

        +------------+-----------+------------+-----------------+---+---------+
        |param_kernel|param_gamma|param_degree|split0_test_score|...|rank_t...|
        +============+===========+============+=================+===+=========+
        |  'poly'    |     --    |      2     |       0.80      |...|    2    |
        +------------+-----------+------------+-----------------+---+---------+
        |  'poly'    |     --    |      3     |       0.70      |...|    4    |
        +------------+-----------+------------+-----------------+---+---------+
        |  'rbf'     |     0.1   |     --     |       0.80      |...|    3    |
        +------------+-----------+------------+-----------------+---+---------+
        |  'rbf'     |     0.2   |     --     |       0.93      |...|    1    |
        +------------+-----------+------------+-----------------+---+---------+

        will be represented by a ``cv_results_`` dict of::

            {
            'param_kernel': masked_array(data = ['poly', 'poly', 'rbf', 'rbf'],
                                         mask = [False False False False]...)
            'param_gamma': masked_array(data = [-- -- 0.1 0.2],
                                        mask = [ True  True False False]...),
            'param_degree': masked_array(data = [2.0 3.0 -- --],
                                         mask = [False False  True  True]...),
            'split0_test_score'  : [0.80, 0.70, 0.80, 0.93],
            'split1_test_score'  : [0.82, 0.50, 0.70, 0.78],
            'mean_test_score'    : [0.81, 0.60, 0.75, 0.85],
            'std_test_score'     : [0.01, 0.10, 0.05, 0.08],
            'rank_test_score'    : [2, 4, 3, 1],
            'split0_train_score' : [0.80, 0.92, 0.70, 0.93],
            'split1_train_score' : [0.82, 0.55, 0.70, 0.87],
            'mean_train_score'   : [0.81, 0.74, 0.70, 0.90],
            'std_train_score'    : [0.01, 0.19, 0.00, 0.03],
            'mean_fit_time'      : [0.73, 0.63, 0.43, 0.49],
            'std_fit_time'       : [0.01, 0.02, 0.01, 0.01],
            'mean_score_time'    : [0.01, 0.06, 0.04, 0.04],
            'std_score_time'     : [0.00, 0.00, 0.00, 0.01],
            'params'             : [{'kernel': 'poly', 'degree': 2}, ...],
            }

        NOTE

        The key ``'params'`` is used to store a list of parameter
        settings dicts for all the parameter candidates.

        The ``mean_fit_time``, ``std_fit_time``, ``mean_score_time`` and
        ``std_score_time`` are all in seconds.

        For multi-metric evaluation, the scores for all the scorers are
        available in the ``cv_results_`` dict at the keys ending with that
        scorer's name (``'_<scorer_name>'``) instead of ``'_score'`` shown
        above. ('split0_test_precision', 'mean_train_precision' etc.)

    best_estimator_ : estimator or dict
        Estimator that was chosen by the search, i.e. estimator
        which gave highest score (or smallest loss if specified)
        on the left out data. Not available if ``refit=False``.

        See ``refit`` parameter for more information on allowed values.

    best_score_ : float
        Mean cross-validated score of the best_estimator

        For multi-metric evaluation, this is present only if ``refit`` is
        specified.

    best_params_ : dict
        Parameter setting that gave the best results on the hold out data.

        For multi-metric evaluation, this is present only if ``refit`` is
        specified.

    best_index_ : int
        The index (of the ``cv_results_`` arrays) which corresponds to the best
        candidate parameter setting.

        The dict at ``search.cv_results_['params'][search.best_index_]`` gives
        the parameter setting for the best model, that gives the highest
        mean score (``search.best_score_``).

        For multi-metric evaluation, this is present only if ``refit`` is
        specified.

    scorer_ : function or a dict
        Scorer function used on the held out data to choose the best
        parameters for the model.

        For multi-metric evaluation, this attribute holds the validated
        ``scoring`` dict which maps the scorer key to the scorer callable.

    n_splits_ : int
        The number of cross-validation splits (folds/iterations).

    refit_time_ : float
        Seconds used for refitting the best model on the whole dataset.

        This is present only if ``refit`` is not False.

    Notes
    -----
    The parameters selected are those that maximize the score of the left out
    data, unless an explicit score is passed in which case it is used instead.

    If `n_jobs` was set to a value higher than one, the data is copied for each
    point in the grid (and not `n_jobs` times). This is done for efficiency
    reasons if individual jobs take very little time, but may raise errors if
    the dataset is large and not enough memory is available.  A workaround in
    this case is to set `pre_dispatch`. Then, the memory is copied only
    `pre_dispatch` many times. A reasonable value for `pre_dispatch` is `2 *
    n_jobs`.

    See Also
    ---------
    :class:`ParameterGrid`:
        generates all the combinations of a hyperparameter grid.

    :func:`sklearn.model_selection.train_test_split`:
        utility function to split the data into a development set usable
        for fitting a GridSearchCV instance and an evaluation set for
        its final evaluation.

    :func:`sklearn.metrics.make_scorer`:
        Make a scorer from a performance metric or loss function.

    """
    _required_parameters = ["estimator", "param_grid"]

    def __init__(self, estimator, param_grid, scoring=None,
                 n_jobs=None, iid='warn', refit=True, cv='warn', verbose=0,
                 pre_dispatch='2*n_jobs', error_score='raise-deprecating',
                 return_train_score=False, spark=True):
        super().__init__(
            spark=spark, estimator=estimator, scoring=scoring,
            n_jobs=n_jobs, iid=iid, refit=refit, cv=cv, verbose=verbose,
            pre_dispatch=pre_dispatch, error_score=error_score,
            return_train_score=return_train_score)
        self.param_grid = param_grid
        _check_param_grid(param_grid)

    def _run_search(self, evaluate_candidates):
        """Search all candidates in param_grid"""
        evaluate_candidates(ParameterGrid(self.param_grid))

class RandomizedSearchCV(SparkBaseSearchCV):
    """Randomized search on hyper parameters.

    RandomizedSearchCV implements a "fit" and a "score" method.
    It also implements "predict", "predict_proba", "decision_function",
    "transform" and "inverse_transform" if they are implemented in the
    estimator used.

    The parameters of the estimator used to apply these methods are optimized
    by cross-validated search over parameter settings.

    In contrast to GridSearchCV, not all parameter values are tried out, but
    rather a fixed number of parameter settings is sampled from the specified
    distributions. The number of parameter settings that are tried is
    given by n_iter.

    If all parameters are presented as a list,
    sampling without replacement is performed. If at least one parameter
    is given as a distribution, sampling with replacement is used.
    It is highly recommended to use continuous distributions for continuous
    parameters.

    Note that before SciPy 0.16, the ``scipy.stats.distributions`` do not
    accept a custom RNG instance and always use the singleton RNG from
    ``numpy.random``. Hence setting ``random_state`` will not guarantee a
    deterministic iteration whenever ``scipy.stats`` distributions are used to
    define the parameter search space.

    Read more in the :ref:`User Guide <randomized_parameter_search>`.

    Parameters
    ----------
    estimator : estimator object.
        A object of that type is instantiated for each grid point.
        This is assumed to implement the scikit-learn estimator interface.
        Either estimator needs to provide a ``score`` function,
        or ``scoring`` must be passed.

    param_distributions : dict
        Dictionary with parameters names (string) as keys and distributions
        or lists of parameters to try. Distributions must provide a ``rvs``
        method for sampling (such as those from scipy.stats.distributions).
        If a list is given, it is sampled uniformly.

    n_iter : int, default=10
        Number of parameter settings that are sampled. n_iter trades
        off runtime vs quality of the solution.

    scoring : string, callable, list/tuple, dict or None, default: None
        A single string (see :ref:`scoring_parameter`) or a callable
        (see :ref:`scoring`) to evaluate the predictions on the test set.

        For evaluating multiple metrics, either give a list of (unique) strings
        or a dict with names as keys and callables as values.

        NOTE that when using custom scorers, each scorer should return a single
        value. Metric functions returning a list/array of values can be wrapped
        into multiple scorers that return one value each.

        See :ref:`multimetric_grid_search` for an example.

        If None, the estimator's score method is used.

    n_jobs : int or None, optional (default=None)
        Number of jobs to run in parallel.
        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
        for more details.

    pre_dispatch : int, or string, optional
        Controls the number of jobs that get dispatched during parallel
        execution. Reducing this number can be useful to avoid an
        explosion of memory consumption when more jobs get dispatched
        than CPUs can process. This parameter can be:

            - None, in which case all the jobs are immediately
              created and spawned. Use this for lightweight and
              fast-running jobs, to avoid delays due to on-demand
              spawning of the jobs

            - An int, giving the exact number of total jobs that are
              spawned

            - A string, giving an expression as a function of n_jobs,
              as in '2*n_jobs'

    iid : boolean, default='warn'
        If True, return the average score across folds, weighted by the number
        of samples in each test set. In this case, the data is assumed to be
        identically distributed across the folds, and the loss minimized is
        the total loss per sample, and not the mean loss across the folds. If
        False, return the average score across folds. Default is True, but
        will change to False in version 0.22, to correspond to the standard
        definition of cross-validation.

        .. versionchanged:: 0.20
            Parameter ``iid`` will change from True to False by default in
            version 0.22, and will be removed in 0.24.

    cv : int, cross-validation generator or an iterable, optional
        Determines the cross-validation splitting strategy.
        Possible inputs for cv are:

        - None, to use the default 3-fold cross validation,
        - integer, to specify the number of folds in a `(Stratified)KFold`,
        - :term:`CV splitter`,
        - An iterable yielding (train, test) splits as arrays of indices.

        For integer/None inputs, if the estimator is a classifier and ``y`` is
        either binary or multiclass, :class:`StratifiedKFold` is used. In all
        other cases, :class:`KFold` is used.

        Refer :ref:`User Guide <cross_validation>` for the various
        cross-validation strategies that can be used here.

        .. versionchanged:: 0.20
            ``cv`` default value if None will change from 3-fold to 5-fold
            in v0.22.

    refit : boolean, string, or callable, default=True
        Refit an estimator using the best found parameters on the whole
        dataset.

        For multiple metric evaluation, this needs to be a string denoting the
        scorer that would be used to find the best parameters for refitting
        the estimator at the end.

        Where there are considerations other than maximum score in
        choosing a best estimator, ``refit`` can be set to a function which
        returns the selected ``best_index_`` given the ``cv_results``.

        The refitted estimator is made available at the ``best_estimator_``
        attribute and permits using ``predict`` directly on this
        ``RandomizedSearchCV`` instance.

        Also for multiple metric evaluation, the attributes ``best_index_``,
        ``best_score_`` and ``best_params_`` will only be available if
        ``refit`` is set and all of them will be determined w.r.t this specific
        scorer. When refit is callable, ``best_score_`` is disabled.

        See ``scoring`` parameter to know more about multiple metric
        evaluation.

        .. versionchanged:: 0.20
            Support for callable added.

    verbose : integer
        Controls the verbosity: the higher, the more messages.

    random_state : int, RandomState instance or None, optional, default=None
        Pseudo random number generator state used for random uniform sampling
        from lists of possible values instead of scipy.stats distributions.
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    error_score : 'raise' or numeric
        Value to assign to the score if an error occurs in estimator fitting.
        If set to 'raise', the error is raised. If a numeric value is given,
        FitFailedWarning is raised. This parameter does not affect the refit
        step, which will always raise the error. Default is 'raise' but from
        version 0.22 it will change to np.nan.

    return_train_score : boolean, default=False
        If ``False``, the ``cv_results_`` attribute will not include training
        scores.
        Computing training scores is used to get insights on how different
        parameter settings impact the overfitting/underfitting trade-off.
        However computing the scores on the training set can be computationally
        expensive and is not strictly required to select the parameters that
        yield the best generalization performance.

    spark : SparkSession or bool
        If True use Spark (getting the exising spark session or creating one
        if there isn't one running). If False doesn't use Spark acceleration.
        Alternatively pass a SparkSession object.

    Attributes
    ----------
    cv_results_ : dict of numpy (masked) ndarrays
        A dict with keys as column headers and values as columns, that can be
        imported into a pandas ``DataFrame``.

        For instance the below given table

        +--------------+-------------+-------------------+---+---------------+
        | param_kernel | param_gamma | split0_test_score |...|rank_test_score|
        +==============+=============+===================+===+===============+
        |    'rbf'     |     0.1     |       0.80        |...|       2       |
        +--------------+-------------+-------------------+---+---------------+
        |    'rbf'     |     0.2     |       0.90        |...|       1       |
        +--------------+-------------+-------------------+---+---------------+
        |    'rbf'     |     0.3     |       0.70        |...|       1       |
        +--------------+-------------+-------------------+---+---------------+

        will be represented by a ``cv_results_`` dict of::

            {
            'param_kernel' : masked_array(data = ['rbf', 'rbf', 'rbf'],
                                          mask = False),
            'param_gamma'  : masked_array(data = [0.1 0.2 0.3], mask = False),
            'split0_test_score'  : [0.80, 0.90, 0.70],
            'split1_test_score'  : [0.82, 0.50, 0.70],
            'mean_test_score'    : [0.81, 0.70, 0.70],
            'std_test_score'     : [0.01, 0.20, 0.00],
            'rank_test_score'    : [3, 1, 1],
            'split0_train_score' : [0.80, 0.92, 0.70],
            'split1_train_score' : [0.82, 0.55, 0.70],
            'mean_train_score'   : [0.81, 0.74, 0.70],
            'std_train_score'    : [0.01, 0.19, 0.00],
            'mean_fit_time'      : [0.73, 0.63, 0.43],
            'std_fit_time'       : [0.01, 0.02, 0.01],
            'mean_score_time'    : [0.01, 0.06, 0.04],
            'std_score_time'     : [0.00, 0.00, 0.00],
            'params'             : [{'kernel' : 'rbf', 'gamma' : 0.1}, ...],
            }

        NOTE

        The key ``'params'`` is used to store a list of parameter
        settings dicts for all the parameter candidates.

        The ``mean_fit_time``, ``std_fit_time``, ``mean_score_time`` and
        ``std_score_time`` are all in seconds.

        For multi-metric evaluation, the scores for all the scorers are
        available in the ``cv_results_`` dict at the keys ending with that
        scorer's name (``'_<scorer_name>'``) instead of ``'_score'`` shown
        above. ('split0_test_precision', 'mean_train_precision' etc.)

    best_estimator_ : estimator or dict
        Estimator that was chosen by the search, i.e. estimator
        which gave highest score (or smallest loss if specified)
        on the left out data. Not available if ``refit=False``.

        For multi-metric evaluation, this attribute is present only if
        ``refit`` is specified.

        See ``refit`` parameter for more information on allowed values.

    best_score_ : float
        Mean cross-validated score of the best_estimator.

        For multi-metric evaluation, this is not available if ``refit`` is
        ``False``. See ``refit`` parameter for more information.

    best_params_ : dict
        Parameter setting that gave the best results on the hold out data.

        For multi-metric evaluation, this is not available if ``refit`` is
        ``False``. See ``refit`` parameter for more information.

    best_index_ : int
        The index (of the ``cv_results_`` arrays) which corresponds to the best
        candidate parameter setting.

        The dict at ``search.cv_results_['params'][search.best_index_]`` gives
        the parameter setting for the best model, that gives the highest
        mean score (``search.best_score_``).

        For multi-metric evaluation, this is not available if ``refit`` is
        ``False``. See ``refit`` parameter for more information.

    scorer_ : function or a dict
        Scorer function used on the held out data to choose the best
        parameters for the model.

        For multi-metric evaluation, this attribute holds the validated
        ``scoring`` dict which maps the scorer key to the scorer callable.

    n_splits_ : int
        The number of cross-validation splits (folds/iterations).

    refit_time_ : float
        Seconds used for refitting the best model on the whole dataset.

        This is present only if ``refit`` is not False.

    Notes
    -----
    The parameters selected are those that maximize the score of the held-out
    data, according to the scoring parameter.

    If `n_jobs` was set to a value higher than one, the data is copied for each
    parameter setting(and not `n_jobs` times). This is done for efficiency
    reasons if individual jobs take very little time, but may raise errors if
    the dataset is large and not enough memory is available.  A workaround in
    this case is to set `pre_dispatch`. Then, the memory is copied only
    `pre_dispatch` many times. A reasonable value for `pre_dispatch` is `2 *
    n_jobs`.

    See Also
    --------
    :class:`GridSearchCV`:
        Does exhaustive search over a grid of parameters.

    :class:`ParameterSampler`:
        A generator over parameter settings, constructed from
        param_distributions.

    """
    _required_parameters = ["estimator", "param_distributions"]

    def __init__(self, estimator, param_distributions, n_iter=10, scoring=None,
                 n_jobs=None, iid='warn', refit=True,
                 cv='warn', verbose=0, pre_dispatch='2*n_jobs',
                 random_state=None, error_score='raise-deprecating',
                 return_train_score=False, spark=True):
        self.param_distributions = param_distributions
        self.n_iter = n_iter
        self.random_state = random_state
        super().__init__(
            spark=spark, estimator=estimator, scoring=scoring,
            n_jobs=n_jobs, iid=iid, refit=refit, cv=cv, verbose=verbose,
            pre_dispatch=pre_dispatch, error_score=error_score,
            return_train_score=return_train_score)

    def _run_search(self, evaluate_candidates):
        """Search n_iter candidates from param_distributions"""
        evaluate_candidates(ParameterSampler(
            self.param_distributions, self.n_iter,
            random_state=self.random_state))


# -*- coding: utf-8 -*-
# @Author: Song Dejia
# @Date:   2018-11-09 10:06:59
# @Last Modified by:   Song Dejia
# @Last Modified time: 2018-11-23 17:26:44
import os
import os.path as osp
import random
import time
import sys;

sys.path.append('../')
import torch
import torch.nn as nn
import numpy as np
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
import argparse
from PIL import Image, ImageOps, ImageStat, ImageDraw
from test_data_loader import TestDataLoader
from net import SiameseRPN
from torch.nn import init
from shapely.geometry import Polygon

parser = argparse.ArgumentParser(description='PyTorch SiameseRPN Training')

parser.add_argument('--train_path', default='/home/akaruvally/scratch/vot2013', metavar='DIR', help='path to dataset')

parser.add_argument('--weight_dir', default='/home/akaruvally/scratch/weights', metavar='DIR', help='path to weight')

parser.add_argument('--checkpoint_path', default=None, help='resume')

parser.add_argument('--max_epoches', default=10000, type=int, metavar='N', help='number of total epochs to run')

parser.add_argument('--max_batches', default=0, type=int, metavar='N', help='number of batch in one epoch')

parser.add_argument('--init_type', default='xavier', type=str, metavar='INIT', help='init net')

parser.add_argument('--lr', default=0.001, type=float, metavar='LR', help='initial learning rate')

parser.add_argument('--momentum', default=0.9, type=float, metavar='momentum', help='momentum')

parser.add_argument('--weight_decay', '--wd', default=5e-5, type=float, metavar='W',
                    help='weight decay (default: 1e-4)')

parser.add_argument('--debug', default=False, type=bool, help='whether to debug')

parser.add_argument('--save_frequency', default=10000, type=int, help='what frequency to save models')

parser.add_argument('--temp_dir', default='/home/akaruvally/scratch_dir/tmp/visualization/7_check_train_phase_debug_pos_anchors',
                    type=str, help='Temporary directory to save the trained models')

def main():
    """ train dataloader """
    args = parser.parse_args()
    data_loader = TestDataLoader(args.train_path, check=args.debug, tmp_dir=args.temp_dir)
    if not os.path.exists(args.weight_dir):
        print("Weights directory not found but not required!!")

    """ compute max_batches """
    for root, dirs, files in os.walk(args.train_path):
        for dirname in dirs:
            dir_path = os.path.join(root, dirname)
            args.max_batches += len(os.listdir(dir_path))

    """ Model on gpu """
    model = SiameseRPN()
    model = model.cuda()
    cudnn.benchmark = True

    """ loss and optimizer """
    criterion = MultiBoxLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)

    """ load weights """
    init_weights(model)
    if not args.checkpoint_path == None:
        assert os.path.isfile(args.checkpoint_path), '{} is not valid checkpoint_path'.format(args.checkpoint_path)
        try:
            checkpoint = torch.load(args.checkpoint_path)
            start = checkpoint['epoch']
            model.load_state_dict(checkpoint['state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer'])
        except:
            start = 0
            init_weights(model)
    else:
        start = 0

    """ train phase """
    closses, rlosses, tlosses = AverageMeter(), AverageMeter(), AverageMeter()
    steps = 0

    # cur_lr = adjust_learning_rate(args.lr, optimizer, epoch, gamma=0.1)
    index_list = range(data_loader.__len__())
    for example in range(data_loader.n_frame_list[0]):
        ret = data_loader.__get__(0, example)
        template = ret['template_tensor'].cuda()
        detection = ret['detection_tensor'].cuda()
        pos_neg_diff = ret['pos_neg_diff_tensor'].cuda()
        cout, rout = model(template, detection)
        predictions, targets = (cout, rout), pos_neg_diff
        closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index = criterion(predictions, targets)
        closs_ = closs.cpu().item()

        if np.isnan(closs_):
            sys.exit(0)

        closses.update(closs.cpu().item())
        rlosses.update(rloss.cpu().item())
        tlosses.update(loss.cpu().item())

        # optimizer.zero_grad()
        # loss.backward()
        # optimizer.step()
        steps += 1

        cout = cout.cpu().detach().numpy()
        score = 1 / (1 + np.exp(cout[:, 0] - cout[:, 1]))

        # ++++++++++++ post process below just for debug ++++++++++++++++++++++++
        # ++++++++++++++++++++ v1.0 add penalty +++++++++++++++++++++++++++++++++
        if ret['pos_anchors'] is not None:
            penalty_k = 0.055
            tx, ty, tw, th = ret['template_target_xywh'].copy()
            tw *= ret['template_cropprd_resized_ratio']
            th *= ret['template_cropprd_resized_ratio']

            anchors = ret['anchors'].copy()
            w = anchors[:, 2] * np.exp(reg_pred[:, 2].cpu().detach().numpy())
            h = anchors[:, 3] * np.exp(reg_pred[:, 3].cpu().detach().numpy())

            eps = 1e-2
            change_w = np.maximum(w / (tw + eps), tw / (w + eps))
            change_h = np.maximum(h / (th + eps), th / (h + eps))
            penalty = np.exp(-(change_w + change_h - 1) * penalty_k)
            pscore = score * penalty
        else:
            pscore = score

        # +++++++++++++++++++ v1.0 add window default cosine ++++++++++++++++++++++
        score_size = 17
        window_influence = 0.42
        window = (np.outer(np.hanning(score_size), np.hanning(score_size)).reshape(17, 17, 1) + np.zeros(
            (1, 1, 5))).reshape(-1)
        pscore = pscore * (1 - window_influence) + window * window_influence
        score_old = score
        score = pscore  # from 0.2 - 0.7

        # +++++++++++++++++++ v1.0 add nms ++++++++++++++++++++++++++++++++++++++++++++
        nms = False
        nms_threshold = 0.6
        start = time.time()
        anchors = ret['anchors'].copy()
        x = anchors[:, 0] + anchors[:, 2] * reg_pred[:, 0].cpu().detach().numpy()
        y = anchors[:, 1] + anchors[:, 3] * reg_pred[:, 1].cpu().detach().numpy()
        w = anchors[:, 2] * np.exp(reg_pred[:, 2].cpu().detach().numpy())
        h = anchors[:, 3] * np.exp(reg_pred[:, 3].cpu().detach().numpy())
        x1 = np.clip(x - w // 2, 0, 256)
        x2 = np.clip(x + w // 2, 0, 256)
        x3 = np.clip(x + w // 2, 0, 256)
        x4 = np.clip(x - w // 2, 0, 256)
        y1 = np.clip(y - h // 2, 0, 256)
        y2 = np.clip(y - h // 2, 0, 256)
        y3 = np.clip(y + h // 2, 0, 256)
        y4 = np.clip(y + h // 2, 0, 256)
        slist = map(reshape, [x1, y1, x2, y2, x3, y3, x4, y4, score])
        s = np.hstack(slist)
        maxscore = max(s[:, 8])
        if nms and maxscore > nms_threshold:
            proposals = standard_nms(s, nms_threshold)
            proposals = proposals if proposals.shape[0] != 0 else s
            print('nms spend {:.2f}ms'.format(1000 * (time.time() - start)))
        else:
            proposals = s
        # ++++++++++++++++++++ debug for class ++++++++++++++++++++++++++++++++++++
        # print(score[pos_index])  # this should tend to be 1
        # print(score[neg_index])  # this should tend to be 0

        # ++++++++++++++++++++ debug for reg ++++++++++++++++++++++++++++++++++++++
        tmp_dir = args.temp_dir
        if not os.path.exists(tmp_dir):
            os.makedirs(tmp_dir)
        detection = ret['detection_cropped_resized'].copy()
        draw = ImageDraw.Draw(detection)
        pos_anchors = ret['pos_anchors'].copy() if ret['pos_anchors'] is not None else None

        if pos_anchors is not None:
            # draw pos anchors
            x = pos_anchors[:, 0]
            y = pos_anchors[:, 1]
            w = pos_anchors[:, 2]
            h = pos_anchors[:, 3]
            x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2
            for i in range(16):
                x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i]
                # draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1, fill='white')  # pos anchor

            # # pos anchor transform to red box after prediction
            x = pos_anchors[:, 0] + pos_anchors[:, 2] * reg_pred[pos_index, 0].cpu().detach().numpy()
            y = pos_anchors[:, 1] + pos_anchors[:, 3] * reg_pred[pos_index, 1].cpu().detach().numpy()
            w = pos_anchors[:, 2] * np.exp(reg_pred[pos_index, 2].cpu().detach().numpy())
            h = pos_anchors[:, 3] * np.exp(reg_pred[pos_index, 3].cpu().detach().numpy())
            x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2
            for i in range(16):
                x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i]
                draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1,
                          fill='red')  # predict(white -> red)
                break

            # # pos anchor should be transformed to green gt box, if red and green is same, it is overfitting
            # x = pos_anchors[:, 0] + pos_anchors[:, 2] * reg_target[pos_index, 0].cpu().detach().numpy()
            # y = pos_anchors[:, 1] + pos_anchors[:, 3] * reg_target[pos_index, 1].cpu().detach().numpy()
            # w = pos_anchors[:, 2] * np.exp(reg_target[pos_index, 2].cpu().detach().numpy())
            # h = pos_anchors[:, 3] * np.exp(reg_target[pos_index, 3].cpu().detach().numpy())
            # x1s, y1s, x2s, y2s = x - w // 2, y - h // 2, x + w // 2, y + h // 2
            # for i in range(16):
            #     x1, y1, x2, y2 = x1s[i], y1s[i], x2s[i], y2s[i]
            #     draw.line([(x1, y1), (x2, y1), (x2, y2), (x1, y2), (x1, y1)], width=1,
            #               fill='green')  # gt  (white -> green)
            x1, y1, x3, y3 = x1s[0], y1s[0], x2s[0], y2s[0]
        else:
            x1, y1, x3, y3 = 0, 0, 0, 0
        # top1 proposal after nms (white)
        if nms:
            index = np.argsort(proposals[:, 8])[::-1][0]
            x1, y1, x2, y2, x3, y3, x4, y4, _ = proposals[index]
            draw.line([(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)], width=3, fill='yellow')
        save_path = osp.join(tmp_dir, 'example_{:010d}_anchor_pred.jpg'.format(example))
        detection.save(save_path)

        # +++++++++++++++++++ v1.0 restore ++++++++++++++++++++++++++++++++++++++++
        ratio = ret['detection_cropped_resized_ratio']
        detection_cropped = ret['detection_cropped'].copy()
        detection_cropped_resized = ret['detection_cropped_resized'].copy()
        original = Image.open(ret['detection_img_path'])
        x_, y_ = ret['detection_tlcords_of_original_image']
        draw = ImageDraw.Draw(original)
        w, h = original.size
        """ un resized """
        x1, y1, x3, y3 = x1 / ratio, y1 / ratio, y3 / ratio, y3 / ratio

        """ un cropped """
        x1 = np.clip(x_ + x1, 0, w - 1).astype(np.int32)  # uncropped #target_of_original_img
        y1 = np.clip(y_ + y1, 0, h - 1).astype(np.int32)
        x3 = np.clip(x_ + x3, 0, w - 1).astype(np.int32)
        y3 = np.clip(y_ + y3, 0, h - 1).astype(np.int32)

        draw.line([(x1, y1), (x3, y1), (x3, y3), (x1, y3), (x1, y1)], width=3, fill='yellow')
        save_path = osp.join(tmp_dir, 'example_{:010d}_restore.jpg'.format(example))
        # original.save(save_path)

        print(
            "example:{:06d}/{:06d}({:.2f})%\tsteps:{:010d}\tcloss:{:.4f}\trloss:{:.4f}\ttloss:{:.4f}".format(
                example + 1, data_loader.n_frame_list[0], 100 * (example + 1) / args.max_batches, steps,
                closses.avg, rlosses.avg, tlosses.avg))

def intersection(g, p):
    g = Polygon(g[:8].reshape((4, 2)))
    p = Polygon(p[:8].reshape((4, 2)))
    if not g.is_valid or not p.is_valid:
        return 0
    inter = Polygon(g).intersection(Polygon(p)).area
    union = g.area + p.area - inter
    if union == 0:
        return 0
    else:
        return inter / union

def standard_nms(S, thres):
    """ use pre_thres to filter """
    index = np.where(S[:, 8] > thres)[0]
    S = S[index]  # ~ 100, 4

    # Then use standard nms
    order = np.argsort(S[:, 8])[::-1]
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        ovr = np.array([intersection(S[i], S[t]) for t in order[1:]])

        inds = np.where(ovr <= thres)[0]
        order = order[inds + 1]
    return S[keep]

def reshape(x):
    t = np.array(x, dtype=np.float32)
    return t.reshape(-1, 1)

def init_weights(net, init_type='normal', gain=0.02):
    def init_func(m):
        # this will apply to each layer
        classname = m.__class__.__name__
        if hasattr(m, 'weight') and (classname.find('conv') != -1 or classname.find('Linear') != -1):
            if init_type == 'normal':
                init.normal_(m.weight.data, 0.0, gain)
            elif init_type == 'xavier':
                init.xavier_normal_(m.weight.data, gain=gain)
            elif init_type == 'kaiming':
                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')  # good for relu
            elif init_type == 'orthogonal':
                init.orthogonal_(m.weight.data, gain=gain)
            else:
                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)

            if hasattr(m, 'bias') and m.bias is not None:
                init.constant_(m.bias.data, 0.0)
        elif classname.find('BatchNorm2d') != -1:
            init.normal_(m.weight.data, 1.0, gain)
            init.constant_(m.bias.data, 0.0)

    # print('initialize network with %s' % init_type)
    net.apply(init_func)

class MultiBoxLoss(nn.Module):
    def __init__(self):
        super(MultiBoxLoss, self).__init__()

    def forward(self, predictions, targets):
        print('+++++++++++++++++++++++++++++++++++++++++++++++++++')
        cout, rout = predictions
        """ class """
        class_pred, class_target = cout, targets[:, 0].long()
        pos_index, neg_index = list(np.where(class_target == 1)[0]), list(np.where(class_target == 0)[0])
        pos_num, neg_num = len(pos_index), len(neg_index)
        class_pred, class_target = class_pred[pos_index + neg_index], class_target[pos_index + neg_index]

        closs = F.cross_entropy(class_pred, class_target, size_average=False, reduce=False)
        closs = torch.div(torch.sum(closs), 64)

        """ regression """
        reg_pred = rout
        reg_target = targets[:, 1:]
        rloss = F.smooth_l1_loss(reg_pred, reg_target, size_average=False, reduce=False)  # 1445, 4
        rloss = torch.div(torch.sum(rloss, dim=1), 4)
        rloss = torch.div(torch.sum(rloss[pos_index]), 16)

        loss = closs + rloss
        return closs, rloss, loss, reg_pred, reg_target, pos_index, neg_index

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def adjust_learning_rate(lr, optimizer, epoch, gamma=0.1):
    """Sets the learning rate to the initial LR decayed 0.9 every 50 epochs"""
    lr = lr * (0.9 ** (epoch // 1))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

if __name__ == '__main__':
    main()


import os
import sys
import numpy as np
import matplotlib as plt
from keras.preprocessing import sequence
from keras.preprocessing.text import Tokenizer
from keras.preprocessing.sequence import pad_sequences
from keras.layers import Dense, Embedding, LSTM, SpatialDropout1D
from keras.models import Sequential
import keras
import tensorflow as tf
import pandas as pd
import tensorflow_datasets as tfds

#path to whatsapp chat
filepath = 'domE.txt'

with open(filepath) as fp:
   lines = fp.readlines()

#remove unlabeled lines (a single message can often get split into multiple lines)
#for whatever reason, it doesn't catch all offending lines on one attempt, so it will run 20 times.
for i in range(20):
    for id, line in enumerate(lines):
        print(str(id))
        if (len(line)>0):
            if (line[0] != "8" and line[0] != "7" and line[0] != "6" and line[0] != "5" and line[0] != "4"):
                lines.remove(line)
                print(line)
        if (len(line)==0):
            lines.remove(line)

#strip lines of whitespace
for id, line in enumerate(lines):
    if (line[-1] == '\n'):
        lines[id] = line.rstrip()

#gather all messages into a list, without names
messages = []
for cnt, line in enumerate(lines):
    #print(line)
    if (len(line.split(": "))>1):
        txt = line.split(": ")[1]
        messages.append(txt)

#gather, for each message, the speaker
speakers = []
for cnt, line in enumerate(lines):
    #print(line)
    if (len(line.split(": "))>1):
        print(str(cnt))
        dateperson = line.split(": ")[0]
        person = dateperson.split(" - ")[1]
        speakers.append(person)

max_features = 100000
tokenizer = Tokenizer(num_words=max_features, split=' ')
tokenizer.fit_on_texts(messages)
X = tokenizer.texts_to_sequences(messages)
X = pad_sequences(X)

#get a list of every unique speaker
speakerset = set(speakers)
uspeakers = []
for i in speakerset:
    uspeakers.append(i)

#convert the categories into numbers (index in uspeakers)
Y = np.zeros(shape=(len(speakers),))
for cnt, person in enumerate(speakers):
    Y[cnt] = uspeakers.index(person)

#build the model
embed_dim = 256
lstm_out = 64
model = Sequential()
model.add(Embedding(max_features, embed_dim,input_length = X.shape[1]))
model.add(SpatialDropout1D(0.2))
model.add(LSTM(lstm_out, dropout=0.2, recurrent_dropout=0.2))
model.add(Dense(len(uspeakers),activation='softmax'))
model.compile(optimizer='sgd', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
print(model.summary())

# Split dataset
train_size = int(0.6 * len(X))
train_X = X[:train_size]
test_X = X[train_size:]
train_Y = Y[:train_size]
test_Y = Y[train_size:]

#train the model
batch_size = 16
model.fit(train_X, train_Y, epochs = 7, batch_size=batch_size, verbose = 2)

#validation
validation_size = int(0.3 * len(X))
X_validate = test_X[-validation_size:]
Y_validate = test_Y[-validation_size:]
X_test = X_test[:-validation_size]
Y_test = Y_test[:-validation_size]
score,acc = model.evaluate(test_X, test_Y, verbose = 2, batch_size = batch_size)
print("score: %.2f" % (score))
print("acc: %.2f" % (acc))

#this function runs predictions on new material
def guess(txt):
    #vectorizing the message by the pre-fitted tokenizer instance
    txt = [txt]
    txt = tokenizer.texts_to_sequences(txt)
    #padding the tweet to have exactly the same shape as `embedding` input
    txt = pad_sequences(txt, maxlen=X.shape[1], dtype='int32', value=0)
    result = model.predict(txt,batch_size=1,verbose = 2)[0]
    #print("Confidence: " + str(sentiment[np.argmax(sentiment)]*100) + "%")
    print(uspeakers[np.argmax(result)] + "\nConfidence: "+str(result[np.argmax(result)]*100) + "%")

guess("this is a message")