from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy as np
import argparse
import time
import torch
import torch.nn as nn

from torch.autograd import Variable
from torch.utils.data import DataLoader
from dataset.factory import get_imdb
from dataset.roidb import RoiDataset, detection_collate
from yolov2 import Yolov2
from torch import optim
from util.network import adjust_learning_rate
from tensorboardX import SummaryWriter
from config import config as cfg

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Yolo v2')
    parser.add_argument('--max_epochs', dest='max_epochs',
                        help='number of epochs to train',
                        default=160, type=int)
    parser.add_argument('--start_epoch', dest='start_epoch',
                        default=1, type=int)
    parser.add_argument('--dataset', dest='dataset',
                        default='voc0712trainval', type=str)
    parser.add_argument('--nw', dest='num_workers',
                        help='number of workers to load training data',
                        default=8, type=int)
    parser.add_argument('--output_dir', dest='output_dir',
                        default='output', type=str)
    parser.add_argument('--use_tfboard', dest='use_tfboard',
                        default=False, type=bool)
    parser.add_argument('--display_interval', dest='display_interval',
                        default=10, type=int)
    parser.add_argument('--mGPUs', dest='mGPUs',
                        default=False, type=bool)
    parser.add_argument('--save_interval', dest='save_interval',
                        default=20, type=int)
    parser.add_argument('--cuda', dest='use_cuda',
                        default=False, type=bool)
    parser.add_argument('--resume', dest='resume',
                        default=False, type=bool)
    parser.add_argument('--checkpoint_epoch', dest='checkpoint_epoch',
                        default=160, type=int)
    parser.add_argument('--exp_name', dest='exp_name',
                        default='default', type=str)

    args = parser.parse_args()
    return args

def get_dataset(datasetnames):
    names = datasetnames.split('+')
    dataset = RoiDataset(get_imdb(names[0]))
    print('load dataset {}'.format(names[0]))
    for name in names[1:]:
        tmp = RoiDataset(get_imdb(name))
        dataset += tmp
        print('load and add dataset {}'.format(name))
    return dataset

def train():

    # define the hyper parameters first
    args = parse_args()
    args.lr = cfg.lr
    args.decay_lrs = cfg.decay_lrs
    args.weight_decay = cfg.weight_decay
    args.momentum = cfg.momentum
    args.batch_size = cfg.batch_size
    args.pretrained_model = os.path.join('data', 'pretrained', 'darknet19_448.weights')

    print('Called with args:')
    print(args)

    lr = args.lr

    # initial tensorboardX writer
    if args.use_tfboard:
        if args.exp_name == 'default':
            writer = SummaryWriter()
        else:
            writer = SummaryWriter('runs/' + args.exp_name)

    if args.dataset == 'voc07train':
        args.imdb_name = 'voc_2007_train'
        args.imdbval_name = 'voc_2007_train'

    elif args.dataset == 'voc07trainval':
        args.imdb_name = 'voc_2007_trainval'
        args.imdbval_name = 'voc_2007_trainval'

    elif args.dataset == 'voc0712trainval':
        args.imdb_name = 'voc_2007_trainval+voc_2012_trainval'
        args.imdbval_name ='voc_2007_test'
    else:
        raise NotImplementedError

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # load dataset
    print('loading dataset....')
    train_dataset = get_dataset(args.imdb_name)

    print('dataset loaded.')

    print('training rois number: {}'.format(len(train_dataset)))
    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size,
                                  shuffle=True, num_workers=args.num_workers,
                                  collate_fn=detection_collate, drop_last=True)

    # initialize the model
    print('initialize the model')
    tic = time.time()
    model = Yolov2(weights_file=args.pretrained_model)
    toc = time.time()
    print('model loaded: cost time {:.2f}s'.format(toc-tic))

    # initialize the optimizer
    optimizer = optim.SGD(model.parameters(), lr=args.lr,
                          momentum=args.momentum, weight_decay=args.weight_decay)

    if args.resume:
        print('resume training enable')
        resume_checkpoint_name = 'yolov2_epoch_{}.pth'.format(args.checkpoint_epoch)
        resume_checkpoint_path = os.path.join(output_dir, resume_checkpoint_name)
        print('resume from {}'.format(resume_checkpoint_path))
        checkpoint = torch.load(resume_checkpoint_path)
        model.load_state_dict(checkpoint['model'])
        args.start_epoch = checkpoint['epoch'] + 1
        lr = checkpoint['lr']
        print('learning rate is {}'.format(lr))
        adjust_learning_rate(optimizer, lr)

    if args.use_cuda:
        model.cuda()

    if args.mGPUs:
        model = nn.DataParallel(model)

    # set the model mode to train because we have some layer whose behaviors are different when in training and testing.
    # such as Batch Normalization Layer.
    model.train()

    iters_per_epoch = int(len(train_dataset) / args.batch_size)

    # start training
    for epoch in range(args.start_epoch, args.max_epochs+1):
        loss_temp = 0
        tic = time.time()
        train_data_iter = iter(train_dataloader)

        if epoch in args.decay_lrs:
            lr = args.decay_lrs[epoch]
            adjust_learning_rate(optimizer, lr)
            print('adjust learning rate to {}'.format(lr))

        if cfg.multi_scale and epoch in cfg.epoch_scale:
            cfg.scale_range = cfg.epoch_scale[epoch]
            print('change scale range to {}'.format(cfg.scale_range))

        for step in range(iters_per_epoch):

            if cfg.multi_scale and (step + 1) % cfg.scale_step == 0:
                scale_index = np.random.randint(*cfg.scale_range)
                cfg.input_size = cfg.input_sizes[scale_index]
                print('change input size {}'.format(cfg.input_size))

            im_data, boxes, gt_classes, num_obj = next(train_data_iter)
            if args.use_cuda:
                im_data = im_data.cuda()
                boxes = boxes.cuda()
                gt_classes = gt_classes.cuda()
                num_obj = num_obj.cuda()

            im_data_variable = Variable(im_data)

            box_loss, iou_loss, class_loss = model(im_data_variable, boxes, gt_classes, num_obj, training=True)

            loss = box_loss.mean()+ iou_loss.mean() \
                   + class_loss.mean()

            optimizer.zero_grad()

            loss.backward()
            optimizer.step()

            loss_temp += loss.item()

            if (step + 1) % args.display_interval == 0:
                toc = time.time()
                loss_temp /= args.display_interval

                iou_loss_v = iou_loss.mean().item()
                box_loss_v = box_loss.mean().item()
                class_loss_v = class_loss.mean().item()

                print("[epoch %2d][step %4d/%4d] loss: %.4f, lr: %.2e, time cost %.1fs" \
                      % (epoch, step+1, iters_per_epoch, loss_temp, lr, toc - tic))
                print("\t\t\tiou_loss: %.4f, box_loss: %.4f, cls_loss: %.4f" \
                      % (iou_loss_v, box_loss_v, class_loss_v))

                if args.use_tfboard:

                    n_iter = (epoch - 1) * iters_per_epoch + step + 1
                    writer.add_scalar('losses/loss', loss_temp, n_iter)
                    writer.add_scalar('losses/iou_loss', iou_loss_v, n_iter)
                    writer.add_scalar('losses/box_loss', box_loss_v, n_iter)
                    writer.add_scalar('losses/cls_loss', class_loss_v, n_iter)

                loss_temp = 0
                tic = time.time()

        if epoch % args.save_interval == 0:
            save_name = os.path.join(output_dir, 'yolov2_epoch_{}.pth'.format(epoch))
            torch.save({
                'model': model.module.state_dict() if args.mGPUs else model.state_dict(),
                'epoch': epoch,
                'lr': lr
                }, save_name)

if __name__ == '__main__':
    train()


import cv2
import numpy as np
cap = cv2.VideoCapture(0)

Wht = 320
confThreshold = 0.5
nmsThreshold = 0.3
classesFile = 'coco names.txt'
classNames = []

with open(classesFile,'rt') as f:
    classNames = f.read().rstrip('\n').split('\n')

modelConfiguration = 'yolov3.cfg'
modelWeights = 'yolov3.weights'

net = cv2.dnn.readNetFromDarknet(modelConfiguration,modelWeights)
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)

def findObjects(outputs,img):
    ht,wt,ct = img.shape
    bbox = []
    classIds = []
    confs = []

    for output in outputs:
        for det in output:
            scores = det[5:]
            classId = np.argmax(scores)
            confidence = scores[classId]

            if confidence>confThreshold:
                w,h = int(det[2]*wt),int(det[3]*ht)
                x,y = int((det[0]*wt)-w/2),int((det[1]*ht)-h/2)
                bbox.append([x,y,w,h])
                classIds.append(classId)
                confs.append(float(confidence))

    indices = cv2.dnn.NMSBoxes(bbox,confs,confThreshold,nmsThreshold)
    print(indices)
    for i in indices:
        i = i[0]
        box = bbox[i]
        x,y,w,h = box[0],box[1],box[2],box[3]
        cv2.rectangle(img,(x,y),(x+w,y+h),(255,0,255),2)
        cv2.putText(img,f"{classNames[classIds[i]].upper()}{int(confs[i]*100)}%",(x,y-10),cv2.FONT_HERSHEY_SIMPLEX,0.6,(255,0,255),2)

while True:
    success,img = cap.read()
    blob = cv2.dnn.blobFromImage(img,1/255,(Wht,Wht),[0,0,0],1,crop=False)
    net.setInput(blob)
    layerNames = net.getLayerNames()
    # print(layerNames)

    outputNames = [layerNames[i[0]-1] for i in net.getUnconnectedOutLayers()]
    # print(outputNames)
    # print(net.getUnconnectedOutLayers)

    outputs = net.forward(outputNames)

    findObjects(outputs,img)
    cv2.imshow('Image',img)
    if cv2.waitKey(1)==13:
        break
cap.release()
cv2.destroyAllWindows()


from utils import *
import numpy as np
import time
import math
import argparse
from keras.initializations import normal, identity, uniform
from keras.models import model_from_json
from keras.models import Sequential, Model
from keras.layers import Dense, BatchNormalization, Activation, Convolution2D, MaxPooling2D, Flatten, Input, merge, Lambda
from keras.layers.advanced_activations import LeakyReLU
from keras.optimizers import Adam, RMSprop
import tensorflow as tf
import keras.backend as K
import json

from keras.applications.resnet50 import ResNet50
from keras.preprocessing import image
from keras.applications.resnet50 import preprocess_input
from keras.utils.np_utils import to_categorical

parser = argparse.ArgumentParser(description="TRPO")
parser.add_argument("--paths_per_collect", type=int, default=10)
parser.add_argument("--max_step_limit", type=int, default=300)
parser.add_argument("--min_step_limit", type=int, default=100)
parser.add_argument("--pre_step", type=int, default=100)
parser.add_argument("--n_iter", type=int, default=1000)
parser.add_argument("--gamma", type=float, default=.95)
parser.add_argument("--lam", type=float, default=.97)
parser.add_argument("--max_kl", type=float, default=0.01)
parser.add_argument("--cg_damping", type=float, default=0.1)
parser.add_argument("--lr_discriminator", type=float, default=5e-5)
parser.add_argument("--d_iter", type=int, default=100)
parser.add_argument("--clamp_lower", type=float, default=-0.01)
parser.add_argument("--clamp_upper", type=float, default=0.01)
parser.add_argument("--lr_baseline", type=float, default=1e-4)
parser.add_argument("--b_iter", type=int, default=25)
parser.add_argument("--lr_posterior", type=float, default=1e-4)
parser.add_argument("--p_iter", type=int, default=50)
parser.add_argument("--buffer_size", type=int, default=75)
parser.add_argument("--sample_size", type=int, default=50)
parser.add_argument("--batch_size", type=int, default=500)

args = parser.parse_args()

class TRPOAgent(object):
    config = dict2(paths_per_collect = args.paths_per_collect,
                   max_step_limit = args.max_step_limit,
                   min_step_limit = args.min_step_limit,
                   pre_step = args.pre_step,
                   n_iter = args.n_iter,
                   gamma = args.gamma,
                   lam = args.lam,
                   max_kl = args.max_kl,
                   cg_damping = args.cg_damping,
                   lr_discriminator = args.lr_discriminator,
                   d_iter = args.d_iter,
                   clamp_lower = args.clamp_lower,
                   clamp_upper = args.clamp_upper,
                   lr_baseline = args.lr_baseline,
                   b_iter = args.b_iter,
                   lr_posterior = args.lr_posterior,
                   p_iter = args.p_iter,
                   buffer_size = args.buffer_size,
                   sample_size = args.sample_size,
                   batch_size = args.batch_size)

    def __init__(self, env, sess, feat_dim, aux_dim, encode_dim, action_dim,
                 img_dim, pre_actions):
        self.env = env
        self.sess = sess
        self.buffer = ReplayBuffer(self.config.buffer_size)
        self.feat_dim = feat_dim
        self.aux_dim = aux_dim
        self.encode_dim = encode_dim
        self.action_dim = action_dim
        self.img_dim = img_dim
        self.pre_actions = pre_actions
        self.feats = feats = tf.placeholder(
            dtype, shape=[None, feat_dim[0], feat_dim[1], feat_dim[2]]
        )
        self.auxs = auxs = tf.placeholder(dtype, shape=[None, aux_dim])
        self.encodes = encodes = tf.placeholder(dtype, shape=[None, encode_dim])
        self.actions = actions = tf.placeholder(dtype, shape=[None, action_dim])

        self.advants = advants = tf.placeholder(dtype, shape=[None])
        self.oldaction_dist_mu = oldaction_dist_mu = \
                tf.placeholder(dtype, shape=[None, action_dim])
        self.oldaction_dist_logstd = oldaction_dist_logstd = \
                tf.placeholder(dtype, shape=[None, action_dim])

        # Create neural network.
        print "Now we build trpo generator"
        self.generator = self.create_generator(feats, auxs, encodes)
        print "Now we build discriminator"
        self.discriminator, self.discriminate = \
            self.create_discriminator(img_dim, aux_dim, action_dim)
        print "Now we build posterior"
        self.posterior = \
            self.create_posterior(img_dim, aux_dim, action_dim, encode_dim)
        self.posterior_target = \
            self.create_posterior(img_dim, aux_dim, action_dim, encode_dim)

        self.demo_idx = 0

        action_dist_mu = self.generator.outputs[0]
        # self.action_dist_logstd_param = action_dist_logstd_param = \
        #         tf.placeholder(dtype, shape=[1, action_dim])
        # action_dist_logstd = tf.tile(action_dist_logstd_param,
        #                              tf.pack((tf.shape(action_dist_mu)[0], 1)))
        action_dist_logstd = tf.placeholder(dtype, shape=[None, action_dim])

        eps = 1e-8
        self.action_dist_mu = action_dist_mu
        self.action_dist_logstd = action_dist_logstd
        N = tf.shape(feats)[0]
        # compute probabilities of current actions and old actions
        log_p_n = gauss_log_prob(action_dist_mu, action_dist_logstd, actions)
        log_oldp_n = gauss_log_prob(oldaction_dist_mu, oldaction_dist_logstd, actions)

        ratio_n = tf.exp(log_p_n - log_oldp_n)
        Nf = tf.cast(N, dtype)
        surr = -tf.reduce_mean(ratio_n * advants) # Surrogate loss
        var_list = self.generator.trainable_weights

        kl = gauss_KL(oldaction_dist_mu, oldaction_dist_logstd,
                      action_dist_mu, action_dist_logstd) / Nf
        ent = gauss_ent(action_dist_mu, action_dist_logstd) / Nf

        self.losses = [surr, kl, ent]
        self.pg = flatgrad(surr, var_list)
        # KL divergence where first arg is fixed
        kl_firstfixed = gauss_selfKL_firstfixed(action_dist_mu,
                                                action_dist_logstd) / Nf
        grads = tf.gradients(kl_firstfixed, var_list)
        self.flat_tangent = tf.placeholder(dtype, shape=[None])
        shapes = map(var_shape, var_list)
        start = 0
        tangents = []
        for shape in shapes:
            size = np.prod(shape)
            param = tf.reshape(self.flat_tangent[start:(start + size)], shape)
            tangents.append(param)
            start += size
        gvp = [tf.reduce_sum(g * t) for (g, t) in zip(grads, tangents)]
        self.fvp = flatgrad(gvp, var_list)
        self.gf = GetFlat(self.sess, var_list)
        self.sff = SetFromFlat(self.sess, var_list)
        self.baseline = NNBaseline(sess, feat_dim, aux_dim, encode_dim,
                                   self.config.lr_baseline, self.config.b_iter,
                                   self.config.batch_size)
        self.sess.run(tf.global_variables_initializer())

        # Create feature extractor
        self.base_model = ResNet50(weights='imagenet', include_top=False)
        self.feat_extractor = Model(
            input=self.base_model.input,
            output=self.base_model.get_layer('activation_40').output
        )

    def create_generator(self, feats, auxs, encodes):
        feats = Input(tensor=feats)
        x = Convolution2D(256, 3, 3)(feats)
        x = LeakyReLU()(x)
        x = Convolution2D(256, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Flatten()(x)
        auxs = Input(tensor=auxs)
        h = merge([x, auxs], mode='concat')
        h = Dense(256)(h)
        h = LeakyReLU()(h)
        h = Dense(128)(h)
        encodes = Input(tensor=encodes)
        c = Dense(128)(encodes)
        h = merge([h, c], mode='sum')
        h = LeakyReLU()(h)

        steer = Dense(1, activation='tanh', init=lambda shape, name:
                      normal(shape, scale=1e-4, name=name))(h)
        accel = Dense(1, activation='sigmoid', init=lambda shape, name:
                             normal(shape, scale=1e-4, name=name))(h)
        brake = Dense(1, activation='sigmoid', init=lambda shape, name:
                      normal(shape, scale=1e-4, name=name))(h)
        actions = merge([steer, accel, brake], mode='concat')
        model = Model(input=[feats, auxs, encodes], output=actions)
        return model

    def create_discriminator(self, img_dim, aux_dim, action_dim):
        imgs = Input(shape=[img_dim[0], img_dim[1], img_dim[2]])
        x = Convolution2D(32, 3, 3, subsample=(2, 2))(imgs)
        x = LeakyReLU()(x)
        x = Convolution2D(64, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Convolution2D(128, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Flatten()(x)
        auxs = Input(shape=[aux_dim])
        actions = Input(shape=[action_dim])
        h = merge([x, auxs, actions], mode='concat')
        h = Dense(256)(h)
        h = LeakyReLU()(h)
        h = Dense(128)(h)
        h = LeakyReLU()(h)
        p = Dense(1)(h)
        discriminate = Model(input=[imgs, auxs, actions], output=p)

        imgs_n = Input(shape=[img_dim[0], img_dim[1], img_dim[2]])
        imgs_d = Input(shape=[img_dim[0], img_dim[1], img_dim[2]])
        auxs_n = Input(shape=[aux_dim])
        auxs_d = Input(shape=[aux_dim])
        actions_n = Input(shape=[action_dim])
        actions_d = Input(shape=[action_dim])

        p_n = discriminate([imgs_n, auxs_n, actions_n])
        p_d = discriminate([imgs_d, auxs_d, actions_d])
        p_d = Lambda(lambda x: -x)(p_d)
        p_output = merge([p_n, p_d], mode='sum')

        model = Model(input=[imgs_n, auxs_n, actions_n,
                             imgs_d, auxs_d, actions_d],
                      output=p_output)
        rmsprop = RMSprop(lr=self.config.lr_discriminator)
        model.compile(
            # little trick to use Keras predefined lambda loss function
            loss=lambda y_pred, p_true: K.mean(y_pred * p_true), optimizer=rmsprop
        )

        return model, discriminate

    def create_posterior(self, img_dim, aux_dim, action_dim, encode_dim):
        imgs = Input(shape=[img_dim[0], img_dim[1], img_dim[2]])
        x = Convolution2D(32, 3, 3, subsample=(2, 2))(imgs)
        x = LeakyReLU()(x)
        x = Convolution2D(64, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Convolution2D(128, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Flatten()(x)
        auxs = Input(shape=[aux_dim])
        actions = Input(shape=[action_dim])
        h = merge([x, auxs, actions], mode='concat')
        h = Dense(256)(h)
        h = LeakyReLU()(h)
        h = Dense(128)(h)
        h = LeakyReLU()(h)
        c = Dense(encode_dim, activation='softmax')(h)

        model = Model(input=[imgs, auxs, actions], output=c)
        adam = Adam(lr=self.config.lr_posterior)
        model.compile(loss='categorical_crossentropy', optimizer=adam,
                      metrics=['accuracy'])
        return model

    def act(self, feats, auxs, encodes, logstds, *args):
        action_dist_mu = \
                self.sess.run(
                    self.action_dist_mu,
                    {self.feats: feats, self.auxs: auxs, self.encodes: encodes}
                )

        act = action_dist_mu + np.exp(logstds) * \
                np.random.randn(*logstds.shape)

        act[:, 0] = np.clip(act[:, 0], -1, 1)
        act[:, 1] = np.clip(act[:, 1], 0, 1)
        act[:, 2] = np.clip(act[:, 2], 0, 1)

        return act

    def learn(self, demo):
        config = self.config
        start_time = time.time()
        numeptotal = 0

        # Set up for training discrimiator
        print "Loading data ..."
        imgs_d, auxs_d, actions_d = demo["imgs"], demo["auxs"], demo["actions"]
        numdetotal = imgs_d.shape[0]
        idx_d = np.arange(numdetotal)
        np.random.shuffle(idx_d)

        imgs_d = imgs_d[idx_d]
        auxs_d = auxs_d[idx_d]
        actions_d = actions_d[idx_d]
        print "Resizing img for demo ..."
        imgs_reshaped_d = []
        for i in xrange(numdetotal):
            imgs_reshaped_d.append(np.expand_dims(cv2.resize(imgs_d[i],
                (self.img_dim[0], self.img_dim[1])), axis=0))
        imgs_d = np.concatenate(imgs_reshaped_d, axis=0).astype(np.float32)
        imgs_d = (imgs_d - 128.) / 128.
        print "Shape of resized demo images:", imgs_d.shape

        for i in xrange(38, config.n_iter):

            # Generating paths.
            # if i == 1:
            if i == 38:
                paths_per_collect = 30
            else:
                paths_per_collect = 10
            rollouts = rollout_contin(
                self.env,
                self,
                self.feat_extractor,
                self.feat_dim,
                self.aux_dim,
                self.encode_dim,
                config.max_step_limit,
                config.min_step_limit,
                config.pre_step,
                paths_per_collect,
                self.pre_actions,
                self.discriminate,
                self.posterior_target)

            for path in rollouts:
                self.buffer.add(path)
            print "Buffer count:", self.buffer.count()

            paths = self.buffer.get_sample(config.sample_size)

            print "Calculating actions ..."
            for path in paths:
                path["mus"] = self.sess.run(
                    self.action_dist_mu,
                    {self.feats: path["feats"],
                     self.auxs: path["auxs"],
                     self.encodes: path["encodes"]}
                )

            mus_n = np.concatenate([path["mus"] for path in paths])
            logstds_n = np.concatenate([path["logstds"] for path in paths])
            feats_n = np.concatenate([path["feats"] for path in paths])
            auxs_n = np.concatenate([path["auxs"] for path in paths])
            encodes_n = np.concatenate([path["encodes"] for path in paths])
            actions_n = np.concatenate([path["actions"] for path in paths])
            imgs_n = np.concatenate([path["imgs"] for path in paths])

            print "Epoch:", i, "Total sampled data points:", feats_n.shape[0]

            # Train discriminator
            numnototal = feats_n.shape[0]
            batch_size = config.batch_size
            start_d = self.demo_idx
            start_n = 0
            if i <= 5:
                d_iter = 120 - i * 20
            else:
                d_iter = 10
            for k in xrange(d_iter):
                loss = self.discriminator.train_on_batch(
                    [imgs_n[start_n:start_n + batch_size],
                     auxs_n[start_n:start_n + batch_size],
                     actions_n[start_n:start_n + batch_size],
                     imgs_d[start_d:start_d + batch_size],
                     auxs_d[start_d:start_d + batch_size],
                     actions_d[start_d:start_d + batch_size]],
                    np.ones(batch_size)
                )
                # print self.discriminator.summary()
                for l in self.discriminator.layers:
                    weights = l.get_weights()
                    weights = [np.clip(w, config.clamp_lower, config.clamp_upper)
                               for w in weights]
                    l.set_weights(weights)

                start_d = self.demo_idx = self.demo_idx + batch_size
                start_n = start_n + batch_size

                if start_d + batch_size >= numdetotal:
                    start_d = self.demo_idx = (start_d + batch_size) % numdetotal
                if start_n + batch_size >= numnototal:
                    start_n = (start_n + batch_size) % numnototal

                print "Discriminator step:", k, "loss:", loss

            idx = np.arange(numnototal)
            np.random.shuffle(idx)
            train_val_ratio = 0.7
            # Training data for posterior
            numno_train = int(numnototal * train_val_ratio)
            imgs_train = imgs_n[idx][:numno_train]
            auxs_train = auxs_n[idx][:numno_train]
            actions_train = actions_n[idx][:numno_train]
            encodes_train = encodes_n[idx][:numno_train]
            # Validation data for posterior
            imgs_val = imgs_n[idx][numno_train:]
            auxs_val = auxs_n[idx][numno_train:]
            actions_val = actions_n[idx][numno_train:]
            encodes_val = encodes_n[idx][numno_train:]

            start_n = 0
            for j in xrange(config.p_iter):
                loss = self.posterior.train_on_batch(
                    [imgs_train[start_n:start_n + batch_size],
                     auxs_train[start_n:start_n + batch_size],
                     actions_train[start_n:start_n + batch_size]],
                    encodes_train[start_n:start_n + batch_size]
                )
                start_n += batch_size
                if start_n + batch_size >= numno_train:
                    start_n = (start_n + batch_size) % numno_train

                posterior_weights = self.posterior.get_weights()
                posterior_target_weights = self.posterior_target.get_weights()
                for k in xrange(len(posterior_weights)):
                    posterior_target_weights[k] = 0.5 * posterior_weights[k] +\
                            0.5 * posterior_target_weights[k]
                self.posterior_target.set_weights(posterior_target_weights)

                output_p = self.posterior_target.predict(
                    [imgs_val, auxs_val, actions_val])
                val_loss = -np.average(
                    np.sum(np.log(output_p) * encodes_val, axis=1))
                print "Posterior step:", j, "loss:", loss, val_loss

            # Computing returns and estimating advantage function.
            path_idx = 0
            for path in paths:
                file_path = "/home/yunzhu/Desktop/log/iter_%d_path_%d.txt" % (i, path_idx)
                f = open(file_path, "w")
                path["baselines"] = self.baseline.predict(path)
                output_d = self.discriminate.predict(
                    [path["imgs"], path["auxs"], path["actions"]])
                output_p = self.posterior_target.predict(
                    [path["imgs"], path["auxs"], path["actions"]])
                path["rewards"] = np.ones(path["raws"].shape[0]) * 2 + \
                        output_d.flatten() * 0.1 + \
                        np.sum(np.log(output_p) * path["encodes"], axis=1)

                path_baselines = np.append(path["baselines"], 0 if
                                           path["baselines"].shape[0] == 100 else
                                           path["baselines"][-1])
                deltas = path["rewards"] + config.gamma * path_baselines[1:] -\
                        path_baselines[:-1]
                # path["returns"] = discount(path["rewards"], config.gamma)
                # path["advants"] = path["returns"] - path["baselines"]
                path["advants"] = discount(deltas, config.gamma * config.lam)
                path["returns"] = discount(path["rewards"], config.gamma)

                f.write("Baseline:\n" + np.array_str(path_baselines) + "\n")
                f.write("Returns:\n" + np.array_str(path["returns"]) + "\n")
                f.write("Advants:\n" + np.array_str(path["advants"]) + "\n")
                f.write("Mus:\n" + np.array_str(path["mus"]) + "\n")
                f.write("Actions:\n" + np.array_str(path["actions"]) + "\n")
                f.write("Logstds:\n" + np.array_str(path["logstds"]) + "\n")
                path_idx += 1

            # Standardize the advantage function to have mean=0 and std=1
            advants_n = np.concatenate([path["advants"] for path in paths])
            # advants_n -= advants_n.mean()
            advants_n /= (advants_n.std() + 1e-8)

            # Computing baseline function for next iter.
            self.baseline.fit(paths)

            feed = {self.feats: feats_n,
                    self.auxs: auxs_n,
                    self.encodes: encodes_n,
                    self.actions: actions_n,
                    self.advants: advants_n,
                    self.action_dist_logstd: logstds_n,
                    self.oldaction_dist_mu: mus_n,
                    self.oldaction_dist_logstd: logstds_n}

            thprev = self.gf()

            def fisher_vector_product(p):
                feed[self.flat_tangent] = p
                return self.sess.run(self.fvp, feed) + p * config.cg_damping

            g = self.sess.run(self.pg, feed_dict=feed)
            stepdir = conjugate_gradient(fisher_vector_product, -g)
            shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
            assert shs > 0

            lm = np.sqrt(shs / config.max_kl)
            fullstep = stepdir / lm
            neggdotstepdir = -g.dot(stepdir)

            def loss(th):
                self.sff(th)
                return self.sess.run(self.losses[0], feed_dict=feed)
            theta = linesearch(loss, thprev, fullstep, neggdotstepdir / lm)
            self.sff(theta)

            surrafter, kloldnew, entropy = self.sess.run(
                self.losses, feed_dict=feed
            )

            episoderewards = np.array([path["rewards"].sum() for path in paths])
            stats = {}
            numeptotal += len(episoderewards)
            stats["Total number of episodes"] = numeptotal
            stats["Average sum of rewards per episode"] = episoderewards.mean()
            stats["Entropy"] = entropy
            stats["Time elapsed"] = "%.2f mins" % ((time.time() - start_time) / 60.0)
            stats["KL between old and new distribution"] = kloldnew
            stats["Surrogate loss"] = surrafter
            print("\n********** Iteration {} **********".format(i))
            for k, v in stats.iteritems():
                print(k + ": " + " " * (40 - len(k)) + str(v))
            if entropy != entropy:
                exit(-1)

            param_dir = "/home/yunzhu/Desktop/params/"
            print("Now we save model")
            self.generator.save_weights(
                param_dir + "generator_model_%d.h5" % i, overwrite=True)
            with open(param_dir + "generator_model_%d.json" % i, "w") as outfile:
                json.dump(self.generator.to_json(), outfile)

            self.discriminator.save_weights(
                param_dir + "discriminator_model_%d.h5" % i, overwrite=True)
            with open(param_dir + "discriminator_model_%d.json" % i, "w") as outfile:
                json.dump(self.discriminator.to_json(), outfile)

            self.baseline.model.save_weights(
                param_dir + "baseline_model_%d.h5" % i, overwrite=True)
            with open(param_dir + "baseline_model_%d.json" % i, "w") as outfile:
                json.dump(self.baseline.model.to_json(), outfile)

            self.posterior.save_weights(
                param_dir + "posterior_model_%d.h5" % i, overwrite=True)
            with open(param_dir + "posterior_model_%d.json" % i, "w") as outfile:
                json.dump(self.posterior.to_json(), outfile)

            self.posterior_target.save_weights(
                param_dir + "posterior_target_model_%d.h5" % i, overwrite=True)
            with open(param_dir + "posterior_target_model_%d.json" % i, "w") as outfile:
                json.dump(self.posterior_target.to_json(), outfile)

class Generator(object):
    def __init__(self, sess, feat_dim, aux_dim, encode_dim, action_dim):
        self.sess = sess
        self.lr = tf.placeholder(tf.float32, shape=[])

        K.set_session(sess)

        self.model, self.weights, self.feats, self.auxs, self.encodes = \
                self.create_generator(feat_dim, aux_dim, encode_dim)

        self.action_gradient = tf.placeholder(tf.float32, [None, action_dim])
        self.params_grad = tf.gradients(self.model.output, self.weights,
                                        self.action_gradient)
        grads = zip(self.params_grad, self.weights)
        self.optimize = tf.train.AdamOptimizer(self.lr).apply_gradients(grads)

        self.sess.run(tf.global_variables_initializer())

    def train(self, feats, auxs, encodes, action_grads, lr):
        self.sess.run(self.optimize, feed_dict={
            self.feats: feats,
            self.auxs: auxs,
            self.encodes: encodes,
            self.lr: lr,
            self.action_gradient: action_grads,
            K.learning_phase(): 1
        })

    def create_generator(self, feat_dim, aux_dim, encode_dim):
        feats = Input(shape=[feat_dim[0], feat_dim[1], feat_dim[2]])
        x = Convolution2D(256, 3, 3)(feats)
        x = LeakyReLU()(x)
        x = Convolution2D(256, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Flatten()(x)
        auxs = Input(shape=[aux_dim])
        h = merge([x, auxs], mode='concat')
        h = Dense(256)(h)
        h = LeakyReLU()(h)
        h = Dense(128)(h)
        encodes = Input(shape=[encode_dim])
        c = Dense(128)(encodes)
        h = merge([h, c], mode='sum')
        h = LeakyReLU()(h)

        steer = Dense(1, activation='tanh', init=lambda shape, name:
                      normal(shape, scale=1e-4, name=name))(h)
        accel = Dense(1, activation='sigmoid', init=lambda shape, name:
                             normal(shape, scale=1e-4, name=name))(h)
        brake = Dense(1, activation='sigmoid', init=lambda shape, name:
                      normal(shape, scale=1e-4, name=name))(h)
        actions = merge([steer, accel, brake], mode='concat')
        model = Model(input=[feats, auxs, encodes], output=actions)
        return model, model.trainable_weights, feats, auxs, encodes

class Posterior(object):
    def __init__(self, sess, img_dim, aux_dim, action_dim, encode_dim):
        self.sess = sess
        self.lr = tf.placeholder(tf.float32, shape=[])

        K.set_session(sess)

        self.model = self.create_posterior(img_dim, aux_dim, action_dim, encode_dim)

    def create_posterior(self, img_dim, aux_dim, action_dim, encode_dim):
        imgs = Input(shape=[img_dim[0], img_dim[1], img_dim[2]])
        x = Convolution2D(32, 3, 3, subsample=(2, 2))(imgs)
        x = LeakyReLU()(x)
        x = Convolution2D(64, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Convolution2D(128, 3, 3, subsample=(2, 2))(x)
        x = LeakyReLU()(x)
        x = Flatten()(x)
        auxs = Input(shape=[aux_dim])
        actions = Input(shape=[action_dim])
        h = merge([x, auxs, actions], mode='concat')
        h = Dense(256)(h)
        h = LeakyReLU()(h)
        h = Dense(128)(h)
        h = LeakyReLU()(h)
        c = Dense(encode_dim, activation='softmax')(h)

        model = Model(input=[imgs, auxs, actions], output=c)
        return model


# coding: utf-8
#
#    Project: Azimuthal integration
#             https://github.com/silx-kit/pyFAI
#
#    Copyright (C) 2015-2021 European Synchrotron Radiation Facility, Grenoble, France
#    Copyright (C)      2016 Synchrotron SOLEIL - L'Orme des Merisiers Saint-Aubin
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

"""Module for treating simultaneously multiple detector configuration
within a single integration"""

__author__ = "Jérôme Kieffer"
__contact__ = "Jerome.Kieffer@ESRF.eu"
__license__ = "MIT"
__copyright__ = "European Synchrotron Radiation Facility, Grenoble, France"
__date__ = "19/01/2021"
__status__ = "stable"
__docformat__ = 'restructuredtext'

import collections
import logging
logger = logging.getLogger(__name__)
from .azimuthalIntegrator import AzimuthalIntegrator
from .containers import Integrate1dResult
from .containers import Integrate2dResult
from . import units
import threading
import numpy
from .method_registry import IntegrationMethod
error = None

class MultiGeometry(object):
    """This is an Azimuthal integrator containing multiple geometries,
    for example when the detector is on a goniometer arm
    """

    def __init__(self, ais, unit="2th_deg",
                 radial_range=(0, 180), azimuth_range=(-180, 180),
                 wavelength=None, empty=0.0, chi_disc=180):
        """
        Constructor of the multi-geometry integrator
        :param ais: list of azimuthal integrators
        :param radial_range: common range for integration
        :param azimuthal_range: common range for integration
        :param empty: value for empty pixels
        :param chi_disc: if 0, set the chi_discontinuity at
        """
        self._sem = threading.Semaphore()
        self.abolute_solid_angle = None
        self.ais = [ai if isinstance(ai, AzimuthalIntegrator)
                    else AzimuthalIntegrator.sload(ai)
                    for ai in ais]
        self.wavelength = None
        if wavelength:
            self.set_wavelength(wavelength)
        self.radial_range = tuple(radial_range[:2])
        self.azimuth_range = tuple(azimuth_range[:2])
        self.unit = units.to_unit(unit)
        self.abolute_solid_angle = None
        self.empty = empty
        if chi_disc == 0:
            for ai in self.ais:
                ai.setChiDiscAtZero()
        elif chi_disc == 180:
            for ai in self.ais:
                ai.setChiDiscAtPi()
        else:
            logger.warning("Unable to set the Chi discontinuity at %s", chi_disc)

    def __repr__(self, *args, **kwargs):
        return "MultiGeometry integrator with %s geometries on %s radial range (%s) and %s azimuthal range (deg)" % \
            (len(self.ais), self.radial_range, self.unit, self.azimuth_range)

    def integrate1d(self, lst_data, npt=1800,
                    correctSolidAngle=True,
                    lst_variance=None, error_model=None,
                    polarization_factor=None,
                    normalization_factor=None,
                    lst_mask=None, lst_flat=None,
                    method="splitpixel"):
        """Perform 1D azimuthal integration

        :param lst_data: list of numpy array
        :param npt: number of points int the integration
        :param correctSolidAngle: correct for solid angle (all processing are then done in absolute solid angle !)
        :param lst_variance: list of array containing the variance of the data. If not available, no error propagation is done
        :type lst_variance: list of ndarray
        :param error_model: When the variance is unknown, an error model can be given: "poisson" (variance = I), "azimuthal" (variance = (I-<I>)^2)
        :type error_model: str
        :param polarization_factor: Apply polarization correction ? is None: not applies. Else provide a value from -1 to +1
        :param normalization_factor: normalization monitors value (list of floats)
        :param all: return a dict with all information in it (deprecated, please refer to the documentation of Integrate1dResult).
        :param lst_mask: numpy.Array or list of numpy.array which mask the lst_data.
        :param lst_flat: numpy.Array or list of numpy.array which flat the lst_data.
        :param method: integration method, a string or a registered method
        :return: 2th/I or a dict with everything depending on "all"
        :rtype: Integrate1dResult, dict
        """
        method = IntegrationMethod.select_one_available(method, dim=1)

        if len(lst_data) == 0:
            raise RuntimeError("List of images cannot be empty")
        if normalization_factor is None:
            normalization_factor = [1.0] * len(self.ais)
        elif not isinstance(normalization_factor, collections.Iterable):
            normalization_factor = [normalization_factor] * len(self.ais)
        if lst_variance is None:
            lst_variance = [None] * len(self.ais)
        if lst_mask is None:
            lst_mask = [None] * len(self.ais)
        elif isinstance(lst_mask, numpy.ndarray):
            lst_mask = [lst_mask] * len(self.ais)
        if lst_flat is None:
            lst_flat = [None] * len(self.ais)
        elif isinstance(lst_flat, numpy.ndarray):
            lst_flat = [lst_flat] * len(self.ais)
        signal = numpy.zeros(npt, dtype=numpy.float64)
        normalization = numpy.zeros_like(signal)
        count = numpy.zeros_like(signal)
        variance = None
        for ai, data, monitor, var, mask, flat in zip(self.ais, lst_data, normalization_factor, lst_variance, lst_mask, lst_flat):
            res = ai.integrate1d_ng(data, npt=npt,
                                    correctSolidAngle=correctSolidAngle,
                                    variance=var, error_model=error_model,
                                    polarization_factor=polarization_factor,
                                    radial_range=self.radial_range,
                                    azimuth_range=self.azimuth_range,
                                    method=method, unit=self.unit, safe=True,
                                    mask=mask, flat=flat, normalization_factor=monitor)
            sac = (ai.pixel1 * ai.pixel2 / ai.dist ** 2) if correctSolidAngle else 1.0
            count += res.count
            normalization += res.sum_normalization * sac
            signal += res.sum_signal
            if res.sigma is not None:
                if variance is None:
                    variance = res.sum_variance.astype(dtype=numpy.float64)  # explicit copy
                else:
                    variance += res.sum_variance

        tiny = numpy.finfo("float32").tiny
        norm = numpy.maximum(normalization, tiny)
        invalid = count <= 0.0
        I = signal / norm
        I[invalid] = self.empty

        if variance is not None:
            sigma = numpy.sqrt(variance) / norm
            sigma[invalid] = self.empty
            result = Integrate1dResult(res.radial, I, sigma)
        else:
            result = Integrate1dResult(res.radial, I)
        result._set_compute_engine(res.compute_engine)
        result._set_unit(self.unit)
        result._set_sum_signal(signal)
        result._set_sum_normalization(normalization)
        result._set_sum_variance(variance)
        result._set_count(count)
        return result

    def integrate2d(self, lst_data, npt_rad=1800, npt_azim=3600,
                    correctSolidAngle=True,
                    lst_variance=None, error_model=None,
                    polarization_factor=None,
                    normalization_factor=None, lst_mask=None,
                    lst_flat=None, method="splitpixel"):
        """Performs 2D azimuthal integration of multiples frames, one for each geometry

        :param lst_data: list of numpy array
        :param npt: number of points int the integration
        :param correctSolidAngle: correct for solid angle (all processing are then done in absolute solid angle !)
        :param lst_variance: list of array containing the variance of the data. If not available, no error propagation is done
        :type lst_variance: list of ndarray
        :param error_model: When the variance is unknown, an error model can be given: "poisson" (variance = I), "azimuthal" (variance = (I-<I>)^2)
        :type error_model: str
        :param polarization_factor: Apply polarization correction ? is None: not applies. Else provide a value from -1 to +1
        :param normalization_factor: normalization monitors value (list of floats)
        :param all: return a dict with all information in it (deprecated, please refer to the documentation of Integrate2dResult).
        :param lst_mask: numpy.Array or list of numpy.array which mask the lst_data.
        :param lst_flat: numpy.Array or list of numpy.array which flat the lst_data.
        :param method: integration method (or its name)
        :return: I/2th/chi or a dict with everything depending on "all"
        :rtype: Integrate2dResult, dict
        """
        if len(lst_data) == 0:
            raise RuntimeError("List of images cannot be empty")
        if normalization_factor is None:
            normalization_factor = [1.0] * len(self.ais)
        elif not isinstance(normalization_factor, collections.Iterable):
            normalization_factor = [normalization_factor] * len(self.ais)
        if lst_variance is None:
            lst_variance = [None] * len(self.ais)
        if lst_mask is None:
            lst_mask = [None] * len(self.ais)
        elif isinstance(lst_mask, numpy.ndarray):
            lst_mask = [lst_mask] * len(self.ais)
        if lst_flat is None:
            lst_flat = [None] * len(self.ais)
        elif isinstance(lst_flat, numpy.ndarray):
            lst_flat = [lst_flat] * len(self.ais)

        method = IntegrationMethod.select_one_available(method, dim=2)

        signal = numpy.zeros((npt_azim, npt_rad), dtype=numpy.float64)
        count = numpy.zeros_like(signal)
        normalization = numpy.zeros_like(signal)
        variance = None

        for ai, data, monitor, var, mask, flat in zip(self.ais, lst_data, normalization_factor, lst_variance, lst_mask, lst_flat):
            res = ai.integrate2d_ng(data, npt_rad=npt_rad, npt_azim=npt_azim,
                                    correctSolidAngle=correctSolidAngle,
                                    variance=var, error_model=error_model,
                                    polarization_factor=polarization_factor,
                                    radial_range=self.radial_range,
                                    azimuth_range=self.azimuth_range,
                                    method=method, unit=self.unit, safe=True,
                                    mask=mask, flat=flat, normalization_factor=monitor)
            sac = (ai.pixel1 * ai.pixel2 / ai.dist ** 2) if correctSolidAngle else 1.0
            count += res.count
            signal += res.sum_signal
            normalization += res.sum_normalization * sac
            if var is not None:
                if variance is None:
                    variance = res.sum_variance.astype(numpy.float64)  # explicit copy !
                else:
                    variance += res.sum_variance

        tiny = numpy.finfo("float32").tiny
        norm = numpy.maximum(normalization, tiny)
        invalid = count <= 0
        I = signal / norm
        I[invalid] = self.empty

        if variance is not None:
            sigma = numpy.sqrt(variance) / norm
            sigma[invalid] = self.empty
            result = Integrate2dResult(I, res.radial, res.azimuthal, sigma)
        else:
            result = Integrate2dResult(I, res.radial, res.azimuthal)
        result._set_sum(signal)
        result._set_compute_engine(res.compute_engine)
        result._set_unit(self.unit)
        result._set_sum_signal(signal)
        result._set_sum_normalization(normalization)
        result._set_sum_variance(variance)
        result._set_count(count)
        return result

    def set_wavelength(self, value):
        """
        Changes the wavelength of a group of azimuthal integrators
        """
        self.wavelength = float(value)
        for ai in self.ais:
            ai.set_wavelength(self.wavelength)


#!/usr/bin/env python
import datetime
import math
import os
import sys
import threading
import time
import yaml

import numpy as np
from pprint import pprint

import mxnet
from mxnet import gluon, nd, gpu
from mxboard import SummaryWriter

from gluoncv.model_zoo.yolo.yolo3 import YOLODetectionBlockV3
from gluoncv.model_zoo.yolo.yolo3 import _upsample

from yolo_modules import basic_yolo

# self define modules
from yolo_modules import yolo_gluon
from yolo_modules import yolo_cv
from yolo_modules import licence_plate_render
from yolo_modules import global_variable

from car import utils
from car.render_car import RenderCar
import car.YOLO as car_YOLO
import licence_plate.LP_detection as LP_detection

render_thread_pre_load = True
export_onnx = False

def main():
    args = utils.yolo_Parser()
    yolo = YOLO(args)

    assert args.mode in car_YOLO.available_mode, \
        'Available Modes Are {}'.format(available_mode)

    exec "yolo.%s()" % args.mode

class CarLPNet(basic_yolo.BasicYOLONet):
    def __init__(self, spec, num_sync_bn_devices=-1, **kwargs):
        super(CarLPNet, self).__init__(spec, num_sync_bn_devices, **kwargs)

        LP_channel = spec['channels'][-3]
        LP_slice_point = spec['LP_slice_point']

        self.LP_branch = mxnet.gluon.nn.HybridSequential()
        self.LP_branch.add(YOLODetectionBlockV3(LP_channel, num_sync_bn_devices))
        self.LP_branch.add(YOLODetectionBlockV3(LP_channel, num_sync_bn_devices))
        self.LP_branch.add(YOLODetectionBlockV3(LP_channel, num_sync_bn_devices))
        self.LP_branch.add(YOLODetectionBlockV3(LP_channel, num_sync_bn_devices))
        self.LP_branch.add(YOLODetectionBlockV3(LP_channel, num_sync_bn_devices))
        self.LP_branch.add(gluon.nn.Conv2D(LP_slice_point[-1], kernel_size=1))

    def hybrid_forward(self, F, x, *args):
        routes = []
        all_output = []
        for i, stage in enumerate(self.stages):
            x = stage(x)
            if i >= len(self.stages) - self.num_pyrmaid_layers:
                routes.append(x)

        # the YOLO output layers are used in reverse order, i.e., from very deep layers to shallow
        end = False
        for i, block, output in zip(range(len(routes)), self.yolo_blocks, self.yolo_outputs):
            if i >= len(routes) - 1:
                _, LP_output = self.LP_branch[0](x)
                _, LP_output = self.LP_branch[1](LP_output)
                _, LP_output = self.LP_branch[2](LP_output)
                _, LP_output = self.LP_branch[3](LP_output)
                _, LP_output = self.LP_branch[4](LP_output)
                LP_output = self.LP_branch[5](LP_output)
                LP_output = LP_output.transpose((0, 2, 3, 1))
                end = True

            x, tip = block(x)
            all_output.append(output(tip))

            if end:
                break
            # add transition layers
            x = self.transitions[i](x)

            # upsample feature map reverse to shallow layers
            upsample = _upsample(x, stride=2)
            x = F.concat(upsample, routes[::-1][i + 1], dim=1)

        return all_output[::-1], [LP_output]

class YOLO(car_YOLO.YOLO, LP_detection.LicencePlateDetectioin):
    def __init__(self, args):
        car_YOLO.YOLO.__init__(self, args)
        # pprint(self.classes)

    def _init_net(self, spec, weight):
        # Do not set num_sync_bn_devices=len(self.ctx)
        # because No conversion function for contrib_SyncBatchNorm yet.
        # (ONNX)
        # self.net = CarLPNet(spec, num_sync_bn_devices=len(self.cx))

        self.net = CarLPNet(spec, num_sync_bn_devices=-1)
        self.num_downsample = len(spec['layers']) - 2  # For LP training

        # self.backup_dir = os.path.join(self.version, 'backup')
        self.backup_dir = os.path.join(
            '/media/nolan/SSD1/YOLO_backup/car_and_LP',
            self.version,
            'backup')

        if weight is None:
            weight = yolo_gluon.get_latest_weight_from(self.backup_dir)

        yolo_gluon.init_NN(self.net, weight, self.ctx)

    # -------------------- LP -------------------- #
    def _score_weight_LP(self, mask, ctx):
        n = self.LP_negative_weight
        p = self.LP_positive_weight

        ones = nd.ones_like(mask)
        score_weight = nd.where(mask > 0, ones*p, ones*n, ctx=ctx)

        return score_weight

    def predict_LP(self, LP_batch_out):
        # LP_batch_out = self.fp16_2_fp32(LP_batch_out)
        LP_batch_out = self.merge_and_slice(LP_batch_out, self.LP_slice_point)

        LP_score = nd.sigmoid(LP_batch_out[0])
        LP_pose_xy = LP_batch_out[1]
        LP_pose_z = LP_batch_out[2]
        LP_pose_r = LP_batch_out[3]
        LP_batch_out = nd.concat(
            LP_score, LP_pose_xy, LP_pose_z, LP_pose_r, dim=-1)

        LP_batch_out = nd.split(LP_batch_out, axis=0, num_outputs=len(LP_batch_out))

        LP_batch_pred = []
        for i, out in enumerate(LP_batch_out):
            best_index = LP_score[i].reshape(-1).argmax(axis=0)
            out = out.reshape((-1, 7))

            pred = out[best_index][0]  # best out
            pred[1:7] = self.LP_pose_activation(pred[1:7])
            LP_batch_pred.append(nd.expand_dims(pred, axis=0))

        LP_batch_pred = nd.concat(*LP_batch_pred, dim=0)

        return LP_batch_pred.asnumpy()

    def LP_pose_activation(self, data_in):
        data_out = nd.zeros(6)

        data_out[0] = data_in[0] * 1000
        data_out[1] = data_in[1] * 1000
        data_out[2] = data_in[2] * 1000
        for i in range(3):
            data = (nd.sigmoid(data_in[i+3]) - 0.5) * 2 * self.LP_r_max[i]
            data_out[i+3] = data * math.pi / 180.

        return data_out

    # -------------------- Training Main -------------------- #
    def render_and_train(self):
        print(global_variable.green)
        print('Render And Train')
        print(global_variable.reset_color)
        # -------------------- show training image # --------------------
        '''
        self.batch_size = 1
        ax = yolo_cv.init_matplotlib_figure()
        '''
        h, w = self.size
        bs = self.batch_size
        # -------------------- background -------------------- #
        self.bg_iter_valid = yolo_gluon.load_background('val', bs, h, w)
        self.bg_iter_train = yolo_gluon.load_background('train', bs, h, w)

        self.car_renderer = RenderCar(h, w, self.classes, self.ctx[0], pre_load=False)
        LP_generator = licence_plate_render.LPGenerator(h, w)

        # -------------------- main loop -------------------- #
        while True:
            if (self.backward_counter % 10 == 0 or 'bg' not in locals()):
                bg = yolo_gluon.ImageIter_next_batch(self.bg_iter_train)
                bg = bg.as_in_context(self.ctx[0])

            # -------------------- render dataset -------------------- #
            imgs, labels = self.car_renderer.render(
                bg, 'train', render_rate=0.5, pascal_rate=0.1)

            imgs, LP_labels = LP_generator.add(imgs, self.LP_r_max, add_rate=0.5)

            batch_xs = yolo_gluon.split_render_data(imgs, self.ctx)
            car_batch_ys = yolo_gluon.split_render_data(labels, self.ctx)
            LP_batch_ys = yolo_gluon.split_render_data(LP_labels, self.ctx)

            self._train_batch(batch_xs, car_batch_ys, LP_batch_ys)

            # -------------------- show training image # --------------------
            '''
            img = yolo_gluon.batch_ndimg_2_cv2img(batch_xs[0])[0]
            img = yolo_cv.cv2_add_bbox(img, car_batch_ys[0][0, 0].asnumpy(), 4, use_r=0)
            yolo_cv.matplotlib_show_img(ax, img)
            print(car_batch_ys[0][0])
            raw_input()
            '''

    def _render_thread(self):
        h, w = self.size
        bs = self.batch_size

        self.bg_iter_valid = yolo_gluon.load_background('val', bs, h, w)
        bg_iter_train = yolo_gluon.load_background('train', bs, h, w)
        LP_generator = licence_plate_render.LPGenerator(h, w)

        self.car_renderer = RenderCar(
            h, w, self.classes, self.ctx[0], pre_load=render_thread_pre_load)

        while not self.shutdown_training:
            if self.rendering_done:
                #print('render done')
                time.sleep(0.01)
                continue

            # ready to render new images
            if (self.backward_counter % 10 == 0 or 'bg' not in locals()):
                bg = yolo_gluon.ImageIter_next_batch(bg_iter_train)
                bg = bg.as_in_context(self.ctx[0])

            # change an other batch of background
            imgs, self.labels = self.car_renderer.render(
                bg, 'train', render_rate=0.5, pascal_rate=0.2)

            self.imgs, self.LP_labels = LP_generator.add(
                imgs, self.LP_r_max, add_rate=0.5)
            self.rendering_done = True

    def _train_thread(self):
        while not self.shutdown_training:
            if not self.rendering_done:
                # training images are not ready
                #print('rendering')
                time.sleep(0.01)
                continue

            batch_xs = self.imgs.copy()
            car_batch_ys = self.labels.copy()
            LP_batch_ys = self.LP_labels.copy()
            batch_xs = yolo_gluon.split_render_data(batch_xs, self.ctx)
            car_batch_ys = yolo_gluon.split_render_data(car_batch_ys, self.ctx)
            LP_batch_ys = yolo_gluon.split_render_data(LP_batch_ys, self.ctx)

            self.rendering_done = False
            self._train_batch(batch_xs, car_batch_ys, LP_batch_ys)

    def _train_batch(self, bxs, car_bys, LP_bys, car_rotate=False):
        all_gpu_loss = []
        with mxnet.autograd.record():
            for gpu_i in range(len(bxs)):
                all_gpu_loss.append([])  # new loss list for gpu_i
                ctx = self.ctx[gpu_i]  # gpu_i = GPU index

                bx = bxs[gpu_i]  # .astype('float16', copy=False)

                car_by = car_bys[gpu_i]
                LP_by = LP_bys[gpu_i]

                x, LP_x = self.net(bx)  # [x, x, x], [x]
                x = self.merge_and_slice(x, self.slice_point)  # from car_YOLO
                LP_x = self.merge_and_slice(LP_x, self.LP_slice_point)

                with mxnet.autograd.pause():

                    y, mask = self._loss_mask(car_by, gpu_i)
                    s_weight = self._score_weight(mask, ctx)
                    # y = self.fp32_2_fp16(y)
                    # s_weight, mask = self.fp32_2_fp16([s_weight, mask])

                    LP_y, LP_mask = self._loss_mask_LP(LP_by, gpu_i)
                    LP_s_weight = self._score_weight_LP(LP_mask, ctx)
                    # LP_y = self.fp32_2_fp16(LP_y)
                    # LP_s_weight, LP_mask = self.fp32_2_fp16([LP_s_weight, LP_mask])

                car_loss = self._get_loss(x, y, s_weight, mask, car_rotate=car_rotate)
                all_gpu_loss[gpu_i].extend(car_loss)

                LP_loss = self._get_loss_LP(LP_x, LP_y, LP_s_weight, LP_mask)
                all_gpu_loss[gpu_i].extend(LP_loss)

                sum(all_gpu_loss[gpu_i]).backward()

        self.trainer.step(self.batch_size)

        if self.record:
            self._record_to_tensorboard_and_save(all_gpu_loss[0])

    def _valid_iou(self):
        for pascal_rate in [1, 0]:
            iou_sum = 0
            c = 0
            for bg in self.bg_iter_valid:
                c += 1
                bg = bg.data[0].as_in_context(self.ctx[0])
                imgs, labels = self.car_renderer.render(bg, 'valid', pascal_rate=pascal_rate)

                # imgs = imgs.astype('float16', copy=False)
                x, LP_x = self.net(imgs)
                outs = self.predict(x)

                pred = nd.zeros((self.batch_size, 4))
                pred[:, 0] = outs[:, 2] - outs[:, 4] / 2
                pred[:, 1] = outs[:, 1] - outs[:, 3] / 2
                pred[:, 2] = outs[:, 2] + outs[:, 4] / 2
                pred[:, 3] = outs[:, 1] + outs[:, 3] / 2
                pred = pred.as_in_context(self.ctx[0])

                for i in range(self.batch_size):
                    label = labels[i, 0, 0:5]
                    iou_sum += yolo_gluon.get_iou(pred[i], label, mode=2)

            mean_iou = iou_sum.asnumpy() / float(self.batch_size * c)
            self.sw.add_scalar(
                'Mean_IOU',
                (self.exp + 'PASCAL %r' % pascal_rate, mean_iou),
                self.backward_counter)

            self.bg_iter_valid.reset()

    # -------------------- Validation Part -------------------- #
    def valid(self):
        print(global_variable.cyan)
        print('Valid')

        bs = 1  # batch size = 1
        h, w = self.size

        # init two matplotlib figures
        ax1 = yolo_cv.init_matplotlib_figure()
        ax2 = yolo_cv.init_matplotlib_figure()

        # init radar figure for vizualizing class distribution
        radar_prob = yolo_cv.RadarProb(self.num_class, self.classes)

        # init background, LP adder, car renderer
        BG_iter = yolo_gluon.load_background('val', bs, h, w)
        LP_generator = licence_plate_render.LPGenerator(h, w)
        car_renderer = RenderCar(h, w, self.classes, self.ctx[0], pre_load=False)

        for bg in BG_iter:
            # select background
            bg = bg.data[0].as_in_context(self.ctx[0])  # b*RGB*w*h
            # render images, type(imgs) = mxnet.ndarray
            imgs, labels = car_renderer.render(bg, 'valid', pascal_rate=0.5, render_rate=0.9)
            imgs, LP_labels = LP_generator.add(imgs, self.LP_r_max, add_rate=0.8)

            # return all_output[::-1], [LP_output]
            x1, x2, x3, LP_x = self.net.forward(is_train=False, data=imgs)
            outs = self.predict([x1, x2, x3])
            LP_outs = self.predict_LP([LP_x])

            # convert ndarray to np.array
            img = yolo_gluon.batch_ndimg_2_cv2img(imgs)[0]

            # draw licence plate border
            img, clipped_LP = LP_generator.project_rect_6d.add_edges(img, LP_outs[0, 1:])
            yolo_cv.matplotlib_show_img(ax2, clipped_LP)

            # draw car border
            img = yolo_cv.cv2_add_bbox(img, labels[0, 0].asnumpy(), 4, use_r=0)
            img = yolo_cv.cv2_add_bbox(img, outs[0], 5, use_r=0)
            yolo_cv.matplotlib_show_img(ax1, img)

            # vizualize class distribution
            radar_prob.plot(outs[0, 0], outs[0, -self.num_class:])
            raw_input('next')

    def export(self):
        yolo_gluon.export(
            self.net,
            (1, 3, self.size[0], self.size[1]),
            self.ctx[0],
            self.export_folder,
            onnx=False,
            fp16=False)

# -------------------- Main -------------------- #
if __name__ == '__main__':
    main()


from transformers import RobertaTokenizer, RobertaModel
from transformers import DeiTFeatureExtractor, DeiTModel

import torch
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import os, sys
import pdb

class BaseModel(nn.Module):
    def __init__(self, post_back):
        super(BaseModel, self).__init__()
        """RoBERTa setting"""
        text_model_path = '/data/project/rw/rung/02_source/model/roberta-large' # "roberta-large" #
        self.text_model = RobertaModel.from_pretrained(text_model_path)
        self.text_tokenizer = RobertaTokenizer.from_pretrained(text_model_path)
        special_token_list = ['[USER]', '[SYSTEM]']
        special_tokens = {'additional_special_tokens': special_token_list}
        self.text_tokenizer.add_special_tokens(special_tokens)
        self.text_model.resize_token_embeddings(len(self.text_tokenizer))

        system_pos = self.text_tokenizer.additional_special_tokens.index('[SYSTEM]')
        self.system_token_id = self.text_tokenizer.additional_special_tokens_ids[system_pos]

        """Deit setting"""
        image_model_path = '/data/project/rw/rung/02_source/model/deit-base-distilled-patch16-224' # "facebook/deit-base-distilled-patch16-224" #
        self.image_model = DeiTModel.from_pretrained(image_model_path, add_pooling_layer=False)

        if post_back:
            self.back_model = DeiTModel.from_pretrained(image_model_path, add_pooling_layer=False)
            post_back_path = "../BTM/bg_model"
            post_back_model = os.path.join(post_back_path, 'model.pt')
            checkpoint = torch.load(post_back_model)
            self.back_model.load_state_dict(checkpoint, strict=False)
            print('PostBack-trained Model Loading!!')

        """similarity"""
        self.hid_dim = 128
        self.text2hid = nn.Linear(self.text_model.config.hidden_size, self.hid_dim) # (1024, 128) self.text_model.config.hidden_size
        self.visual2hid = nn.Linear(self.image_model.config.hidden_size, self.hid_dim) # (768, 128) self.image_model.config.hidden_size
        self.bg2hid = nn.Linear(self.image_model.config.hidden_size, self.hid_dim) # (768, 128) self.image_model.config.hidden_size

        """multi-task"""
        # utterance category prediction
        self.W_category = nn.Linear(self.hid_dim, 4) # (128, 4)
        # system matching
        self.W_system = nn.Linear(self.text_model.config.hidden_size, 2) # (1024, 2)

        """score"""
        self.softmax = nn.Softmax(dim=1)
        self.cos = nn.CosineSimilarity(dim=1, eps=1e-6)
        self.sigmoid = nn.Sigmoid()

    def imgfea2rep(self, image_features):
        patch_represent = self.image_model(image_features)['last_hidden_state'] # (1, 196, 768)
        image_represent = patch_represent[:,0,:] # (1, 768)
        return image_represent

    def backfea2rep(self, bg_features):
        patch_represent = self.back_model(bg_features)['last_hidden_state'] # (1, 196, 768)
        image_represent = patch_represent[:,0,:] # (1, 768)
        return image_represent

    def forward(self, text_tokens, meta_tokens, object_features, background_features, score_type, meta, system_matching, utt_category, background, post_back):
        """ dialogue text projection"""
        text_logits = self.text_model(text_tokens).last_hidden_state # (batch, text_len, 1024)
        text_represent = text_logits[:,0,:] # (batch, 1024)
        t_h = self.text2hid(text_represent) # (batch, 128)

        """ meta text projection"""
        if meta:
            meta_logits = self.text_model(meta_tokens).last_hidden_state # (batch, text_len, 1024)
            meta_represent = meta_logits[:,0,:] # (batch, 1024)
            m_h = self.text2hid(meta_represent) # (batch, 128)

        """ object projection """
        image_represent = self.imgfea2rep(object_features) # (batch, 768)
        o_h = self.visual2hid(image_represent) # (batch, 128)

        """ background proejction """
        if background:
            if post_back:
                bg_represent = self.backfea2rep(background_features) # (batch, 768)
            else:
                bg_represent = self.imgfea2rep(background_features) # (batch, 768)
            bg_h = self.bg2hid(bg_represent) # (batch, 128)
        else:
            bg_h = torch.zeros(o_h.shape).cuda()

        """ total visual represent vector """
        v_h = o_h + bg_h

        """ system matching """
        system_logits_list = []
        if system_matching:
            for batch, text_token in enumerate(text_tokens):
                # text_token: (text_len)
                system_pos_list = []
                for system_pos, token in enumerate(text_token.tolist()):
                    if token == self.system_token_id:
                        system_pos_list.append(system_pos)

                system_hidden = []
                for system_pos in system_pos_list:
                    system_hidden.append(text_logits[batch:batch+1,system_pos,:]) # [(1, 1024), ..., ]

                if len(system_hidden) == 0:
                    system_logits_list.append([])
                else:
                    system_hiddens = torch.cat(system_hidden, 0) # (len, 1024)
                    system_logits = self.W_system(system_hiddens) # (len, 2)
                    system_logits_list.append(system_logits) # [(len1,2), (len2,2)]
        else:
            system_logits_list.append([])

        """ utterance category prediction"""
        if utt_category:
            utt_category_logits = self.W_category(t_h) # (batch, 4)
        else:
            utt_category_logits = torch.zeros([t_h.shape[0],4])

        """ dialog-object score """
        if score_type == 'cos':
            t2v_score = (self.cos(t_h, v_h)+1)/2
        elif score_type == 'sigmoid':
            t2v_score = self.sigmoid(self.cos(t_h, v_h) * 100) # (batch)
        else: ## concat
            final_h = torch.cat([t_h, v_h], 1) # (1, 128*x)
            t2v_logit = self.concat2score2(final_h) # (1,2)
            t2v_score = self.softmax(t2v_logit)[:,1]

        """ meta-object score """
        if meta:
            m2v_score = self.sigmoid(self.cos(m_h, v_h) * 100) # (batch)
            return t2v_score, m2v_score, system_logits_list, utt_category_logits
        else:
            m2v_score = torch.zeros(t2v_score.shape)
            return t2v_score, m2v_score, system_logits_list, utt_category_logits


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import shutil
import sys

import tensorflow as tf
import numpy as np

import mysql.connector
import time

config = {
    'auth_plugin': 'mysql_native_password',
    'user': '',
    'password': '',
    'host': 'localhost',
    'ssl_ca': 'C:\certificates\ca.pem',
    'ssl_cert': 'C:\certificates\client-cert.pem',
    'ssl_key': 'C:\certificates\client-key.pem',
}

parser = argparse.ArgumentParser()

parser.add_argument(
    '--model_dir', type=str, default='/tmp/dota_model',
    help='Base directory for the model.')

parser.add_argument(
    '--train_epochs', type=int, default=32, help='Number of training epochs.')

parser.add_argument(
    '--epochs_per_eval', type=int, default=4,
    help='The number of training epochs to run between evaluations.')

parser.add_argument(
    '--batch_size', type=int, default=40, help='Number of examples per batch.')

parser.add_argument(
    '--train', default=False, action='store_const', const=True,
)

parser.add_argument(
    '--predict', type=str, default='',
    help='Input vector'
)

_NUM_EXAMPLES = {
    'train': 80,
    'validation': 20,
}

cnx = None
cursor = None
num_heroes = 120

def build_model_columns():
    feature_columns = [tf.feature_column.numeric_column('x', shape=num_heroes*2)]

    return feature_columns

def build_estimator(model_dir):
    deep_columns = build_model_columns()
    hidden_units = [180]

    run_config = tf.estimator.RunConfig().replace(
        session_config=tf.ConfigProto(device_count={'GPU': 0}))

    if FLAGS.train:
        return tf.estimator.DNNClassifier(
            model_dir=model_dir,
            feature_columns=deep_columns,
            hidden_units=hidden_units,
            config=run_config)
    else:
        return tf.estimator.DNNClassifier(
            model_dir=model_dir,
            feature_columns=deep_columns,
            hidden_units=hidden_units,
            config=run_config,
            warm_start_from=model_dir)

def get_sample_count():
    query = ("SELECT COUNT(*) FROM matches")
    cursor.execute(query)
    return cursor.fetchone()[0]

def setup_database():
    global cnx
    global cursor
    global num_heroes

    cnx = mysql.connector.connect(**config)
    cursor = cnx.cursor()
    query = ("USE dota2")
    cursor.execute(query)

    query = ("SELECT id FROM heroes")
    cursor.execute(query)
    h_max = -1
    for h in cursor.fetchall():
        if h[0] > h_max:
            h_max = h[0]
    print(h_max)
    num_heroes = h_max

def query_matches(train_test):
    if train_test == "test":
        query = ("SELECT * FROM matches ORDER BY match_id DESC LIMIT %s")
        data = (_NUM_EXAMPLES['validation'],)
        cursor.execute(query, data)
        print("Test set of size: " + str(_NUM_EXAMPLES['validation']))
    elif train_test == "train":
        query = ("SELECT * FROM matches ORDER BY match_id DESC LIMIT %s")
        data = (_NUM_EXAMPLES['validation'],)
        cursor.execute(query, data)
        max_id = cursor.fetchall()[-1:][0][0]
        print("Starting at match_id: " + str(max_id))
        query = ("SELECT * FROM matches WHERE match_id<%s ORDER BY match_id DESC LIMIT %s")
        data = (max_id,_NUM_EXAMPLES['train'])
        cursor.execute(query, data)
        print("Training set of size: " + str(_NUM_EXAMPLES['train']))
    else:
        print("wrong train_test value")
    matches = []
    m_columns = cursor.column_names
    for m in cursor.fetchall():
        m = dict(zip(m_columns, m))
        matches.append({'match_id': m['match_id'], 'radiant_win': m['radiant_win'], 'radiant': [], 'dire': []})
    match_hero_stmt = ("SELECT * FROM match_hero WHERE match_id=%s")
    for m in matches:
        data = (m['match_id'],)
        cursor.execute(match_hero_stmt,data)
        mh_columns = cursor.column_names
        for mh in cursor.fetchall():
            mh = dict(zip(mh_columns, mh))
            if mh['player_slot'] < 128:
                m['radiant'].append(mh)
            else:
                m['dire'].append(mh)
    return matches

def format_matches(matches):
    features = []
    labels = []
    count = 0
    for m in matches:
        r_indices = []
        d_indices = []
        for r in m['radiant']:
            r_indices.append(r['hero_id']-1)
        for d in m['dire']:
            d_indices.append(d['hero_id']-1 + num_heroes)
        depth = 2*num_heroes
        r_indices = np.concatenate((r_indices, d_indices))
        r_tensor = tf.one_hot(r_indices, depth, dtype=tf.int32)
        features.append(sum(r_tensor.eval()))
        labels.append(m['radiant_win'])
        count = count + 1
        if count % 100 == 0:
            print("Formatted " + str(count) + " matches.")
    features = {'x': np.asarray(features)}
    return features, labels

def input_fn(num_epochs, shuffle, batch_size, train_test):

    def query_format_matches():
        print('Querying matches')
        matches = query_matches(train_test)
        print('Formatting matches')
        f, l = format_matches(matches)
        return f, l

    print("Creating dataset from tensor slices")
    sess = tf.InteractiveSession()
    sess.as_default()
    dataset = tf.data.Dataset.from_tensor_slices(query_format_matches())
    sess.close()
    if shuffle:
        dataset = dataset.shuffle(buffer_size=_NUM_EXAMPLES['train'])

    dataset = dataset.repeat(num_epochs)
    dataset = dataset.batch(batch_size)

    return dataset

def predict_fn(input_vector):

    input_v = input_vector.split(',')

    features = []
    r_indices = input_v[:5]
    d_indices = input_v[5:]
    for i in range(5):
        r_indices[i] = int(r_indices[i])-1
        d_indices[i] = int(d_indices[i]) - 1 + num_heroes
    depth = 2*num_heroes
    indices = np.concatenate((r_indices, d_indices))
    tensor = tf.one_hot(indices, depth, dtype=tf.int32)
    sess = tf.InteractiveSession()
    sess.as_default()
    features.append(sum(tensor.eval()))
    sess.close()
    feature = {'x': np.asarray(features)}

    dataset = tf.data.Dataset.from_tensors(feature)

    return dataset

# attempt to use a mapping function for data input
# def input_fn(num_epochs, shuffle, batch_size, train_test):
#
#     millis = int(round(time.time() * 1000))
#
#     def format_match(match):
#         r_indices = []
#         d_indices = []
#         for r in match['radiant']:
#             r_indices.append(r['hero_id']-1)
#         for d in match['dire']:
#             d_indices.append(d['hero_id']-1 + num_heroes)
#         depth = 2*num_heroes
#         r_tensor = tf.one_hot(r_indices, depth, dtype=tf.int32)
#         d_tensor = tf.one_hot(d_indices, depth, dtype=tf.int32)
#         sess = tf.InteractiveSession()
#         features = (tf.convert_to_tensor(sum(r_tensor.eval())+sum(d_tensor.eval())))
#         label = (tf.convert_to_tensor(match['radiant_win']))
#         sess.close()
#         features = {'x': features}
#         return features, label
#
#     print('Querying matches')
#     matches = query_matches(train_test)
#
#     dataset = tf.data.Dataset.from_tensor_slices(matches)
#
#     if shuffle:
#         dataset = dataset.shuffle(buffer_size=_NUM_EXAMPLES['train'])
#
#     dataset = dataset.map(format_match, num_parallel_calls=8)
#
#     dataset = dataset.repeat(num_epochs)
#     dataset = dataset.batch(batch_size)
#
#     millis2 = int(round(time.time() * 1000))
#
#     print('-' * 60)
#     print("input time")
#     print(millis2-millis)
#     print('-' * 60)
#
#     return dataset

def main(unused_argv):

    config_file = open("config_file.txt", 'r')
    db_user = config_file.readline()
    db_pass = config_file.readline()
    config_file.close()

    config['user'] = db_user.rstrip()
    config['password'] = db_pass.rstrip()

    setup_database()

    num_samples = get_sample_count()

    global _NUM_EXAMPLES
    _NUM_EXAMPLES = {
        'train': num_samples - int(0.2 * num_samples),
        'validation': int(0.2 * num_samples),
    }

    if not FLAGS.train and not FLAGS.predict:
        print('No training or prediction called')
        return

    if FLAGS.train:

        # Clean up the model directory if present
        shutil.rmtree(FLAGS.model_dir, ignore_errors=True)
        model = build_estimator(FLAGS.model_dir)

        millis = int(round(time.time() * 1000))

        # Train and evaluate the model every `FLAGS.epochs_per_eval` epochs.
        for n in range(FLAGS.train_epochs // FLAGS.epochs_per_eval):
            model.train(input_fn=lambda: input_fn(
                FLAGS.epochs_per_eval, True, FLAGS.batch_size, "train"))

            results = model.evaluate(input_fn=lambda: input_fn(
                1, False, FLAGS.batch_size, "test"))

            # Display evaluation metrics
            print('Results at epoch', (n + 1) * FLAGS.epochs_per_eval)
            print('-' * 60)

            for key in sorted(results):
                print('%s: %s' % (key, results[key]))
        millis2 = int(round(time.time() * 1000))

        print('-' * 60)
        print("input time")
        print(millis2-millis)
        print('-' * 60)

    else:
        model = build_estimator(FLAGS.model_dir)

    if FLAGS.predict:

        prediction = model.predict(input_fn=lambda: predict_fn(FLAGS.predict))

        print('-' * 60)
        print('Prediction')
        print('-' * 60)
        for value in prediction:
            print(str(value['logistic'])[1:-1])

if __name__ == '__main__':
    tf.logging.set_verbosity(tf.logging.INFO)
    FLAGS, unparsed = parser.parse_known_args()
    tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)


import numpy as np
import pyaudio as pa
import struct
import matplotlib.pyplot as plt

CHUNK = 1024 * 2
FORMAT = pa.paInt16
CHANNELS = 1
RATE = 44100 # in Hz

p = pa.PyAudio()

stream = p.open(
    format = FORMAT,
    channels = CHANNELS,
    rate = RATE,
    input=True,
    output=True,
    frames_per_buffer=CHUNK
)

fig,ax = plt.subplots()
x = np.arange(0,2*CHUNK,2)
line, = ax.plot(x, np.random.rand(CHUNK),'r')
ax.set_ylim(-60000,60000)
ax.ser_xlim = (0,CHUNK)
fig.show()

while 1:
    data = stream.read(CHUNK)
    dataInt = struct.unpack(str(CHUNK) + 'h', data)
    line.set_ydata(dataInt)
    fig.canvas.draw()
    fig.canvas.flush_events()


import warnings

import numpy as np
import pyqtgraph as pg

import acconeer.exptool as et

def main():
    args = et.utils.ExampleArgumentParser(num_sens=1).parse_args()
    et.utils.config_logging(args)

    if args.socket_addr:
        client = et.SocketClient(args.socket_addr)
    elif args.spi:
        client = et.SPIClient()
    else:
        port = args.serial_port or et.utils.autodetect_serial_port()
        client = et.UARTClient(port)

    sensor_config = get_sensor_config()
    processing_config = get_processing_config()
    sensor_config.sensor = args.sensors

    session_info = client.setup_session(sensor_config)

    pg_updater = PGUpdater(sensor_config, processing_config, session_info)
    pg_process = et.PGProcess(pg_updater)
    pg_process.start()

    client.start_session()

    interrupt_handler = et.utils.ExampleInterruptHandler()
    print("Press Ctrl-C to end session")

    processor = Processor(sensor_config, processing_config, session_info)

    while not interrupt_handler.got_signal:
        info, data = client.get_next()
        plot_data = processor.process(data, info)

        if plot_data is not None:
            try:
                pg_process.put_data(plot_data)
            except et.PGProccessDiedException:
                break

    print("Disconnecting...")
    pg_process.close()
    client.disconnect()

def get_sensor_config():
    config = et.configs.SparseServiceConfig()
    config.profile = et.configs.SparseServiceConfig.Profile.PROFILE_3
    config.sampling_mode = et.configs.SparseServiceConfig.SamplingMode.A
    config.range_interval = [0.24, 0.48]
    config.sweeps_per_frame = 64
    config.sweep_rate = 3e3
    config.hw_accelerated_average_samples = 60
    return config

class ProcessingConfiguration(et.configbase.ProcessingConfig):
    VERSION = 1

    show_data_plot = et.configbase.BoolParameter(
        label="Show data",
        default_value=True,
        updateable=True,
        order=0,
    )

    show_speed_plot = et.configbase.BoolParameter(
        label="Show speed on FFT y-axis",
        default_value=False,
        updateable=True,
        order=10,
    )

get_processing_config = ProcessingConfiguration

class Processor:
    def __init__(self, sensor_config, processing_config, session_info):
        pass

    def process(self, data, data_info=None):
        if data_info is None:
            warnings.warn(
                "To leave out data_info or set to None is deprecated",
                DeprecationWarning,
                stacklevel=2,
            )

        frame = data

        zero_mean_frame = frame - frame.mean(axis=0, keepdims=True)
        fft = np.fft.rfft(zero_mean_frame.T * np.hanning(frame.shape[0]), axis=1)
        abs_fft = np.abs(fft)

        return {
            "frame": frame,
            "abs_fft": abs_fft,
        }

class PGUpdater:
    def __init__(self, sensor_config, processing_config, session_info):
        self.processing_config = processing_config

        self.downsampling_factor = sensor_config.downsampling_factor
        self.sweeps_per_frame = sensor_config.sweeps_per_frame
        sweep_rate = session_info["sweep_rate"]
        self.depths = et.utils.get_range_depths(sensor_config, session_info)
        self.step_length = session_info["step_length_m"]
        self.f_res = sweep_rate / self.sweeps_per_frame
        self.fft_x_scale = 100 * self.step_length

        self.smooth_max = et.utils.SmoothMax(
            sweep_rate / self.sweeps_per_frame,
            tau_grow=0,
            tau_decay=0.5,
            hysteresis=0.1,
        )

        self.setup_is_done = False

    def setup(self, win):
        self.plots = []
        self.curves = []
        for i in range(self.depths.size):
            title = "{:.0f} cm".format(100 * self.depths[i])
            plot = win.addPlot(row=0, col=i, title=title)
            plot.setMenuEnabled(False)
            plot.setMouseEnabled(x=False, y=False)
            plot.hideButtons()
            plot.showGrid(x=True, y=True)
            plot.setYRange(0, 2 ** 16)
            plot.hideAxis("left")
            plot.hideAxis("bottom")
            plot.plot(np.arange(self.sweeps_per_frame), 2 ** 15 * np.ones(self.sweeps_per_frame))
            curve = plot.plot(pen=et.utils.pg_pen_cycler())
            self.plots.append(plot)
            self.curves.append(curve)

        self.ft_plot = win.addPlot(row=1, col=0, colspan=self.depths.size)
        self.ft_plot.setMenuEnabled(False)
        self.ft_plot.setMouseEnabled(x=False, y=False)
        self.ft_plot.hideButtons()
        self.ft_im = pg.ImageItem(autoDownsample=True)
        self.ft_im.setLookupTable(et.utils.pg_mpl_cmap("viridis"))
        self.ft_plot.addItem(self.ft_im)
        self.ft_plot.setLabel("bottom", "Depth (cm)")
        self.ft_plot.getAxis("bottom").setTickSpacing(6 * self.downsampling_factor, 6)

        self.setup_is_done = True
        self.update_processing_config()

    def update_processing_config(self, processing_config=None):
        if processing_config is None:
            processing_config = self.processing_config
        else:
            self.processing_config = processing_config

        if not self.setup_is_done:
            return

        for plot in self.plots:
            plot.setVisible(self.processing_config.show_data_plot)

        half_wavelength = 2.445e-3
        self.ft_im.resetTransform()
        self.ft_im.translate(100 * (self.depths[0] - self.step_length / 2), 0)
        if self.processing_config.show_speed_plot:
            self.ft_plot.setLabel("left", "Speed (m/s)")
            self.ft_im.scale(self.fft_x_scale, self.f_res * half_wavelength)
        else:
            self.ft_plot.setLabel("left", "Frequency (kHz)")
            self.ft_im.scale(self.fft_x_scale, self.f_res * 1e-3)

    def update(self, data):
        frame = data["frame"]

        for i, ys in enumerate(frame.T):
            self.curves[i].setData(ys)

        m = np.max(data["abs_fft"])
        m = max(m, 1e4)
        m = self.smooth_max.update(m)
        self.ft_im.updateImage(data["abs_fft"], levels=(0, m * 1.05))

if __name__ == "__main__":
    main()


from multiprocessing import Process, Queue
import logging
import time

import cv2

import numpy as np

import tensorflow as tf

from ..common import config
from ..common.eye_cropper import EyeCropper

logger = logging.getLogger(__name__)

layers = tf.keras.layers

def _is_stale(timestamp):
  """ Checks if an image is stale.
  Args:
    timestamp: The timestamp of the image to check.
  Returns:
    True if the image is stale, false otherwise. """
  if time.time() - timestamp > config.STALE_THRESHOLD:
    logger.warning("Dropping stale image from %f at %f." % \
                   (timestamp, time.time()))
    return True

  return False

class GazePredictor(object):
  """ Takes in eye images, and uses the model to predict gaze. """

  def __init__(self, model, model_file, phone, display=False, drop_stale=True):
    """
    Args:
      model: The model class to use.
      model_file: The saved model to load for predictions.
      phone: The configuration data for the phone we are using.
      display: If true, it will enable a debug display that shows the image
               crops.
      drop_stale: If true, we want to automatically drop stale images. """
    # Initialize landmark and prediction processes.
    self._prediction_process = _CnnProcess(model, model_file, drop_stale=drop_stale)
    self._landmark_process = _LandmarkProcess(self._prediction_process,
                                              phone,
                                              display=display,
                                              drop_stale=drop_stale)

  def __del__(self):
    # Make sure internal processes have terminated.
    self._landmark_process.release()
    self._prediction_process.release()

  def predict_gaze(self):
    """ Predicts the user's gaze based on current frames.
    Returns:
      The predicted gaze point, in cm, the sequence number of the
      corresponding frame, and the timestamp. """
    # Wait for new output from the predictor.
    return self._prediction_process.get_output()

  def process_image(self, image, seq_num):
    """ Adds a new image to the prediction pipeline.
    Args:
      image: The image to add.
      seq_num: The sequence number of the image. """
    # Add a new timestamp for this image. This allows us to drop frames that go
    # stale.
    timestamp = time.time()

    self._landmark_process.add_new_input(image, seq_num, timestamp)

class _CnnProcess(object):
  """ Runs the CNN prediction in a separate process on the GPU, so that it can
  be handled concurrently. """

  def __init__(self, model, model_file, drop_stale=True):
    """
    Args:
      model: The model class to use.
      model_file: The file to load the predictor model from.
      drop_stale: Whether to automatically drop stale images. """
    self.__model_file = model_file
    self.__model_class = model
    self.__drop_stale = drop_stale

    # Create the queues.
    self.__input_queue = Queue()
    self.__output_queue = Queue()

    # Fork the predictor process.
    self.__process = Process(target=_CnnProcess.predict_forever, args=(self,))
    self.__process.start()

  @staticmethod
  def _make_eye_pathway(eye_input):
    """ Generates the preprocessing pathway for a single eye.
    Args:
      eye_input: The raw eye input to be preprocessed.
    Returns:
      The input placeholder, and the output node. """
    # Convert to grayscale.
    eye_gray = tf.image.rgb_to_grayscale(eye_input)
    # Crop it properly.
    eye_cropped = tf.image.crop_to_bounding_box(eye_gray, 3, 3, 218, 218)
    # Resize it properly.
    resize_to = config.EYE_SHAPE[:2]
    eye_resized = tf.image.resize_images(eye_cropped, resize_to,
                                         align_corners=True)
    # Normalize it.
    eye_norm = tf.map_fn(lambda frame: \
                         tf.image.per_image_standardization(frame), eye_resized)

    return eye_norm

  @staticmethod
  def _make_face_pathway(face_input):
    """ Generates the preprocessing pathway for the face image.
    Args:
      face_input: The raw face input to be processed.
    Returns:
      The input placeholder, and the output node. """
    # Crop it properly.
    face_cropped = tf.image.crop_to_bounding_box(face_input, 11, 11, 202, 202)
    # Resize it properly.
    resize_to = config.FACE_SHAPE[:2]
    face_resized = tf.image.resize_images(face_cropped, resize_to,
                                          align_corners=True)

    # Normalize it.
    face_norm = tf.map_fn(lambda frame: \
                          tf.image.per_image_standardization(frame),
                          face_resized)

    return face_norm

  def release(self):
    """ Cleans up and terminates internal process. """
    self.__process.terminate()

  def predict_forever(self):
    """ Generates predictions indefinitely. """
    # Create preprocessing layers.
    eye_preproc_layer = layers.Lambda(_CnnProcess._make_eye_pathway)
    face_preproc_layer = layers.Lambda(_CnnProcess._make_face_pathway)

    # We expect color image inputs.
    color_eye_shape = config.EYE_SHAPE[:2] + (3,)
    # Load the model we trained.
    model = self.__model_class(config.FACE_SHAPE, eye_shape=color_eye_shape,
                               eye_preproc=eye_preproc_layer,
                               face_preproc=face_preproc_layer)
    self.__predictor = model.build()
    self.__predictor.load_weights(self.__model_file)

    while True:
      self.__predict_once()

  def __predict_once(self):
    """ Reads an image from the input queue, processes it, and writes a
    prediction to the output queue. """
    left_eye, right_eye, face, grid, seq_num, timestamp = \
        self.__input_queue.get()
    if seq_num is None:
      # A None tuple means the end of the sequence. Propagate this through the
      # pipeline.
      self.__output_queue.put((None, None, None))
      return
    if left_eye is None:
      # If we have a sequence number but no images, we got a bad detection.
      self.__output_queue.put((None, seq_num, timestamp))
      return
    if (self.__drop_stale and _is_stale(timestamp)):
      # The image is stale, so indicate that it is invalid.
      self.__output_queue.put((None, seq_num, timestamp))
      return

    # Convert everything to floats.
    left_eye = left_eye.astype(np.float32)
    right_eye = right_eye.astype(np.float32)
    face = face.astype(np.float32)

    # Add the batch dimension.
    left_eye = np.expand_dims(left_eye, axis=0)
    right_eye = np.expand_dims(right_eye, axis=0)
    face = np.expand_dims(face, axis=0)
    grid = np.expand_dims(grid, axis=0)

    # Generate a prediction.
    pred = self.__predictor.predict([left_eye, right_eye, face, grid],
                                    batch_size=1)
    # Remove the batch dimension.
    if type(pred) is list:
      pred = [x[0] for x in pred]
    else:
      pred = pred[0]

    self.__output_queue.put((pred, seq_num, timestamp))

  def add_new_input(self, left_eye, right_eye, face, grid, seq_num, timestamp):
    """ Adds a new input to be processed. Will block.
    Args:
      left_eye: The left eye crop.
      rigth_eye: The right eye crop.
      face: The face crop.
      grid: The face grid.
      seq_num: The sequence number of the image.
      timestamp: The timestamp of the image.
    """
    self.__input_queue.put((left_eye, right_eye, face, grid, seq_num,
                            timestamp))

  def get_output(self):
    """ Gets an output from the prediction process. Will block.
    Returns:
      The predicted gaze point and the sequence number. """
    return self.__output_queue.get()

class _LandmarkProcess(object):
  """ Reads images from a queue, and runs landmark detection in a separate
  process. """

  def __init__(self, cnn_process, phone, display=False, drop_stale=True):
    """
    Args:
      phone: The configuration of the phone that we are capturing data on.
      cnn_process: The _CnnProcess to send captured images to.
      display: If true, it will enable a debugging display that shows the
               detected crops on-screen.
      drop_stale: Whether to automatically drop stale frames. """
    self.__cnn_process = cnn_process
    self.__display = display
    self.__phone = phone
    self.__drop_stale = drop_stale

    # Create the queues.
    self.__input_queue = Queue()
    self.__output_queue = Queue()

    # Fork the capture process.
    self.__process = Process(target=_LandmarkProcess.run_forever,
                             args=(self,))
    self.__process.start()

  def release(self):
    """ Cleans up and terminates internal process. """
    self.__process.terminate()

  def __run_once(self):
    """ Reads and crops a single image. It will send it to the predictor
    process when finished. """
    # Get the next input from the queue.
    image, seq_num, timestamp = self.__input_queue.get()
    if image is None:
      # A None tuple means the end of a sequence. Propagate this through the
      # pipeline.
      self.__cnn_process.add_new_input(None, None, None, None, None, None)
      return
    if (self.__drop_stale and _is_stale(timestamp)):
      # Image is stale. Indicate that it is invalid.
      self.__cnn_process.add_new_input(None, None, None, None, seq_num,
                                       timestamp)

    # Crop the image.
    left_eye, right_eye, face = self.__cropper.crop_image(image)
    if face is None:
      # We failed to get an image.
      logger.warning("Failed to get good detection for %d." % (seq_num))
      # Send along the sequence number.
      self.__cnn_process.add_new_input(None, None, None, None, seq_num,
                                       timestamp)
      return

    # Produce face mask.
    mask = self.__cropper.face_grid()

    if self.__display:
      # Show the debugging display.
      mask_sized = cv2.resize(mask, (224, 224))
      mask_sized = np.expand_dims(mask_sized, axis=2)
      mask_sized = np.tile(mask_sized, [1, 1, 3]) * 225.0
      mask_sized = mask_sized.astype(np.uint8)

      combined = np.concatenate((left_eye, right_eye, face, mask_sized), axis=1)
      cv2.imshow("Server Detections", combined)
      cv2.waitKey(1)

    # Send it along.
    self.__cnn_process.add_new_input(left_eye, right_eye, face, mask,
                                     seq_num, timestamp)

  def run_forever(self):
    """ Reads and crops images indefinitely. """
    # Eye cropper to use for eye detection.
    self.__cropper = EyeCropper(self.__phone)

    while True:
      self.__run_once()

  def add_new_input(self, image, seq_num, timestamp):
    """ Adds a new input to be processed. Will block.
    Args:
      image: The image to process.
      seq_num: The sequence number of the image.
      timestamp: The timestamp of the image. """
    self.__input_queue.put((image, seq_num, timestamp))