import argparse
import time
import shutil
import os
import os.path as osp
import csv
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR
from data_cnn60 import NTUDataLoaders, AverageMeter, make_dir, get_cases, get_num_classes
from sklearn.metrics import confusion_matrix
from collections import OrderedDict
import torch.nn.functional as F
from alignment_module import AE, PosMmen, JPosMmen, MMDLoss, multi_class_hinge_loss, triplet_loss, Mmen
import pickle as pkl

parser = argparse.ArgumentParser(description='View adaptive')
parser.add_argument('--ss', type=int, help="split size")
parser.add_argument('--st', type=str, help="split type")
parser.add_argument('--dataset', type=str, help="dataset path")
parser.add_argument('--wdir', type=str, help="directory to save weights path")
parser.add_argument('--le', type=str, help="language embedding model")
parser.add_argument('--ve', type=str, help="visual embedding model")
parser.add_argument('--phase', type=str, help="train or val")
parser.add_argument('--gpu', type=str, help="gpu device number")
parser.add_argument('--ntu', type=int, help="no of classes")
parser.add_argument('--temp', type=int, help="temperature")
parser.add_argument('--thresh', type=float, help="threshold")

args = parser.parse_args()

gpu = args.gpu
ss = args.ss
st = args.st
dataset_path = args.dataset
wdir = args.wdir
le = args.le
ve = args.ve
phase = args.phase
num_classes = args.ntu
temp = args.temp
thresh = args.thresh

# gpu = '0'
# ss = 10
# st = 'r'
# dataset_path = 'ntu_results/shift_10_r'
# # wdir = gating_our
# le = 'bert'
# ve = 'shift'
# phase = 'train'
# num_classes = 120

os.environ["CUDA_VISIBLE_DEVICES"] = gpu
seed = 5
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
device = torch.device("cuda")
print(torch.cuda.device_count())

criterion2 = nn.MSELoss().to(device)

if ve == 'vacnn':
    vis_emb_input_size = 2048
elif ve == 'shift':
    vis_emb_input_size = 256
else:
    pass

text_hidden_size = 100
vis_hidden_size = 512
output_size = 50

if le == 'bert':
    noun_emb_input_size = 1024
    verb_emb_input_size = 1024
elif le == 'w2v':
    noun_emb_input_size = 300
    verb_emb_input_size = 300
else:
    pass

npm_checkpoint =  wdir + '/' + le + '/main_gvacnn_NPM.pth.tar'
vpm_checkpoint = wdir + '/' + le + '/main_gvacnn_VPM.pth.tar'
jm_checkpoint =  wdir + '/' + le + '/main_gvacnn_JM.pth.tar'
NounPosMmen = Mmen(vis_emb_input_size, noun_emb_input_size, output_size).to(device)
NounPosMmen.load_state_dict(torch.load(npm_checkpoint)['npm_state_dict'], strict=False)
NounPosMmen_optimizer = optim.Adam(NounPosMmen.parameters(), lr=0.0001)
NounPosMmen_scheduler = ReduceLROnPlateau(NounPosMmen_optimizer, mode='max', factor=0.1, patience=14, cooldown=6, verbose=True)

VerbPosMmen = Mmen(vis_emb_input_size, verb_emb_input_size, output_size).to(device)
VerbPosMmen.load_state_dict(torch.load(vpm_checkpoint)['vpm_state_dict'], strict=False)
VerbPosMmen_optimizer = optim.Adam(VerbPosMmen.parameters(), lr=0.0001)
VerbPosMmen_scheduler = ReduceLROnPlateau(VerbPosMmen_optimizer, mode='max', factor=0.1, patience=14, cooldown=6, verbose=True)

joint_emb_size = 100
joint_hidden_size = 50
JointMmen = JPosMmen(joint_emb_size, joint_hidden_size).to(device)
JointMmen.load_state_dict(torch.load(jm_checkpoint)['vpm_state_dict'], strict=False)
JointMmen_optimizer = optim.Adam(JointMmen.parameters(), lr=0.0001, weight_decay = 0.001)
JointMmen_scheduler = ReduceLROnPlateau(JointMmen_optimizer, mode='max', factor=0.1, patience=14, cooldown=6, verbose=True)

ntu_loaders = NTUDataLoaders(dataset_path, 'max', 1)
train_loader = ntu_loaders.get_train_loader(1024, 8)
zsl_loader = ntu_loaders.get_val_loader(1024, 8)
val_loader = ntu_loaders.get_test_loader(1024, 8)
zsl_out_loader = ntu_loaders.get_val_out_loader(1024, 8)
val_out_loader = ntu_loaders.get_test_out_loader(1024, 8)
train_size = ntu_loaders.get_train_size()
zsl_size = ntu_loaders.get_val_size()
val_size = ntu_loaders.get_test_size()
print('Train on %d samples, validate on %d samples' % (train_size, zsl_size))

nouns_vocab = np.load('../resources/nouns_vocab.npy')
verbs_vocab = np.load('../resources/verbs_vocab.npy')
nouns = nouns_vocab[np.argmax(np.load('../resources/nouns_ohe.npy'), -1)][:num_classes]
verbs = verbs_vocab[np.argmax(np.load('../resources/verbs_ohe.npy'), -1)][:num_classes]
labels = np.load('../resources/labels.npy')

if phase == 'val':
    gzsl_inds = np.load('../resources/label_splits/'+ st + 's' + str(num_classes - ss) +'.npy')
    unseen_inds = np.sort(np.load('../resources/label_splits/' + st + 'v' + str(ss) + '_0.npy'))
    seen_inds = np.load('../resources/label_splits/'+ st + 's' + str(num_classes - ss -ss) + '_0.npy')
else:
    gzsl_inds = np.arange(num_classes)
    unseen_inds = np.sort(np.load('../resources/label_splits/' + st + 'u' + str(ss) + '.npy'))
    seen_inds = np.load('../resources/label_splits/'+ st + 's' + str(num_classes - ss) + '.npy')

unseen_labels = labels[unseen_inds]
seen_labels = labels[seen_inds]

unseen_nouns = nouns[unseen_inds]
unseen_verbs = verbs[unseen_inds]
seen_nouns = nouns[seen_inds]
seen_verbs = verbs[seen_inds]
verb_corp = np.unique(verbs[gzsl_inds])
noun_corp = np.unique(nouns[gzsl_inds])

verb_emb = torch.from_numpy(np.load('../resources/ntu' + str(num_classes) + '_' + le + '_verb.npy')).view([num_classes, verb_emb_input_size])
verb_emb = verb_emb/torch.norm(verb_emb, dim = 1).view([num_classes, 1]).repeat([1, verb_emb_input_size])
noun_emb = torch.from_numpy(np.load('../resources/ntu' + str(num_classes) + '_' + le + '_noun.npy')).view([num_classes, noun_emb_input_size])
noun_emb = noun_emb/torch.norm(noun_emb, dim = 1).view([num_classes, 1]).repeat([1, noun_emb_input_size])

unseen_verb_emb = verb_emb[unseen_inds, :]
unseen_noun_emb = noun_emb[unseen_inds, :]
seen_verb_emb = verb_emb[seen_inds, :]
seen_noun_emb = noun_emb[seen_inds, :]
print("loaded language embeddings")

def get_text_data(target, verb_emb, noun_emb):
    return verb_emb[target].view(target.shape[0], verb_emb_input_size).float(), noun_emb[target].view(target.shape[0], verb_emb_input_size).float()

def save_checkpoint(state, filename='checkpoint.pth.tar', is_best=False):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

def accuracy(class_embedding, vis_trans_out, target, inds):
    inds = torch.from_numpy(inds).to(device)
    temp_vis = vis_trans_out.unsqueeze(1).expand(vis_trans_out.shape[0], class_embedding.shape[0], vis_trans_out.shape[1])
    temp_cemb = class_embedding.unsqueeze(0).expand(vis_trans_out.shape[0], class_embedding.shape[0], vis_trans_out.shape[1])
    preds = torch.argmax(torch.sum(temp_vis*temp_cemb, axis=2), axis = 1)
    acc = torch.sum(inds[preds] == target).item()/(preds.shape[0])
    return acc, torch.sum(temp_vis*temp_cemb, axis=2)

def ce_accuracy(output, target):
    batch_size = target.size(0)
    _, pred = output.topk(1, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))
    correct = correct.view(-1).float().sum(0, keepdim=True)
    return correct.mul_(100.0 / batch_size)

def zsl_validate(val_loader, epoch, margin):

    with torch.no_grad():
        losses = AverageMeter()
        acces = AverageMeter()
        ce_loss_vals = []
        ce_acc_vals = []
        # trunk.eval()
        NounPosMmen.eval()
        VerbPosMmen.eval()
        JointMmen.eval()
        # PrpPosMmen.eval()
        tars = []
        preds = []
        scores = []
        noun_scores = []
        verb_scores = []
        for i, (inputs, target) in enumerate(val_loader):
            emb = inputs.to(device)
            verb_emb_target, noun_emb_target = get_text_data(target, verb_emb, noun_emb)
            verb_tar = []
            for p in verbs[target]:
                verb_tar.append(np.argwhere(verb_corp == p)[0][0])

            noun_tar = []
            for p in np.unique(nouns[target]):
                noun_tar.append(np.argwhere(noun_corp == p)[0][0])

            emb = emb/torch.norm(emb, dim = 1).view([emb.size(0), 1]).repeat([1, emb.shape[1]])
            verb_embeddings = VerbPosMmen.TextArch(unseen_verb_emb.to(device).float())
            noun_embeddings = NounPosMmen.TextArch(unseen_noun_emb.to(device).float())
            noun_embeddings = noun_embeddings/torch.norm(noun_embeddings, dim = 1).view([noun_embeddings.size(0), 1]).repeat([1, noun_embeddings.shape[1]])
            verb_embeddings = verb_embeddings/torch.norm(verb_embeddings, dim = 1).view([verb_embeddings.size(0), 1]).repeat([1, verb_embeddings.shape[1]])
            joint_text_embedding = torch.cat([verb_embeddings, noun_embeddings], axis=1)
            fin_text_embedding = JointMmen.TextArch(joint_text_embedding)
            fin_text_embedding = fin_text_embedding/torch.norm(fin_text_embedding, dim = 1).view([fin_text_embedding.size(0), 1]).repeat([1, fin_text_embedding.shape[1]])

            vis_verb_transform, verb_transform = VerbPosMmen(emb, verb_emb_target.to(device).float())
            vis_verb_transform = vis_verb_transform/torch.norm(vis_verb_transform, dim = 1).view([vis_verb_transform.size(0), 1]).repeat([1, vis_verb_transform.shape[1]])
            verb_transform = verb_transform/torch.norm(verb_transform, dim = 1).view([verb_transform.size(0), 1]).repeat([1, verb_transform.shape[1]])
            verb_vis_loss = multi_class_hinge_loss(vis_verb_transform, verb_transform, target, verb_embeddings, margin).to(device)
            verb_verb_loss = triplet_loss(vis_verb_transform, torch.tensor(verb_tar), device, margin).to(device)

            vis_noun_transform,  noun_transform = NounPosMmen(emb, noun_emb_target.to(device).float())
            vis_noun_transform = vis_noun_transform/torch.norm(vis_noun_transform, dim = 1).view([vis_noun_transform.size(0), 1]).repeat([1, vis_noun_transform.shape[1]])
            noun_transform = noun_transform/torch.norm(noun_transform, dim = 1).view([noun_transform.size(0), 1]).repeat([1, noun_transform.shape[1]])
            noun_vis_loss = multi_class_hinge_loss(vis_noun_transform, noun_transform, target, noun_embeddings, margin).to(device)
            noun_noun_loss = triplet_loss(vis_noun_transform, torch.tensor(noun_tar), device, margin).to(device)

            joint_vis = torch.cat([vis_verb_transform,  vis_noun_transform], axis = 1)
            joint_text = torch.cat([verb_transform,  noun_transform], axis = 1)
            fin_vis, fin_text = JointMmen(joint_vis, joint_text)
            fin_text = fin_text/torch.norm(fin_text, dim = 1).view([fin_text.size(0), 1]).repeat([1, fin_text.shape[1]])
            fin_vis = fin_vis/torch.norm(fin_vis, dim = 1).view([fin_vis.size(0), 1]).repeat([1, fin_vis.shape[1]])
            joint_loss = multi_class_hinge_loss(fin_vis, fin_text, target, fin_text_embedding, margin).to(device)

            loss = verb_vis_loss + noun_vis_loss
            loss += 0.01*(verb_verb_loss) + 0.01*(noun_noun_loss)
            loss += joint_loss

            # ce acc
            ce_acc, score = accuracy(fin_text_embedding, fin_vis, target.to(device), unseen_inds)
            losses.update(loss.item(), inputs.size(0))
            acces.update(ce_acc, inputs.size(0))
            tars.append(target)
            scores.append(score)
            ce_loss_vals.append(loss.cpu().detach().numpy())
            ce_acc_vals.append(ce_acc)

            if i % 20 == 0:
                print('Epoch-{:<3d} {:3d}/{:3d} batches \t'
                    'loss {loss.val:.4f} ({loss.avg:.4f})\t'
                    'accu {acc.val:.3f} ({acc.avg:.3f})\t'.format(
                    epoch, i + 1, len(val_loader), loss=losses, acc=acces))
        return losses.avg, acces.avg, scores

def validate(train_loader, epoch, margin):

    losses = AverageMeter()
    acces = AverageMeter()
    ce_loss_vals = []
    ce_acc_vals = []
    scores = []
    for i, (inputs, target) in enumerate(train_loader):
        # (inputs, target) = next(iter(train_loader))
        emb = inputs.to(device)
        scores.append(emb)
    return scores

zsl_loss, zsl_acc, test_zs = zsl_validate(val_loader, 0, 0.3)
test_seen = np.load(dataset_path + '/gtest_out.npy')

tars = []
for i, (_, target) in enumerate(val_loader):
    tars += list(target.numpy())

test_y = []
for i in tars:
    if i in unseen_inds:
        test_y.append(0)
    else:
        test_y.append(1)

test_zs = np.array([j.cpu().detach().numpy() for i in test_zs for j in i])

def temp_scale(seen_features, T):
    return np.array([np.exp(i)/np.sum(np.exp(i)) for i in (seen_features + 1e-12)/T])

prob_test_zs = np.array([np.exp(i)/np.sum(np.exp(i)) for i in test_zs])
prob_test_seen = temp_scale(test_seen, temp)
main_prob_test_seen = np.array([np.exp(i)/np.sum(np.exp(i)) for i in test_seen])

feat_test_zs = np.sort(prob_test_zs, 1)[:,::-1][:,:ss]
feat_test_seen = np.sort(prob_test_seen, 1)[:,::-1][:,:ss]

gating_test_x = np.concatenate([feat_test_zs, feat_test_seen], 1)
gating_test_y = test_y
with open(wdir + '/' + le + '/gating_model.pkl', 'rb') as f:
    gating_model = pkl.load(f)

prob_gate = gating_model.predict_proba(gating_test_x)
pred_test = 1 - prob_gate[:, 0]>thresh

a = prob_gate
b = np.zeros(prob_gate.shape[0])
p_gate_seen = prob_gate[:, 1]
prob_y_given_seen = prob_test_seen + (1/num_classes)*np.repeat((1 - p_gate_seen)[:, np.newaxis], num_classes, 1)

p_gate_unseen = prob_gate[:, 0]
prob_y_given_unseen = prob_test_zs + (1/ss)*np.repeat((1 - p_gate_unseen)[:, np.newaxis], ss, 1)

prob_seen = prob_y_given_seen*np.repeat(p_gate_seen[:, np.newaxis], num_classes, 1)
prob_unseen = prob_y_given_unseen*np.repeat(p_gate_unseen[:, np.newaxis], ss, 1)

final_preds = []
seen_count = 0
tot_seen = 0
unseen_count = 0
tot_unseen = 0
gseen_count = 0
gunseen_count = 0
for i in range(len(prob_seen)):
    if pred_test[i] == 1:
        pred = np.argmax(main_prob_test_seen[i, :])
    else:
        pred = unseen_inds[np.argmax(prob_test_zs[i, :])]

    if tars[i] in seen_inds:
        tot_seen += 1
        if pred_test[i] == 1:
            gseen_count += 1
        if pred == tars[i]:
            seen_count += 1
    else:
        if pred_test[i] == 0:
            gunseen_count+=1
        tot_unseen += 1
        if pred == tars[i]:
            unseen_count += 1
    final_preds.append(pred)

seen_acc = seen_count/tot_seen
unseen_acc = unseen_count/tot_unseen
h_mean = 2*seen_acc*unseen_acc/(seen_acc + unseen_acc)
print('seen accuracy', seen_acc)
print('unseen accuracy', unseen_acc)
print('harmonic mean', h_mean)


# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
train_classifier
"""
import os
import argparse
import numpy as np
import mindspore.dataset as ds
from mindspore import nn
from mindspore import save_checkpoint
from mindspore.nn import piecewise_constant_lr
from mindspore import context
from mindspore import ParameterTuple
from mindspore import ops
import src.util as util
from src.data.IterSamplers import CategoriesSampler
from src.data.mini_Imagenet import MiniImageNet
from src.model.classifier import Classifier, ClassifierWithLossCell
from src.model.meta_eval import MetaEval
from tqdm import tqdm

class TrainOneStepCell(nn.Cell):
    """
    TrainOneStepCell
    """

    def __init__(self, network, optimizer, sens=1.0):
        super(TrainOneStepCell, self).__init__(auto_prefix=True)
        self.network = network
        self.optimizer = optimizer
        self.weights = ParameterTuple(network.trainable_params())
        self.grad = ops.GradOperation(get_by_list=True, sens_param=True)
        self.sens = sens

    def set_sens(self, value):
        """
        :param value:
        :return:
        """
        self.sens = value

    def construct(self, data, label):
        """
        :param data:
        :param label:
        :return:
        """
        weights = self.weights
        loss = self.network(data, label)
        sens = ops.Fill()(ops.DType()(loss), ops.Shape()(loss), self.sens)
        grads = self.grad(self.network, weights)(data, label, sens)
        return ops.depend(loss, self.optimizer(grads))

def main():
    """
    train
    :return:
    """
    util.ensure_path(save_path)
    # train set
    root_path = os.path.join(args.root_path, args.dataset)
    n_way, n_query, n_shots = 5, 15, [1, 5]
    if args.dataset == 'mini-imagenet':
        trainset = MiniImageNet(root_path, 'train')
        testset = MiniImageNet(root_path, 'test')
    else:
        print('not found error')

    trainloader = ds.GeneratorDataset(trainset, ['data', 'label'], shuffle=True).batch(
        args.batch_size)
    # test set
    fs_loaders = []
    for n_shot in n_shots:
        fs_loader = CategoriesSampler(testset.data, testset.label, n_way, n_shot + n_query,
                                      args.ep_batch, args.ep_per_batch)
        fs_loaders.append(ds.GeneratorDataset(fs_loader, ['data'], shuffle=True))

    # define network
    net = Classifier(args.n_classes)
    net_with_loss = ClassifierWithLossCell(net)
    # define opt
    train_batch = trainset.len // args.batch_size
    multiStepLR = piecewise_constant_lr([(args.max_epoch - 10) * train_batch,
                                         args.max_epoch * train_batch], [args.lr, args.lr * 0.1])

    net_opt = nn.SGD(params=net.trainable_params(), learning_rate=multiStepLR,
                     weight_decay=args.weight_decay, momentum=0.9, nesterov=True)
    train_cell = TrainOneStepCell(net_with_loss, net_opt)

    eval_loss_fn = MetaEval()
    timer_used = util.Timer()
    timer_epoch = util.Timer()
    max_va = 0.

    for epoch in range(1, args.max_epoch + 1):
        timer_epoch.s()
        aves_keys = ['tl', 'ta']
        for n_shot in n_shots:
            aves_keys += ['fsa-' + str(n_shot)]
        aves = {k: util.Averager() for k in aves_keys}
        # pre train
        net.set_train(mode=True)
        for data in tqdm(trainloader.create_dict_iterator(), desc='train', leave=False):
            loss = train_cell(data['data'], data['label'])
            acc = net_with_loss.acc
            aves['tl'].add(loss.asnumpy())
            aves['ta'].add(acc.asnumpy())
        # few-shot eval
        if epoch == args.max_epoch or epoch % args.save_epoch == 0:
            net.set_train(mode=False)
            for i, n_shot in enumerate(n_shots):
                np.random.seed(0)
                for data in tqdm(fs_loaders[i].create_dict_iterator(), desc='test-' + str(n_shot),
                                 leave=False):
                    x_shot, x_query = data['data'][:, :, :n_shot], data['data'][:, :, n_shot:]
                    img_shape = x_query.shape[-3:]
                    x_query = x_query.view(args.ep_per_batch, -1, *img_shape)
                    label = util.make_nk_label(n_way, n_query, args.ep_per_batch)
                    acc_val, _ = eval_loss_fn.eval(x_shot, x_query, label, net.encoder)
                    aves['fsa-' + str(n_shot)].add(acc_val.asnumpy())
        # post
        for k, v in aves.items():
            aves[k] = v.item()
        t_epoch = util.time_str(timer_epoch.t())
        t_used = util.time_str(timer_used.t())
        t_estimate = util.time_str(timer_used.t() / epoch * args.max_epoch)

        print("epoch {},train loss {:.4f}, train acc {:.4f}".format(str(epoch), aves['tl'],
                                                                    aves['ta']))
        if epoch == args.max_epoch or epoch % args.save_epoch == 0:
            for n_shot in n_shots:
                key = 'fsa-' + str(n_shot)
                print("epoch {}, {}-shot, val acc {:.4f}".format(str(epoch), n_shot, aves[key]))

        if epoch <= args.max_epoch:
            print("{} {}/{}".format(t_epoch, t_used, t_estimate))
        else:
            print("{}".format(t_epoch))

        path = os.path.join(save_path, 'epoch-{}.ckpt'.format(epoch))
        if epoch >= 15 and epoch % args.save_epoch == 0:
            save_checkpoint(net, path)
        if aves['fsa-' + str(5)] > max_va:
            max_va = aves['fsa-' + str(5)]
            save_checkpoint(net, os.path.join(save_path, 'max-va.ckpt'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--name', default=None)
    parser.add_argument('--root_path', default='./dataset/')
    parser.add_argument('--tag', default=None)
    parser.add_argument('--device_target', type=str, default='GPU', choices=['Ascend', 'GPU', 'CPU'])
    parser.add_argument('--dataset', default='mini-imagenet')
    parser.add_argument('--encoder', default='resnet12')
    parser.add_argument('--load_encoder',
                        default='./save/classifier3_mini-imagenet_resnet12/epoch-70.ckpt')
    parser.add_argument('--ep_per_batch', type=int, default=4)
    parser.add_argument('--ep_batch', type=int, default=200)
    parser.add_argument('--max_epoch', type=int, default=25)
    parser.add_argument('--visualize_datasets', default=True)
    parser.add_argument('--lr', type=float, default=0.1)
    parser.add_argument('--weight_decay', type=float, default=5.e-4)
    parser.add_argument('--device_id', type=int, default=0)
    parser.add_argument('--num_gpu', type=int, default=1)
    parser.add_argument('--classifier', default='linear-classifier')
    parser.add_argument('--n_classes', type=int, default=64)
    parser.add_argument('--batch_size', type=int, default=128)
    parser.add_argument('--save_epoch', type=int, default=1)
    parser.add_argument('--eval_fs_epoch', type=int, default=3)
    parser.add_argument('--data_url', default=None, help='Location of data.')
    parser.add_argument('--train_url', default=None, help='Location of training outputs.')
    parser.add_argument('--run_offline', type=str, default="False", help='run in offline')
    args = parser.parse_args()

    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target, save_graphs=False)
    if args.device_target == 'GPU' or args.device_target == 'Ascend':
        context.set_context(device_id=args.device_id)
        if args.run_offline == "True":
            print("run_online--")
            import moxing as mox
            mox.file.copy_parallel(src_url=args.data_url, dst_url=args.root_path)
    else:
        raise ValueError("Unsupported platform.")

    svname = 'classifier_{}'.format(args.dataset)
    svname += '_' + args.encoder
    save_path = os.path.join('./save/', svname)
    main()

    if args.run_offline == "True":
        md_save_path = os.path.join(args.train_url, save_path)
        mox.file.make_dirs(md_save_path)
        mox.file.copy_parallel(src_url=save_path, dst_url=md_save_path)


import cv2
import numpy as np
import HandTrackingModule as htm
import time
import autopy

wCam, hCam = 640, 480
#fps Variable
pTime=0
#frame Size reduction Variable
frameR = 100
#Fluiding variables
fluiding = 5
plocX, plocY = 0, 0
clocX ,clocY = 0, 0

cap = cv2.VideoCapture(0)
cap.set(3,wCam)
cap.set(4,hCam)
detector = htm.handDetector(maxHands=1)
wScr, hScr = autopy.screen.size()

while True:
    # step-1 : Find hand landmarks
    success, img = cap.read()
    img=detector.findHands(img)
    lmList, bbox = detector.findPosition(img)

    # step-2 : Find tip of index finger and middle finger
    #Index mode will move the finger
    if(len(lmList)!=0):
        x1,y1 = lmList[8][1:]
        x2,y2 = lmList[12][1:]

        # step-3 : Check which fingers are up
        fingers = detector.fingersUp()
            #print(fingers)

        # step-4 : Only index Finger : Moving Mode
        cv2.rectangle(img,(frameR,frameR), (wCam-frameR , hCam-frameR), (255,0,255), 2)

        if fingers[1]==1 and fingers[2]==0:

            # Convert Coordinates to get the correct coordinates
            x3 = np.interp(x1, (frameR,wCam-frameR),(0,wScr))
            y3 = np.interp(y1, (frameR,hCam-frameR),(0,hScr))

            # step-5 : Making in Fluid(values)
            clocX = plocX + (x3- plocX) / fluiding
            clocY = plocY + (y3- plocY) / fluiding

            # step-6 : Move Mouse
            autopy.mouse.move(wScr-clocX,clocY)
            plocX, plocY = clocX, clocY
            cv2.circle(img,(x1,y1),15,(255,0,255),cv2.FILLED)

        # step-7 : For Verifing the Mode of mouse
        if fingers[1]==1 and fingers[2]==1:
            #if in clicking mode - find distace b/w fingers
            length, img, lineInfo = detector.findDistance(8, 12, img)
            #print(length)
            if length < 35:
                cv2.circle(img,(lineInfo[4],lineInfo[5]),15,(0,255,0),cv2.FILLED)

                #if distance is short - click mouse
                autopy.mouse.click()

    # step-8 : fix frame rate
    cTime = time.time()
    fps = 1 / (cTime - pTime)
    pTime = cTime
    cv2.putText(img, str(int(fps)), (10, 70), cv2.FONT_HERSHEY_PLAIN, 3,(255, 0, 255), 3)

    # step-9 : display
    cv2.imshow("Image", img)
    cv2.waitKey(1)


import argparse
import string
import os
import time
from multiprocessing import Pool
from itertools import repeat
from scipy.stats import bernoulli
import pandas as pd
import numpy as np
import sys
from tqdm import tqdm

np.random.seed(0)

def parse_args():
    parser = argparse.ArgumentParser(description='Reverse substitution')
    parser.add_argument('-i', '--input_file', required=True, help='path to file containing text')
    parser.add_argument('-n', '--num_of_rows', type=int, help='number of rows to be read from input text file')
    parser.add_argument('-a', '--abbreviation_table', required=True, help='path to abbreviation table')
    parser.add_argument('-o', '--output_file', required=True, help='path to output file')
    parser.add_argument('-p', '--probability', default=0.15, type=float, help='probability of reverse substitution')
    parser.add_argument('--max_cores', default=4, type=int, help='max number of cpu cores')
    return parser.parse_args()

def clean_expansion(text):
    lower = text.lower()
    no_punc = lower.translate(str.maketrans("", "", string.punctuation))
    no_white = " ".join(no_punc.split())
    no_num = no_white.translate(str.maketrans("", "", string.digits))
    return no_num

def reverse_substitute(input_tuple):
    text, expansions_to_ab, p = input_tuple
    expansions = [expans for expans in expansions_to_ab.keys() if expans in text]
    if not expansions:  # no expansion found
        return None
    wcs = [len(expansion.split()) for expansion in expansions]
    first_words = [expansion.split()[0] for expansion in expansions]
    max_wc = np.max(wcs)
    labels = {}
    current_text = text.split()
    for idx, word in enumerate(current_text):
        if word not in first_words: # the word can be the first word
            continue
        if np.sum([' '.join(current_text[idx:idx + wc]) in expansions for wc in range(max_wc, 0, -1)]) == 0: # no expansion detected here
            continue
        current_wcs = [wc for wc in range(max_wc, 0, -1) if ' '.join(current_text[idx:idx + wc]) in expansions] # word counts for all possible expansions
        for wc in current_wcs: # starting with longest
            if wc > len(current_text) - idx: # if remaining text is not long enough
                continue
            if bernoulli.rvs(p, size=1):
                labels[idx] = ' '.join(current_text[idx:idx + wc]) # record location
                current_text[idx] = expansions_to_ab[' '.join(current_text[idx:idx + wc])] # replace with abb
                for _ in range(1, wc):
                    del current_text[idx + 1] # delete remaining expansion
                break # don't search for shorter expansions
    if not labels:
        return None
    return ' '.join(current_text), labels

if __name__ == "__main__":
    args = parse_args()
    print('[' + time.ctime() + ']', 'args parsed')

    if args.num_of_rows:
        text_df = pd.read_csv(args.input_file, header=None, nrows=args.num_of_rows, engine='c')
        print('[' + time.ctime() + ']', args.num_of_rows, 'rows read from', args.input_file)
    else:
        text_df = pd.read_csv(args.input_file, header=None, engine='c')
        print('[' + time.ctime() + ']', 'all rows read from', args.input_file)
    text_df.columns = ['TEXT']

    adam_valid_df = pd.read_csv(args.abbreviation_table, sep='\t')
    print('[' + time.ctime() + ']', 'read abbreviation table from', args.abbreviation_table)

    ab_to_expansions = {ab: adam_valid_df['EXPANSION'][adam_valid_df['PREFERRED_AB'] == ab].to_list() for ab in adam_valid_df['PREFERRED_AB'].unique()}
    expansions_to_ab = {row['EXPANSION']: row['PREFERRED_AB'] for _, row in adam_valid_df.iterrows()}
    print('[' + time.ctime() + ']', 'abbreviation-expansion mapping constructed')

    with Pool(processes=args.max_cores) as pool:
        output_tuples = pool.map(reverse_substitute, zip(text_df['TEXT'], repeat(expansions_to_ab), repeat(args.probability)))
    print('[' + time.ctime() + ']', 'reverse substitute complete')

    abs_id = 0
    abs_ids = []
    rs_texts = []
    locations = []
    true_expansions = []
    print('[' + time.ctime() + ']', 'start parsing outputs')
    for output in tqdm(output_tuples):
        if not output:
            continue
        labels = output[1]
        labels = list(sorted(labels.items()))
        abs_id += 1
        for label in labels:
            abs_ids.append(abs_id)
            rs_texts.append(output[0])
            locations.append(label[0])
            true_expansions.append(label[1])
    output_df = pd.DataFrame({'ABSTRACT_ID': abs_ids, 'TEXT': rs_texts, 'LOCATION': locations, 'LABEL': true_expansions})
    output_df.to_csv(args.output_file, index=None)
    print('[' + time.ctime() + ']', 'output saved to', args.output_file)


from OpenGL.GL import *
from OpenGL.GLU import *
import numpy as np
import pygame

class Background():
    def __init__(self,game, pid, width, height, num):
        self.game=game
        self.pid=pid
        self.kind='Background'
        self.width=width
        self.height=height
        self.num=num
        self.r=1
        self.colour=(127,255,212,255)
        self.sprite=[]
        self.sprite_size=[]
        for i in [0,1,2,3,2,1,0]:
            r=self.r+i
            temp_sprite=pygame.Surface((2*r, 2*r), flags=pygame.SRCALPHA)
            temp_sprite.fill((255,255,255,0))
            temp_sprite_size=temp_sprite.get_size()
            pygame.draw.circle(temp_sprite,self.colour,(r,r),r)
            self.sprite.append(temp_sprite)
            self.sprite_size.append(temp_sprite_size)
        self.num_sprites=len(self.sprite)
        self.zorder=-9999
        self.x=(np.random.rand(self.num)-0.5)*self.width
        self.y=(np.random.rand(self.num)-0.5)*self.height
        self.sid=np.random.randint(0,3, self.num)

    def update(self):
        pass

    def draw(self):
        for star in range(self.num):
            """
            if np.random.rand() < 0.1:
                self.sid[star]+=1
            if self.sid[star] >= self.num_sprites:
                self.sid[star]=0
            dx=self.x[star]-self.game.camera.x
            dy=self.y[star]-self.game.camera.y
            self.game.screen.blit(self.sprite[self.sid[star]], (
                dx*self.game.camera.cangle+dy*self.game.camera.sangle+self.game.screen.width/2-self.sprite_size[self.sid[star]][0]/2+self.game.camera.x_shift,
                -dx*self.game.camera.sangle+dy*self.game.camera.cangle+self.game.screen.height/2-self.sprite_size[self.sid[star]][1]/2+self.game.camera.y_shift
                )
            )
            """
            glColor3f(1, 1, 1);
            glBegin(GL_QUADS)
            glVertex3fv((self.x[star]-5,self.y[star]-5,0))
            glVertex3fv((self.x[star]-5,self.y[star]+5,0))
            glVertex3fv((self.x[star]+5,self.y[star]+5,0))
            glVertex3fv((self.x[star]+5,self.y[star]-5,0))
            glEnd()

    def set_pid(self,pid):
        self.pid=pid


import requests
import _ssl
import ssl
from io import StringIO  # Python3 use: from io import StringIO
import sys
# Direct output to terminal tab
sys.stdout = sys.stderr = mystdout = StringIO()

from matplotlib.pyplot import style
from numpy import indices
from CoinbaseTrade import Coin
from NiceHashTrade import Nice
from window import Ui_MainWindow
from threading import Thread
import sys
from time import time, sleep
from PyQt5.QtCore import pyqtSignal, QObject, QThread, QEvent, Qt
from PyQt5 import QtGui
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc
import random

from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure
import matplotlib.pyplot as plt

marketsNice = []
changePurchaseNice = []
holdingNice = []
changePriceNice = []
spreadNice = []
transactionsNice = []

marketsCoin = []
changePurchaseCoin = []
holdingCoin = []
changePriceCoin = []
spreadCoin = []
transactionsCoin = []

class niceChecking(QObject):
	progress = pyqtSignal()
	def run(self):
		timeNow = (time() // checkingTime) + 1
		try:
			Nice.trade(timeNow, marketsNice, changePurchaseNice, holdingNice, changePriceNice, spreadNice, transactionsNice, True)
		except Exception as inst:
			print('Error when initializing tables and graphs')
			print(inst)
		self.progress.emit()
		while(True):
			sleep(0.1)
			if time() // checkingTime >= timeNow:
				print("Check NiceHash")
				try:
					Nice.trade(timeNow, marketsNice, changePurchaseNice, holdingNice, changePriceNice, spreadNice, transactionsNice)
				except Exception as inst:
					print('Error NiceHash')
					print(inst)
				timeNow = (time() // checkingTime) + 1
				self.progress.emit()

class coinChecking(QObject):
	progress = pyqtSignal()
	def run(self):
		timeNow = (time() // checkingTime) + 1
		try:
			Coin.trade(timeNow, marketsCoin, changePurchaseCoin, holdingCoin, changePriceCoin, spreadCoin, transactionsCoin, True)
		except Exception as inst:
			print('Error when initializing tables and graphs')
			print(inst)
		self.progress.emit()
		while(True):
			sleep(0.1)
			if time() // checkingTime >= timeNow:
				print("Check Coinbase")
				try:
					Coin.trade(timeNow, marketsCoin, changePurchaseCoin, holdingCoin, changePriceCoin, spreadCoin, transactionsCoin)
				except Exception as inst:
					print('Error Coinbase')
					print(inst)
				timeNow = (time() // checkingTime) + 1
				self.progress.emit()

class updateUI(QObject):
	progress = pyqtSignal()
	def run(self):
		while(True):
			sleep(0.1)
			self.progress.emit()

checkingTime = 900  # Check prices and trade if possible every minute

print("modules were imported")

class controllerWindow(qtw.QMainWindow):
	def __init__(self, *args, **kwargs):
		super().__init__(*args, **kwargs)

		self.ui = Ui_MainWindow()
		self.ui.setupUi(self)

		style = self.style()

		self.tray_icon = qtw.QSystemTrayIcon()
		self.tray_icon.setIcon(QtGui.QIcon('btc-logo.png'))
		self.setWindowIcon(QtGui.QIcon('btc-logo.png'))

		# Restore the window when the tray icon is double clicked.
		self.tray_icon.activated.connect(self.restore_window)

		self.ui.TableData.cellClicked.connect(self.marketClickedNice)
		self.ui.TableData_2.cellClicked.connect(self.marketClickedCoin)

		# Set plots for NiceHash
		self.niceMarket = plt.figure()
		self.niceBase = plt.figure()
		self.ui.graphMarket = FigureCanvas(self.niceMarket)
		self.ui.graphBase = FigureCanvas(self.niceBase)
		self.toolbar = NavigationToolbar(self.ui.graphMarket, self.ui.graphMarket)
		self.toolbar2 = NavigationToolbar(self.ui.graphBase, self.ui.graphBase)

		self.ui.ImagesSplitter.addWidget(self.ui.graphMarket)
		self.ui.ImagesSplitter.addWidget(self.ui.graphBase)

		self.baseNice = 'ETHBTC'

		self.niceMarket.clear()
		# create an axis
		ax = self.niceMarket.add_subplot(111)
		# plot data
		Nice.plotBase(self.baseNice, ax)
		# refresh canvas
		self.ui.graphMarket.draw()

		# Set plots for Coinbase
		self.coinMarket = plt.figure()
		self.coinBase = plt.figure()
		self.ui.graphMarket_2 = FigureCanvas(self.coinMarket)
		self.ui.graphBase_2 = FigureCanvas(self.coinBase)
		self.toolbar_2 = NavigationToolbar(self.ui.graphMarket_2, self.ui.graphMarket_2)
		self.toolbar2_2 = NavigationToolbar(self.ui.graphBase_2, self.ui.graphBase_2)

		self.ui.ImagesSplitter_2.addWidget(self.ui.graphMarket_2)
		self.ui.ImagesSplitter_2.addWidget(self.ui.graphBase_2)

		self.baseCoin = 'BTC-EUR'

		self.coinMarket.clear()
		# create an axis
		ax = self.coinMarket.add_subplot(111)
		# plot data
		Coin.plotBase(self.baseCoin, ax)
		# refresh canvas
		self.ui.graphMarket_2.draw()

		# Start threads
		self.threadNice = QThread()
		self.niceChecker = niceChecking()
		self.niceChecker.moveToThread(self.threadNice)
		self.threadNice.started.connect(self.niceChecker.run)
		self.threadNice.finished.connect(self.threadNice.deleteLater)
		self.niceChecker.progress.connect(self.niceTable)
		self.threadNice.start()

		self.threadCoin = QThread()
		self.coinChecker = coinChecking()
		self.coinChecker.moveToThread(self.threadCoin)
		self.threadCoin.started.connect(self.coinChecker.run)
		self.threadCoin.finished.connect(self.threadCoin.deleteLater)
		self.coinChecker.progress.connect(self.coinTable)
		self.threadCoin.start()

		self.thread2 = QThread()
		self.updateUi = updateUI()
		self.updateUi.moveToThread(self.thread2)
		self.thread2.started.connect(self.updateUi.run)
		self.thread2.finished.connect(self.thread2.deleteLater)
		self.updateUi.progress.connect(self.updateImage)
		self.thread2.start()

	def updatePlotsNice(self):
		self.niceMarket.clear()
		# create an axis
		ax = self.niceMarket.add_subplot(111)
		# plot data
		Nice.plotBase(self.baseNice, ax)
		# refresh canvas
		self.ui.graphMarket.draw()

	def updatePlotsCoin(self):
		self.coinMarket.clear()
		# create an axis
		ax = self.coinMarket.add_subplot(111)
		# plot data
		Coin.plotBase(self.baseCoin, ax)
		# refresh canvas
		self.ui.graphMarket_2.draw()

	def changeEvent(self, event):
		if event.type() == QEvent.WindowStateChange:
			if event.oldState() and Qt.WindowMinimized:
				self.hide()
				self.tray_icon.show()
			elif event.oldState() == Qt.WindowNoState or self.windowState() == Qt.WindowMaximized:
				print("WindowMaximized")

	def restore_window(self, reason):
		if reason == qtw.QSystemTrayIcon.DoubleClick:
			# self.showNormal will restore the window even if it was
			# minimized.
			self.show()
			self.showMaximized()

	def marketClickedNice(self):
		# Read selected row
		self.indices = self.ui.TableData.selectionModel().selectedRows()
		for index in sorted(self.indices):
			print('Row %d is selected' % index.row())
			print(self.ui.TableData.item(index.row(), 0).text())
			self.baseNice = self.ui.TableData.item(index.row(), 0).text()
			self.updatePlotsNice()

	def marketClickedCoin(self):
		# Read selected row
		self.indices = self.ui.TableData_2.selectionModel().selectedRows()
		for index in sorted(self.indices):
			print('Row %d is selected' % index.row())
			print(self.ui.TableData_2.item(index.row(), 0).text())
			self.baseCoin = self.ui.TableData_2.item(index.row(), 0).text()
			self.updatePlotsCoin()

	def updateImage(self):
		# Update graphs NiceHash
		# self.graph1 = QtGui.QPixmap(pictureNice1)
		# self.graph1Scaled = self.graph1.scaled(
		# 	self.ui.graphMarket.width(), self.ui.graphMarket.height(), qtc.Qt.AspectRatioMode.KeepAspectRatio)

		# self.graph2 = QtGui.QPixmap(pictureNice2)
		# self.graph2Scaled = self.graph2.scaled(self.ui.graphBase.width(
		# ), self.ui.graphBase.height(), qtc.Qt.AspectRatioMode.KeepAspectRatio)

		# self.ui.graphMarket.setPixmap(self.graph1Scaled)
		# self.ui.graphBase.setPixmap(self.graph2Scaled)

		# # Update graphs Coinbase
		# self.graph1_2 = QtGui.QPixmap(pictureCoin1)
		# self.graph1Scaled_2 = self.graph1_2.scaled(
		# 	self.ui.graphMarket_2.width(), self.ui.graphMarket_2.height(), qtc.Qt.AspectRatioMode.KeepAspectRatio)

		# self.graph2_2 = QtGui.QPixmap(pictureCoin2)
		# self.graph2Scaled_2 = self.graph2_2.scaled(self.ui.graphBase_2.width(
		# ), self.ui.graphBase_2.height(), qtc.Qt.AspectRatioMode.KeepAspectRatio)

		# self.ui.graphMarket_2.setPixmap(self.graph1Scaled_2)
		# self.ui.graphBase_2.setPixmap(self.graph2Scaled_2)

		#Update transactions text
		for i in range(len(transactionsNice)):
			self.ui.transactions.setText(transactionsNice[i])
		for i in range(len(transactionsCoin)):
			self.ui.transactions_2.setText(transactionsCoin[i])

		#Update terminal text
		self.ui.terminalText.setText(mystdout.getvalue())

	def niceTable(self):
		# Update plots NiceHash
		self.updatePlotsNice()

		# Update table NiceHash
		self.ui.TableData.clearContents()
		self.ui.TableData.setRowCount(0)
		self.ui.TableData.setRowCount(len(marketsNice))
		self.ui.TableData.setSortingEnabled(False)
		for i in range(len(marketsNice)):
			self.ui.TableData.setItem(i, 0, qtw.QTableWidgetItem(marketsNice[i]))

			self.ui.TableData.setItem(
							i, 1, qtw.QTableWidgetItem(str(holdingNice[i])))
			if holdingNice[i]:
				self.ui.TableData.item(i, 1).setBackground(QtGui.QColorConstants.Yellow)

			self.ui.TableData.setItem(i, 2, qtw.QTableWidgetItem(
				str("{:.2f}".format(changePurchaseNice[i]))+'%'))
			if(changePurchaseNice[i] > 5):
				self.ui.TableData.item(i, 2).setBackground(QtGui.QColorConstants.Green)
			elif (changePurchaseNice[i] < -4):
				self.ui.TableData.item(i, 2).setBackground(QtGui.QColorConstants.Red)

			self.ui.TableData.setItem(i, 3, qtw.QTableWidgetItem(
				str("{:.2f}".format(changePriceNice[i]))+'%'))
			if(changePriceNice[i] > 0):
				self.ui.TableData.item(i, 3).setBackground(QtGui.QColorConstants.Green)
			elif (changePriceNice[i] < 0):
				self.ui.TableData.item(i, 3).setBackground(QtGui.QColorConstants.Red)

			self.ui.TableData.setItem(i, 4, qtw.QTableWidgetItem(
				str("{:.2f}".format(spreadNice[i]))+'%'))
			if(spreadNice[i] <= 0.5):
				self.ui.TableData.item(i, 4).setBackground(QtGui.QColorConstants.Green)
			else:
				self.ui.TableData.item(i, 4).setBackground(QtGui.QColorConstants.Red)
		self.ui.TableData.setSortingEnabled(True)

	def coinTable(self):
		# Update plots Coinbase
		self.updatePlotsCoin()

		# Update table Coinbase
		self.ui.TableData_2.clearContents()
		self.ui.TableData_2.setRowCount(0)
		self.ui.TableData_2.setRowCount(len(marketsCoin))
		self.ui.TableData_2.setSortingEnabled(False)
		for i in range(len(marketsCoin)):
			self.ui.TableData_2.setItem(i, 0, qtw.QTableWidgetItem(marketsCoin[i]))

			self.ui.TableData_2.setItem(
                            i, 1, qtw.QTableWidgetItem(str(holdingCoin[i])))
			if holdingCoin[i]:
				self.ui.TableData_2.item(i, 1).setBackground(QtGui.QColorConstants.Yellow)

			self.ui.TableData_2.setItem(i, 2, qtw.QTableWidgetItem(
				str("{:.2f}".format(changePurchaseCoin[i]))+'%'))
			if(changePurchaseCoin[i] > 5):
				self.ui.TableData_2.item(i, 2).setBackground(QtGui.QColorConstants.Green)
			elif (changePurchaseCoin[i] < -4):
				self.ui.TableData_2.item(i, 2).setBackground(QtGui.QColorConstants.Red)

			self.ui.TableData_2.setItem(i, 3, qtw.QTableWidgetItem(
				str("{:.2f}".format(changePriceCoin[i]))+'%'))
			if(changePriceCoin[i] > 0):
				self.ui.TableData_2.item(i, 3).setBackground(QtGui.QColorConstants.Green)
			elif (changePriceCoin[i] < 0):
				self.ui.TableData_2.item(i, 3).setBackground(QtGui.QColorConstants.Red)

			self.ui.TableData_2.setItem(i, 4, qtw.QTableWidgetItem(
				str("{:.2f}".format(spreadCoin[i]))+'%'))
			if(spreadCoin[i] <= 0.5):
				self.ui.TableData_2.item(i, 4).setBackground(QtGui.QColorConstants.Green)
			else:
				self.ui.TableData_2.item(i, 4).setBackground(QtGui.QColorConstants.Red)
		self.ui.TableData_2.setSortingEnabled(True)

app = qtw.QApplication([])
widget = controllerWindow()
widget.showMaximized()
widget.updateImage()

sys.exit(app.exec_())


import numpy as np
import cv2 as cv

cap = cv.VideoCapture(0)

while(True):
    ret, frame = cap.read()

    hsv = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
    print(hsv[len(hsv)/2, len(hsv[0])/2])
    mask1 = cv.inRange(hsv, np.array([130, 10, 0]), np.array([180, 255, 255]))
    mask2 = cv.inRange(hsv, np.array([0, 10, 0]), np.array([20, 255, 255]))
    mask = mask1 + mask2
    mask = cv.blur(mask, (10, 10))
    mask = cv.threshold(mask, 200, 255, cv.THRESH_BINARY)[1]
    hsv = cv.bitwise_and(frame, frame, mask=mask)

    contours = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)[1]
    contour = sorted(contours, key=cv.contourArea, reverse=True)[0]
    for cnt in contours:
        x,y,w,h = cv.boundingRect(cnt)
        cv.rectangle(frame,(x,y),(x+w,y+h),(0,255,0),2)
        rect = cv.minAreaRect(cnt)
        box = cv.boxPoints(rect)
        box = np.int0(box)
        cv.drawContours(frame,[box],0,(0,0,255),2)
    cv.drawContours(frame, contours, -1, (0, 255, 0), 3)

    for x in range(-20, 20):
        frame[len(frame)/2 + x, len(frame[0])/2] = [255, 0, 0]
        frame[len(frame)/2, len(frame[0])/2 + x] = [255, 0, 0]
    cv.imshow('frame', frame)
    ##cv.imshow('contours', contour)
    if cv.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv.destroyAllWindows()


import math
import time
import warnings
from typing import Any, Dict, Optional, Type, Union

import gym
import numpy as np
import torch
from gym import spaces
from torch import optim
from torch.nn import functional as F

from paper_rl.architecture.ac.core import ActorCritic, count_vars
from paper_rl.common.buffer import GenericBuffer
from paper_rl.common.rollout import Rollout
from paper_rl.common.utils import to_torch
from paper_rl.logger.logger import Logger
from paper_rl.modelfree.ppo.buffer import PPOBuffer

class PPO:
    def __init__(
        self,
        ac: ActorCritic,
        env: gym.Env,
        num_envs: int,
        observation_space,
        action_space,
        steps_per_epoch: int = 10000,
        train_iters: int = 80,
        gamma: float = 0.99,
        gae_lambda: float = 0.97,
        clip_ratio: float = 0.2,
        ent_coef: float = 0.0,
        pi_coef: float = 1.0,
        vf_coef: float = 1.0,
        dapg_lambda: float = 0.1,
        dapg_damping: float = 0.99,
        # max_grad_norm: float = 0.5, # TODO
        # use_sde: bool = False,
        # sde_sample_freq: int = -1,
        target_kl: Optional[float] = 0.01,
        logger: Logger = None,
        # create_eval_env: bool = False,
        # verbose: int = 0,
        seed: Optional[int] = None,
        device: Union[torch.device, str] = "cpu",

    ) -> None:
        # Random seed
        if seed is None:
            seed = 0
        # seed += 10000 * proc_id()
        # torch.manual_seed(seed)
        # np.random.seed(seed)

        self.n_envs = num_envs
        self.env = env  # should be vectorized
        self.device = device
        self.observation_space = observation_space
        self.action_space = action_space

        # hparams
        self.target_kl = target_kl
        self.pi_coef = pi_coef
        self.ent_coef = ent_coef
        self.vf_coef = vf_coef
        self.clip_ratio = clip_ratio
        self.gae_lambda = gae_lambda
        self.gamma = gamma

        self.dapg_lambda = dapg_lambda
        self.dapg_damping = dapg_damping

        # self.pi_optimizer = optim.Adam(ac.pi.parameters(), lr=pi_lr)
        # self.vf_optimizer = optim.Adam(ac.v.parameters(), lr=vf_lr)

        # exp params
        self.train_iters = train_iters
        self.steps_per_epoch = steps_per_epoch

        self.logger = logger
        self.buffer = PPOBuffer(
            buffer_size=self.steps_per_epoch,
            observation_space=observation_space,
            action_space=action_space,
            n_envs=self.n_envs,
            gamma=self.gamma,
            lam=self.gae_lambda,
        )
        self.ac = ac.to(self.device)
        var_counts = tuple(count_vars(module) for module in [ac.pi, ac.v])
        self.logger.print(
            "\nNumber of parameters: \t pi: %d, \t v: %d\n" % var_counts,
            color="green",
            bold=True,
        )

    def to_state_dict(self):
        state_dict = dict(
            ac_state_dict=self.ac.state_dict(),
            # TODO
        )
        return state_dict

    def train(
        self,
        train_callback=None,
        rollout_callback=None,
        obs_to_tensor=None,
        max_ep_len=None,
        start_epoch: int = 0,
        n_epochs: int = 10,
        critic_warmup_epochs: int = 0, # number of epochs to collect rollouts and update critic only, freezing policy
        pi_optimizer: torch.optim.Optimizer = None,
        vf_optimizer: torch.optim.Optimizer = None,
        batch_size=1000,
        accumulate_grads=False,
        demo_trajectory_sampler = None,
        dapg_nll_loss=False,
    ):
        """
        Parameters
        ----------

        max_ep_len - max episode length. Can be set to infinity if environment returns a done signal to cap max episode length

        demo_trajectory_sampler - optional function that returns sampled demo trajectories to provide for DAPG training
        """
        dapg = False
        if demo_trajectory_sampler is not None: dapg = True
        if max_ep_len is None:
            # TODO: infer this
            raise ValueError("max_ep_len is missing")
        ac = self.ac
        env = self.env
        buf = self.buffer
        logger = self.logger
        clip_ratio = self.clip_ratio
        train_iters = self.train_iters
        target_kl = self.target_kl
        n_envs = self.n_envs
        device = self.device
        rollout = Rollout()
        def policy(o):
            if obs_to_tensor is None:
                o = torch.as_tensor(o, dtype=torch.float32)
            else:
                o = obs_to_tensor(o)
            return ac.step(o)

        def update(update_pi=True):
            all_demo_trajectories = None
            if dapg:
                all_demo_trajectories = []
                steps_per_train_iter = int(math.ceil(buf.buffer_size * buf.n_envs / batch_size))
                ac.eval()
                with torch.no_grad():
                    for _ in range(steps_per_train_iter):
                        demo_trajectories = demo_trajectory_sampler(batch_size)
                        if not dapg_nll_loss:
                            # collect old logp
                            pi, logp = ac.pi(to_torch(demo_trajectories["observations"], device=device), to_torch(demo_trajectories["actions"],  device=device))
                            demo_trajectories["logp"] = logp
                        all_demo_trajectories.append(demo_trajectories)
                ac.train()

            update_res = ppo_update(
                buffer=buf,
                ac=ac,
                pi_optimizer=pi_optimizer,
                vf_optimizer=vf_optimizer,
                pi_coef=self.pi_coef,
                vf_coef=self.vf_coef,
                clip_ratio=clip_ratio,
                train_iters=train_iters,
                batch_size=batch_size,
                target_kl=target_kl,
                logger=logger,
                device=device,
                accumulate_grads=accumulate_grads,
                update_pi=update_pi,
                demo_trajectories=all_demo_trajectories,
                dapg_lambda=self.dapg_lambda,
                dapg_nll_loss=dapg_nll_loss
            )
            pi_info, loss_pi, loss_v, update_step = (
                update_res["pi_info"],
                update_res["loss_pi"],
                update_res["loss_v"],
                update_res["update_step"],
            )
            logger.store(tag="train", StopIter=update_step, append=False)

            if loss_v is not None:
                logger.store(
                    tag="train",
                    LossV=loss_v.cpu().item(),
                )
            if pi_info is not None:
                kl, ent, cf, loss_pi, ppo_loss = pi_info["kl"], pi_info["ent"], pi_info["cf"], pi_info["loss_pi"], pi_info["ppo_loss"]
                logger.store(
                    tag="train",
                    LossPi=loss_pi,
                    PPOLoss=ppo_loss,
                    KL=kl,
                    Entropy=ent,
                    ClipFrac=cf,
                )
                if "dapg_actor_loss" in pi_info:
                    logger.store("train", LossDAPGActor=pi_info["dapg_actor_loss"])

        for epoch in range(start_epoch, start_epoch + n_epochs):

            # rollout
            rollout_start_time = time.time_ns()
            buf.reset()
            ac.eval()
            rollout.collect(policy=policy, env=env, n_envs=n_envs, buf=buf, steps=self.steps_per_epoch, rollout_callback=rollout_callback, max_ep_len=max_ep_len, logger=logger)
            ac.train()
            rollout_end_time = time.time_ns()
            rollout_delta_time = (rollout_end_time - rollout_start_time) * 1e-9
            logger.store("train", RolloutTime=rollout_delta_time, critic_warmup_epochs=critic_warmup_epochs, append=False)

            update_start_time = time.time_ns()
            # advantage normalization before update
            update_pi = epoch >= critic_warmup_epochs
            update(update_pi = update_pi)
            update_end_time = time.time_ns()

            if dapg:
                logger.store("train", dapg_lambda=self.dapg_lambda, append=False)
                if update_pi: self.dapg_lambda *= self.dapg_damping

            logger.store("train", UpdateTime=(update_end_time - update_start_time) * 1e-9, append=False)
            logger.store("train", Epoch=epoch, append=False)
            logger.store("train", TotalEnvInteractions=self.steps_per_epoch * self.n_envs * (epoch + 1), append=False)
            if train_callback is not None:
                early_stop = train_callback(epoch=epoch)
                if early_stop is not None and early_stop:
                    break

def compute_gae(
    buffer: GenericBuffer,
):
    # TODO move out the gae computations from within buffer's finish_path function
    buffer.buffers["adv_buf"] = (buffer.buffers["adv_buf"] - buffer.buffers["adv_buf"].mean()) / (buffer.buffers["adv_buf"].std())

def ppo_update(
    buffer: GenericBuffer,
    ac: ActorCritic,
    pi_optimizer,
    vf_optimizer,
    pi_coef,
    vf_coef,
    clip_ratio,
    train_iters,
    batch_size,
    target_kl,
    logger=None,
    device=torch.device("cpu"),
    accumulate_grads=False,
    update_pi=True,
    demo_trajectories=None,
    dapg_lambda=0.1,
    dapg_nll_loss=False,
):
    """

    demo_trajectories - dict of ndarray or list of dicts
        if list, treats each element as a batch

    dapg_nll_loss - bool
        if true, uses negative log likelihood as the demo augmented loss. Otherwise uses PPO loss
    """
    dapg = False
    if demo_trajectories is not None: dapg = True
    # TODO ABSTRACT OBS TO DEVICE TO A SEPARATE FUNC
    def compute_loss_pi(data):
        obs, act, adv, logp_old = data["obs_buf"], data["act_buf"].to(device), data["adv_buf"].to(device), data["logp_buf"].to(device)
        obs = to_torch(obs, device=device)

        # Policy loss)
        ac.pi.eval()
        pi, logp = ac.pi(obs, act)
        ac.pi.train()
        ratio = torch.exp(logp - logp_old)
        clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio) * adv
        loss_pi = -(torch.min(ratio * adv, clip_adv)).mean() * pi_coef

        # Entropy loss for some basic extra exploration
        entropy = pi.entropy()
        with torch.no_grad():
            # Useful extra info
            logr = logp - logp_old
            approx_kl = (torch.exp(logr) - 1 - logr).mean().cpu().item()
            clipped = ratio.gt(1 + clip_ratio) | ratio.lt(1 - clip_ratio)
            clipfrac = torch.as_tensor(clipped, dtype=torch.float32).mean().item()
        pi_info = dict(kl=approx_kl, ent=entropy.mean().item(), cf=clipfrac, ppo_loss=loss_pi.cpu().item())

        # add dapg loss, the negative log likelihood in particular
        if dapg:
            demo_trajectories = data["demo_trajectories"]
            ac.pi.eval()
            demo_pi, demo_logp = ac.pi(to_torch(demo_trajectories["observations"], device=device), to_torch(demo_trajectories["actions"], device=device))
            ac.pi.train()

            if dapg_nll_loss:
                dapg_actor_loss = -demo_logp.mean()
            else:
                # mark advantage as 1 since advantage is normalized
                demo_logp_old = demo_trajectories["logp"]
                ratio = torch.exp(demo_logp - demo_logp_old)
                clip_adv = torch.clamp(ratio, 1 - clip_ratio, 1 + clip_ratio)
                dapg_actor_loss = -(torch.min(ratio, clip_adv)).mean()
            dapg_actor_loss = dapg_actor_loss * dapg_lambda
            pi_info["dapg_actor_loss"] = dapg_actor_loss.cpu().item()
            loss_pi = loss_pi + dapg_actor_loss
        pi_info["loss_pi"] = loss_pi.cpu().item()

        return dict(loss_pi=loss_pi, logp=logp, entropy=entropy, pi_info=pi_info)

    # Set up function for computing value loss
    def compute_loss_v(data):
        obs, ret = data["obs_buf"], data["ret_buf"].to(device)
        if isinstance(obs, dict):
            for k in obs.keys():
                obs[k] = obs[k].to(device)
        else:
            obs = obs.to(device)
        return ((ac.v(obs) - ret) ** 2).mean()

    # Train policy with multiple steps of gradient descent
    update_step = 0
    early_stop_update = False
    loss_pi = None
    pi_info = None

    # compute gae here
    compute_gae(buffer=buffer)

    for _ in range(train_iters):
        if early_stop_update:
            break
        N = buffer.buffer_size * buffer.n_envs

        steps_per_train_iter = int(math.ceil(N / batch_size))
        if accumulate_grads:
            if update_pi: pi_optimizer.zero_grad()
            vf_optimizer.zero_grad()
            average_kl = 0
        for batch_idx in range(steps_per_train_iter):
            batch_data = buffer.sample_batch(batch_size=batch_size)
            if dapg:
                batch_demo_trajectories = demo_trajectories[batch_idx]
                batch_data["demo_trajectories"] = batch_demo_trajectories
            if update_pi:
                loss_pi_data = compute_loss_pi(batch_data)
                loss_pi, logp, entropy, pi_info = loss_pi_data["loss_pi"], loss_pi_data["logp"], loss_pi_data["entropy"], loss_pi_data["pi_info"]
                kl = pi_info["kl"]
                if accumulate_grads:
                    average_kl += kl
                if not accumulate_grads and target_kl is not None:
                    if kl > 1.5 * target_kl:
                        logger.print("Early stopping at step %d due to reaching max kl." % update_step)
                        early_stop_update = True
                        break
            loss_v = compute_loss_v(batch_data) * vf_coef
            # TODO - entropy loss
            # if entropy is None:
            #     # Approximate entropy when no analytical form
            #     entropy_loss = -torch.mean(-logp)
            # else:
            #     entropy_loss = -torch.mean(entropy)
            # torch.nn.utils.clip_grad.clip_grad_norm_() # TODO
            if not accumulate_grads:
                if update_pi:
                    pi_optimizer.zero_grad()
                vf_optimizer.zero_grad()
                if update_pi:
                    loss_pi.backward()
                    pi_optimizer.step()
                loss_v.backward()
                vf_optimizer.step()
                update_step += 1
            if accumulate_grads:
                # scale loss down
                if update_pi: loss_pi /= steps_per_train_iter
                loss_v /= steps_per_train_iter
                if update_pi: loss_pi.backward()
                loss_v.backward()

        if accumulate_grads and target_kl is not None:
            average_kl /= steps_per_train_iter
            if average_kl > 1.5 * target_kl:
                logger.print("Early stopping at step %d due to reaching max kl." % update_step)
                early_stop_update = True
                break
        if accumulate_grads:
            if update_pi: pi_optimizer.step()
            vf_optimizer.step()
            update_step += 1
        if early_stop_update:
            break
    return dict(
        pi_info=pi_info,
        update_step=update_step,
        loss_v=loss_v,
        loss_pi=loss_pi,
    )


from keras.models import Sequential, Model, model_from_json
import matplotlib.pyplot as plt
import keras
import pickle
import wave  # !pip install wave
import os
import pandas as pd
import numpy as np
import sys
import warnings
import librosa
import librosa.display
import IPython.display as ipd  # To play sound in the notebook

from tensorflow.keras.utils import to_categorical
from tensorflow.keras import optimizers

import sounddevice as sd

from flask import send_from_directory

# ignore warnings
if not sys.warnoptions:
    warnings.simplefilter("ignore")

def run():

    # OUTPUT_FILE = send_from_directory("../client/public/", 'record.wav')() #File name to save
    OUTPUT_FILE = '../recordings/record.wav'
    DURATION = 3  #Recording length (seconds)
    RATE = 44100  #Sampling frequency
    CHANNELS = 2

    #Start recording (wave_length Record for seconds. Wait until the recording is finished with wait)
    print("* Recording ..")
    sd.default.samplerate = RATE
    sd.default.channels = CHANNELS
    data = sd.rec(int(DURATION * RATE))
    sd.wait()
    print("* Done Recording.")

    #Normalize. Since it is recorded with 16 bits of quantization bit, it is maximized in the range of int16.
    print("* Normalization ..")
    data = data / data.max() * np.iinfo(np.int16).max

    # float -> int
    data = data.astype(np.int16)

    #Save file
    with wave.open(OUTPUT_FILE, mode='wb') as wb:
        wb.setnchannels(CHANNELS)  #monaural
        wb.setsampwidth(2)  # 16bit=2byte
        wb.setframerate(RATE)
        wb.writeframes(data.tobytes())  #Convert to byte string

    print("* WAV file saved successfully.")

    return True


import cv2
import numpy as np
import matplotlib
matplotlib.use('Agg')
# matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt
from hyperpose import Config,Model,Dataset
from hyperpose.Dataset import imread_rgb_float,imwrite_rgb_float
matplotlib.get_backend()
import os
import sys
import time
import math
import json
import glob
import argparse
import matplotlib
import numpy as np
import tensorflow as tf
import tensorlayer as tl
os.environ['CUDA_VISIBLE_DEVICES']='0,1'

gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[0], True)
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[1], True)

model_type = "Openpose"
# model_backbone = "Default"
model_backbone = "Mobilenetv2"
# MobilenetV1.MobilenetV2.Vggtiny.Vgg16.Vgg19.Resnet18.Resnet50
dataset_type = "USERDEF"
# model_name = dataset_type+'_'+model_type+'_' + model_backbone + '_'+time.strftime("%Y%m%d_%H%M")
model_name = dataset_type+'_'+time.strftime("%Y%m%d_%H%M")
dataset_version ="2020"
dataset_path = "./data"
train_type = "Single_train"
kf_optimizer ='Sync_avg'
userdef_dataset = 1
useradd_data_path = "./data"

domainadapt_data_path = None

optim_type ="Adam"
learning_rate =1e-4
log_interval =1e2
save_interval =5e3

Config.set_model_name(model_name)
Config.set_model_type(Config.MODEL[model_type])
Config.set_model_backbone(Config.BACKBONE[model_backbone])
Config.set_log_interval(log_interval)
Config.set_save_interval(save_interval)
# config train
Config.set_train_type(Config.TRAIN[train_type])
Config.set_learning_rate(learning_rate)
Config.set_optim_type(Config.OPTIM[optim_type])
Config.set_kungfu_option(Config.KUNGFU[kf_optimizer])

# config dataset
print(f"test enabling dataset:{userdef_dataset}")
Config.set_userdef_dataset(userdef_dataset)
Config.set_dataset_type(Config.DATA[dataset_type])
Config.set_dataset_path(dataset_path)
Config.set_dataset_version(dataset_version)

print('\n' * 4)

# train
config = Config.get_config()

config.pretrain.batch_size=4
config.eval.batch_size=1
config.train.batch_size=4

config.train.save_interval = 5000
config.train.n_step = 100000
print(config.model)
print('\n' * 2)
print(config.train)
print('\n' * 2)
print(config.data)
print('\n' * 4)

model = Model.get_model(config)

# Config.set_model_name("default_name")
# # Config.set_model_name("openpose_test_20201028")
# # Config.set_model_type(Config.MODEL.Openpose)
# Config.set_model_type(Config.MODEL.LightweightOpenpose)
# Config.set_model_backbone(Config.BACKBONE.Resnet50)
# config=Config.get_config()
# print(config.model)
# # config.model
#
#
# #get and load model
# model=Model.get_model(config)

weight_path = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_20201121_1621/model_dir/newest_model.npz'
weight_path = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_Openpose_Mobilenetv2_20201124_2030/model_dir/newest_model.npz'

'''
import os
import cv2
import sys
import time
import math
import json
import glob
import argparse
import matplotlib
import numpy as np
import tensorflow as tf
import tensorlayer as tl
from hyperpose import Config,Model,Dataset

os.environ['CUDA_VISIBLE_DEVICES']='0,1'
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[0], True)
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[1], True)

model_type = "Openpose"
# model_backbone = "Default"
# model_backbone = "Mobilenetv2"
model_backbone = "Resnet18"
# MobilenetV1.MobilenetV2.Vggtiny.Vgg16.Vgg19.Resnet18.Resnet50

dataset_type = "USERDEF"
model_name = dataset_type+'_'+model_type+'_' + model_backbone + '_'+time.strftime("%Y%m%d_%H%M")
dataset_version ="2020"
dataset_path = "./data"
train_type = "Single_train"
kf_optimizer ='Sync_avg'
userdef_dataset = 1
useradd_data_path = "./data"

domainadapt_data_path = None

optim_type ="Adam"
learning_rate =1e-4
log_interval =1e2
save_interval =5e3

Config.set_model_name(model_name)
Config.set_model_type(Config.MODEL[model_type])
Config.set_model_backbone(Config.BACKBONE[model_backbone])
Config.set_log_interval(log_interval)
Config.set_save_interval(save_interval)
# config train
Config.set_train_type(Config.TRAIN[train_type])
Config.set_learning_rate(learning_rate)
Config.set_optim_type(Config.OPTIM[optim_type])
Config.set_kungfu_option(Config.KUNGFU[kf_optimizer])

# config dataset
print(f"test enabling dataset:{userdef_dataset}")
Config.set_userdef_dataset(userdef_dataset)
Config.set_dataset_type(Config.DATA[dataset_type])
Config.set_dataset_path(dataset_path)
Config.set_dataset_version(dataset_version)

print('\n' * 4)

# train
config = Config.get_config()

config.pretrain.batch_size=4
config.eval.batch_size=1
config.train.batch_size=4
config.model.hin = 48*2
config.model.win = 64*2
config.model.hout = 6*2 #48/8
config.model.wout = 8*2 #64/8
config.train.save_interval = 5000
config.train.n_step = 60000
# config.train.lr_decay_every_step = 2000

print(config.model)
print('\n' * 2)
print(config.train)
print('\n' * 2)
print(config.data)
print('\n' * 4)

model = Model.get_model(config)
'''
weight_path = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_Openpose_Resnet18_20201201_1927/model_dir/newest_model.npz' # 这个结果很不错
# 只有四個點，最後一個點不准

'''

import os
import cv2
import sys
import time
import math
import json
import glob
import argparse
import matplotlib
import numpy as np
import tensorflow as tf
import tensorlayer as tl
from hyperpose import Config,Model,Dataset

os.environ['CUDA_VISIBLE_DEVICES']='0,1'
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[0], True)
gpu = tf.config.experimental.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(gpu[1], True)

model_type = "Openpose"
# model_backbone = "Default"
# model_backbone = "Mobilenetv2"
model_backbone = "Resnet18"
# MobilenetV1.MobilenetV2.Vggtiny.Vgg16.Vgg19.Resnet18.Resnet50

dataset_type = "USERDEF"
model_name = dataset_type+'_'+model_type+'_' + model_backbone + '_'+time.strftime("%Y%m%d_%H%M")
dataset_version ="2020"
dataset_path = "./data"
train_type = "Single_train"
kf_optimizer ='Sync_avg'
userdef_dataset = 1
useradd_data_path = "./data"

domainadapt_data_path = None

optim_type ="Adam"
learning_rate =1e-4
log_interval =1e2
save_interval =5e3

Config.set_model_name(model_name)
Config.set_model_type(Config.MODEL[model_type])
Config.set_model_backbone(Config.BACKBONE[model_backbone])
Config.set_log_interval(log_interval)
Config.set_save_interval(save_interval)
# config train
Config.set_train_type(Config.TRAIN[train_type])
Config.set_learning_rate(learning_rate)
Config.set_optim_type(Config.OPTIM[optim_type])
Config.set_kungfu_option(Config.KUNGFU[kf_optimizer])

# config dataset
print(f"test enabling dataset:{userdef_dataset}")
Config.set_userdef_dataset(userdef_dataset)
Config.set_dataset_type(Config.DATA[dataset_type])
Config.set_dataset_path(dataset_path)
Config.set_dataset_version(dataset_version)

print('\n' * 4)

# train
config = Config.get_config()

config.pretrain.batch_size=4
config.eval.batch_size=1
config.train.batch_size=4
config.model.hin = 48*2
config.model.win = 64*2
config.model.hout = 6*2 #48/8
config.model.wout = 8*2 #64/8
config.train.save_interval = 5000
config.train.n_step = 60000
# config.train.lr_decay_every_step = 2000

print(config.model)
print('\n' * 2)
print(config.train)
print('\n' * 2)
print(config.data)
print('\n' * 4)

model = Model.get_model(config)

# train = Model.get_train(config)
#
# dataset = Dataset.get_dataset(config)
#
#
# train(model, dataset)

'''

weight_path = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_Openpose_Resnet18_20201202_2325/model_dir/newest_model.npz' # 这个结果很不错

model.load_weights(weight_path)

import matplotlib
matplotlib.use('Agg')
# matplotlib.use('Qt5Agg')
import matplotlib.pyplot as plt
# C:\Users\ps\Desktop\djz\hyperpose\save_dir\USERDEF_Openpose_Resnet18_20201130_2259\train_vis_dir
# UserdefColor = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0]]
UserdefColor = [[1, 0, 0], [0, 0, 1], [0, 1, 0], [0, 0, 1],[0, 0, 1], [0, 1, 0]]

save_all_pic = 1
if save_all_pic == 1:
    path_from = './data/egg/fly/train2020/'
    # C:\Users\ps\Desktop\djz\hyperpose\data\egg\fly\train2020
    # path_from = './data/fly/train2020/'
    # path_from = 'H:/data/picture_fly/2/'
    # path_to = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_20201121_1621/val/'
    # path_to = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_20201121_1621/val_head_output/'
    # path_to = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_Openpose_Mobilenetv2_20201124_2030/val_output/'
    # path_to = 'C:/Users/ps/Desktop/djz/hyperpose/save_dir/USERDEF_Openpose_Mobilenetv2_20201124_2030/val_head_output/'
    # C:\Users\ps\Desktop\djz\hyperpose\save_dir\USERDEF_Openpose_Resnet18_20201130_1642\model_dir
    # path_to = 'C:/Users/ps/Desktop/djz/hyperpose/rubbish/val_head_output/'
    # path_ta = 'C:\Users\ps\Desktop\djz\hyperpose\save_dir\USERDEF_Openpose_Resnet18_20201130_2259\val_head_output/'
    # path_to = './save_dir/USERDEF_Openpose_Resnet18_20201202_2325/1/'
    path_to = './save_dir/USERDEF_Openpose_Resnet18_20201223_2319/train_all/'
    path_to = './save_dir/USERDEF_Openpose_Resnet18_20201223_2319/train_all_3/'

    read_name = os.listdir(path_from)
    number_of_pic = len(read_name)

    start_time = time.time()
    for K_0 in range(number_of_pic):#(10):#(number_of_pic):# K_0=0
        img_name = read_name[K_0]
        img0 = cv2.imread(path_from+img_name)
        # img1 = cv2.resize(img0, (432, 368))
        # img1 = cv2.resize(img0, (384, 288))
        img1 = cv2.resize(img0, (128, 96))
        # img1 = np.copy(img0)
        # img1 = cv2.resize(img0, (96, 96))#TODO(JZ)
        # plt.imshow(img1)
        ori_image = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
        input_image = ori_image.astype(np.float32) / 255.0
        if (model.data_format == "channels_first"):
            input_image = np.transpose(input_image, [2, 0, 1])
        img_c, img_h, img_w = input_image.shape

        conf_map, paf_map = model.infer(input_image[np.newaxis, :, :, :])
        find_all_part = 0
        if find_all_part == 1:
            visualize = Model.get_visualize(Config.MODEL.Openpose)
            # def visualize(img,conf_map,paf_map,save_name="maps",save_dir="./save_dir/vis_dir",data_format="channels_first",save_tofile=True):
            vis_parts_heatmap, vis_limbs_heatmap = visualize(input_image, conf_map[0], paf_map[0],
                                                         save_name=img_name[:-4],
                                                         save_dir=path_to,
                                                         data_format=model.data_format )

        find_all = 1
        if find_all == 1:
            conf_map_head = np.transpose(conf_map[0], [2, 0, 1])
            for K_1 in range(5):
                conf_head = np.abs(conf_map_head[K_1, :, :])
                heatmap_avg = cv2.resize(conf_head,
                                         (input_image.shape[1], input_image.shape[0]),
                                         interpolation=cv2.INTER_CUBIC)
                x, y = np.unravel_index(np.argmax(heatmap_avg), heatmap_avg.shape)
                center = y, x  # x, y      #int(x*8),int(y*8)    #y, x
                radius = 3  # int(radius)
                # cv2.circle(input_image, center, radius, (1, 0, 0), -1)
                cv2.circle(input_image, center, radius, UserdefColor[K_1], -1)
            plt.imsave(path_to + img_name, input_image)
        find_head = 0
        if find_head == 1:
            # np.abs(conf_map[0, :, :])

            conf_map_head = np.transpose(conf_map[0], [2, 0, 1])
            conf_head=np.abs(conf_map_head[0, :, :])
            # show_conf_map = np.abs(conf_map[0, :, :])
            # plt.imshow(conf_head)
            heatmap_avg = cv2.resize(conf_head,
                                     (input_image.shape[1], input_image.shape[0]),
                                     interpolation=cv2.INTER_CUBIC)
            x, y = np.unravel_index(np.argmax(heatmap_avg), heatmap_avg.shape)
            # x, y = np.unravel_index(np.argmax(conf_head), conf_head.shape)
            center = y,x#x, y      #int(x*8),int(y*8)    #y, x
            radius = 3  # int(radius)
            cv2.circle(input_image, center, radius, (1, 0, 0), -1)
            plt.imsave(path_to + img_name, input_image)
            # plt.imsave(img_name, input_image)

        vis_and_head = 0
        if vis_and_head==1:

            img, conf_map, paf_map = input_image, conf_map[0], paf_map[0]
            save_name = img_name[:-4]
            save_dir = path_to
            data_format = model.data_format
            save_tofile = True
            if (type(img) != np.ndarray):
                img = img.numpy()
            if (type(conf_map) != np.ndarray):
                conf_map = conf_map.numpy()
            if (type(paf_map) != np.ndarray):
                paf_map = paf_map.numpy()

            if (data_format == "channels_last"):
                conf_map = np.transpose(conf_map, [2, 0, 1])
                paf_map = np.transpose(paf_map, [2, 0, 1])
            elif (data_format == "channels_first"):
                img = np.transpose(img, [1, 2, 0])
            os.makedirs(save_dir, exist_ok=True)
            ori_img = np.clip(img * 255.0, 0.0, 255.0).astype(np.uint8)
            vis_img = ori_img.copy()
            fig = plt.figure(figsize=(8, 8))
            # show input image
            a = fig.add_subplot(2, 2, 1)
            a.set_title("input image")
            plt.imshow(vis_img)
            # show conf_map
            show_conf_map = np.abs(conf_map[0, :, :])
            #np.amax(np.abs(conf_map[:-1, :, :]), axis=0)
            a = fig.add_subplot(2, 2, 3)
            a.set_title("conf_map")
            plt.imshow(show_conf_map)
            # show paf_map
            show_paf_map = np.amax(np.abs(paf_map[0:2, :, :]), axis=0)
            a = fig.add_subplot(2, 2, 4)
            a.set_title("paf_map")
            plt.imshow(show_paf_map)
            # save
            if (save_tofile):
                plt.savefig(f"{save_dir}/{save_name}_visualize.png")
                plt.close('all')

        if K_0%10 == 0:
            print(K_0)
    end_time = time.time()
    print("Time used: ", end_time - start_time, 's')
else:
    matplotlib.use('Agg')
    # infer on single image
    # ori_image=cv2.cvtColor(cv2.imread("./sample.jpg"),cv2.COLOR_BGR2RGB)
    img_name = "./data/fly/train2020/000000000182.jpg"

    ori_image = cv2.cvtColor(cv2.imread(img_name), cv2.COLOR_BGR2RGB)
    input_image = ori_image.astype(np.float32) / 255.0
    if (model.data_format == "channels_first"):
        input_image = np.transpose(input_image, [2, 0, 1])

    img_c, img_h, img_w = input_image.shape

    conf_map, paf_map = model.infer(input_image[np.newaxis, :, :, :])

    visualize = Model.get_visualize(Config.MODEL.Openpose)

    # def visualize(img,conf_map,paf_map,save_name="maps",save_dir="./save_dir/vis_dir",data_format="channels_first",save_tofile=True):
    vis_parts_heatmap, vis_limbs_heatmap = visualize(input_image, conf_map[0], paf_map[0],
                                                     data_format=model.data_format, save_tofile=False, )

    postprocess = Model.get_postprocess(Config.MODEL.Openpose)
    humans = postprocess(conf_map[0], paf_map[0], img_h, img_w, model.parts, model.limbs, model.data_format,
                         model.colors)
    # humans=postprocess(conf_map,paf_map,img_h,img_w,model.parts,model.limbs,model.data_format,model.colors)
    #

    # img = input_image
    conf_map = conf_map[0]
    paf_map = paf_map[0]
    if (type(conf_map) != np.ndarray):
        conf_map = conf_map.numpy()
    if (type(paf_map) != np.ndarray):
        paf_map = paf_map.numpy()
    parts = model.parts
    limbs = model.limbs
    colors = model.colors
    if (colors == None):
        colors = [[255, 0, 0]] * len(parts)
    from hyperpose.Model.openpose.infer import Post_Processor

    # from .infer import Post_Processor
    post_processor = Post_Processor(parts, limbs, colors)
    data_format = model.data_format
    humans = post_processor.process(conf_map, paf_map, img_h, img_w, data_format=data_format)

    # #draw all detected skeletons
    output_img = ori_image.copy()
    for human in humans:
        output_img = human.draw_human(output_img)

    # if you want to visualize all the images in one plot:
    # show image,part heatmap,limb heatmap and detected image
    # here we use 'transpose' because our data_format is 'channels_first'
    fig = plt.figure(figsize=(8, 8))
    # origin image
    origin_fig = fig.add_subplot(2, 2, 1)
    origin_fig.set_title("origin image")
    origin_fig.imshow(ori_image)
    # parts heatmap
    parts_fig = fig.add_subplot(2, 2, 2)
    parts_fig.set_title("parts heatmap")
    parts_fig.imshow(vis_parts_heatmap)
    # limbs heatmap
    limbs_fig = fig.add_subplot(2, 2, 3)
    limbs_fig.set_title("limbs heatmap")
    limbs_fig.imshow(vis_limbs_heatmap)
    # detected results
    result_fig = fig.add_subplot(2, 2, 4)
    result_fig.set_title("detect result")
    result_fig.imshow(output_img)
    # save fig
    plt.savefig("./sample_custome_infer.png")
    plt.close()

# import time
#
# start_time = time.time()
# for K in range(10):
#     conf_map, paf_map = model.infer(input_image[np.newaxis, :, :, :])
# end_time = time.time()
# print("Time used: ", end_time-start_time, 's')
#
#
# import matplotlib.pyplot as plt
# conf_temp=conf_map[0,:,:,0]
# plt.imshow(conf_temp,"gray")
# plt.imsave('a1234343124.jpg',conf_temp)
#get visualize function, which is able to get visualized part and limb heatmap image from inferred heatmaps
#
#
#
# img = input_image
# conf_map = conf_map[0]
# paf_map = paf_map[0]
#
# if(type(img)!=np.ndarray):
#     img=img.numpy()
# if(type(conf_map)!=np.ndarray):
#     conf_map=conf_map.numpy()
# if(type(paf_map)!=np.ndarray):
#     paf_map=paf_map.numpy()
#
# # import numpy as np
# data_format = model.data_format
# if(data_format=="channels_last"):
#     conf_map=np.transpose(conf_map,[2,0,1])
#     paf_map=np.transpose(paf_map,[2,0,1])
# elif(data_format=="channels_first"):
#     img=np.transpose(img,[1,2,0])
#
# # import os
# # os.makedirs(save_dir,exist_ok=True)
# #
# # import matplotlib
# # matplotlib.use('Qt5Agg')
#
# ori_img=np.clip(img*255.0,0.0,255.0).astype(np.uint8)
# vis_img=ori_img.copy()
# fig=plt.figure(figsize=(8,8))
# #show input image
# a=fig.add_subplot(2,2,1)
# a.set_title("input image")
# plt.imshow(vis_img)
# #show conf_map
# show_conf_map=np.amax(np.abs(conf_map[:-1,:,:]),axis=0)
# # show_conf_map=np.amax(conf_map[:-1,:,:],axis=0)
# # show_conf_map = np.abs(conf_map[0,:,:])
# a=fig.add_subplot(2,2,3)
# a.set_title("conf_map")
# plt.imshow(show_conf_map, alpha=0.8)
# #show paf_map
# show_paf_map=np.amax(np.abs(paf_map[:,:,:]),axis=0)
# # show_paf_map = np.abs(paf_map[0,:,:])
# a=fig.add_subplot(2,2,4)
# a.set_title("paf_map")
# plt.imshow(show_paf_map)

# plt.savefig("save_name_visualize.png")

#save
# if(save_tofile):
#     plt.savefig(f"{save_dir}/{save_name}_visualize.png")
#     plt.close('all')
#
#
#
#
#
#
#
# config.model
# #get postprocess function, which is able to get humans that contains assembled detected parts from inferred heatmaps
# postprocess=Model.get_postprocess(Config.MODEL.Openpose)
# humans=postprocess(conf_map[0],paf_map[0],img_h,img_w,model.parts,model.limbs,model.data_format,model.colors)
# # humans=postprocess(conf_map,paf_map,img_h,img_w,model.parts,model.limbs,model.data_format,model.colors)
# #
#
# # img = input_image
# conf_map = conf_map[0]
# paf_map = paf_map[0]
# if(type(conf_map)!=np.ndarray):
#     conf_map=conf_map.numpy()
# if(type(paf_map)!=np.ndarray):
#     paf_map=paf_map.numpy()
# parts= model.parts
# limbs = model.limbs
# colors = model.colors
# if(colors==None):
#     colors=[[255,0,0]]*len(parts)
# from hyperpose.Model.openpose.infer import Post_Processor
# # from .infer import Post_Processor
# post_processor=Post_Processor(parts,limbs,colors)
# data_format = model.data_format
# humans=post_processor.process(conf_map,paf_map,img_h,img_w,data_format=data_format)
#
# # #draw all detected skeletons
# output_img=ori_image.copy()
# for human in humans:
#     output_img=human.draw_human(output_img)
#
# #if you want to visualize all the images in one plot:
# #show image,part heatmap,limb heatmap and detected image
# #here we use 'transpose' because our data_format is 'channels_first'
# fig=plt.figure(figsize=(8,8))
# #origin image
# origin_fig=fig.add_subplot(2,2,1)
# origin_fig.set_title("origin image")
# origin_fig.imshow(ori_image)
# #parts heatmap
# parts_fig=fig.add_subplot(2,2,2)
# parts_fig.set_title("parts heatmap")
# parts_fig.imshow(vis_parts_heatmap)
# #limbs heatmap
# limbs_fig=fig.add_subplot(2,2,3)
# limbs_fig.set_title("limbs heatmap")
# limbs_fig.imshow(vis_limbs_heatmap)
# #detected results
# result_fig=fig.add_subplot(2,2,4)
# result_fig.set_title("detect result")
# result_fig.imshow(output_img)
# #save fig
# plt.savefig("./sample_custome_infer.png")
# plt.close()


# -*- coding: utf-8 -*-
"""
face recognition defines using dlib with guidance of wuhuikai/FaceSwap

@author: John Watson and Damen Kelly
"""

import cv2
import dlib
import numpy as np
from scipy.spatial import Delaunay

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor("../dlib-models/shape_predictor_68_face_landmarks.dat")
recogniser = dlib.face_recognition_model_v1("../dlib-models/dlib_face_recognition_resnet_model_v1.dat")

def match(src_descriptor, descriptors):
    distances = np.linalg.norm(descriptors - src_descriptor, axis=1)
    if True in list(distances <= 0.6):
        return np.argmin(distances)
    else:
        return None

def recognise(img, shape):
    descriptor = np.array(recogniser.compute_face_descriptor(img, shape))
    return descriptor

def swap(src_img, src_points, src_bbox, input_img, input_points):
    if src_points is None or src_bbox is None:  # no face detected
        print("No face to swap")
        return src_img

    result_img = warp(src_img, src_points, src_bbox, input_img, input_points)

    """src_img = draw(src_img)
    for point in input_points:
        cv2.circle(input_img, (point[0], point[1]), 2, (0, 255, 0), -1)
    bbox = find(input_img)[1]
    cv2.line(input_img, tuple(bbox[0]), tuple(bbox[1]), (0, 0, 255), 2)
    cv2.line(input_img, tuple(bbox[0]), tuple(bbox[2]), (0, 0, 255), 2)
    cv2.line(input_img, tuple(bbox[2]), tuple(bbox[3]), (0, 0, 255), 2)
    cv2.line(input_img, tuple(bbox[1]), tuple(bbox[3]), (0, 0, 255), 2)
    cv2.imshow("John", src_img)
    cv2.imshow("Andrew", input_img)
    cv2.waitKey(0)  # Demonstration"""

    ## Mask for blending
    h, w = src_img.shape[:2]
    mask = mask_from_points((h, w), src_points)
    #cv2.imshow("mask", mask)
    mask_src = np.mean(result_img, axis=2) > 0
    mask = np.asarray(mask * mask_src, dtype=np.uint8)
    #cv2.imshow("mask1", mask)

    # colour correction
    #Apply Masks
    result_img = cv2.bitwise_and(result_img, result_img, mask=mask)
    #cv2.imshow("result_img", result_img)
    dst_face_masked = cv2.bitwise_and(src_img, src_img, mask=mask)
    #cv2.imshow("dst_face_masked", dst_face_masked)
    result_img = correct_colours(dst_face_masked, result_img, src_points)
    #cv2.imshow("result_img2", result_img)
    ## Shrink the mask
    kernel = np.ones((10, 10), np.uint8)
    mask = cv2.erode(mask, kernel, iterations=1)
    #cv2.imshow("mask3", mask)
    ##Poisson Blending
    r = cv2.boundingRect(mask)
    center = ((r[0] + int(r[2] / 2), r[1] + int(r[3] / 2)))
    result_img = cv2.seamlessClone(result_img, src_img, mask, center, cv2.NORMAL_CLONE)
    #cv2.imshow("result_img3", result_img)
    #cv2.waitKey(0)
    return result_img

def correct_colours(im1, im2, landmarks1):
    COLOUR_CORRECT_BLUR_FRAC = 0.75     #Blending amount
    LEFT_EYE_POINTS = list(range(42, 48))
    RIGHT_EYE_POINTS = list(range(36, 42))
    #determine kernal size
    blur_amount = COLOUR_CORRECT_BLUR_FRAC * np.linalg.norm(np.mean(landmarks1[LEFT_EYE_POINTS], axis=0) - np.mean(landmarks1[RIGHT_EYE_POINTS], axis=0))
    blur_amount = int(blur_amount)
    #Ensure kernal size value is not even
    if blur_amount % 2 == 0:
        blur_amount += 1
    # Apply blur
    im1_blur = cv2.medianBlur(im1, blur_amount, 0)
    im2_blur = cv2.medianBlur(im2, blur_amount, 0)

    # Avoid divide-by-zero errors.
    im2_blur = im2_blur.astype(int)
    im2_blur = im2_blur + 128*(im2_blur <= 1)
    #Blend images
    result = np.float64(im2) * np.float64(im1_blur) / np.float64(im2_blur)
    #Remove overflown pixel values
    result = np.uint8(np.clip(result, 0, 255))

    return result

def mask_from_points(size, points):
    #Create kernel array
    kernel = np.ones((10, 10), np.uint8)
    #Create mask of zeroes
    mask = np.zeros(size, np.uint8)
    cv2.fillConvexPoly(mask, cv2.convexHull(points), 255)
    mask = cv2.erode(mask, kernel, iterations=1)
    return mask

def find(img):
    detections = detector(img)  # detect faces and store rectangles
    if len(detections) == 0:
        print("No face detected")
        return None, None, None
    shape = predictor(img, detections[0])  # find the shape of the face (use first detected face)
    points = np.array(list([point.x, point.y] for point in shape.parts()))

    points = find_pupils(img, points)

    bbox = [[min(points[:, 0]), min(points[:, 1])], [max(points[:, 0]), min(points[:, 1])],
            [min(points[:, 0]), max(points[:, 1])], [max(points[:, 0]), max(points[:, 1])]]
    '''
    0   1
    2   3
    '''

    return points, bbox, shape

def find_pupils(img, points):
    left_bbox = [[min(points[36:41, 0]), min(points[36:41, 1])], [max(points[36:41, 0]), min(points[36:41, 1])],
                 [min(points[36:41, 0]), max(points[36:41, 1])], [max(points[36:41, 0]), max(points[36:41, 1])]]
    right_bbox = [[min(points[42:47, 0]), min(points[42:47, 1])], [max(points[42:47, 0]), min(points[42:47, 1])],
                  [min(points[42:47, 0]), max(points[42:47, 1])], [max(points[42:47, 0]), max(points[42:47, 1])]]
    # find bounding box for eyes

    bboxs = [left_bbox, right_bbox]

    for bbox in bboxs:
        eye = img[bbox[0][1]:bbox[3][1], bbox[0][0]:bbox[3][0]]  # crop image to each eye
        #cv2.imshow("eye", eye)
        eye = cv2.cvtColor(eye, cv2.COLOR_BGR2GRAY)  # convert to single channel
        #cv2.imshow("eye1", eye)
        #eye = cv2.inRange(eye, (0, 0, 0), (50, 50, 50))
        eye = cv2.GaussianBlur(eye, (3, 3), 0)
        #cv2.imshow("eye2", eye)
        eye = cv2.erode(eye, (3, 3), iterations=3)
        #cv2.imshow("eye3", eye)
        ret, _ = cv2.threshold(eye, 0, 255, cv2.THRESH_OTSU)
        _, eye = cv2.threshold(eye, ret*0.7, 255, cv2.THRESH_BINARY_INV)
        #cv2.imshow("eye4", eye)
        #cv2.waitKey(0)

        try:
            m = cv2.moments(eye)
            x = int(m["m10"] / m["m00"])
            y = int(m["m01"] / m["m00"])
            eye = cv2.cvtColor(eye, cv2.COLOR_GRAY2BGR)
            cv2.circle(eye, (x, y), 1, (0, 0, 255), -1)
            #cv2.imshow("eye", eye)
            #cv2.waitKey(0)
            points = np.vstack((points, [bbox[0][0] + x, bbox[0][1] + y]))  # absolute coordinates
        except (IndexError, ZeroDivisionError):
            points = np.vstack((points, [int((bbox[0][0] + bbox[3][0]) / 2), int((bbox[0][1] + bbox[3][1]) / 2)]))

        """contours, _ = cv2.findContours(eye, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
        contours = sorted(contours, key=cv2.contourArea)
        try:
            m = cv2.moments(contours[-2])
            x = int(m["m10"] / m["m00"])
            y = int(m["m01"] / m["m00"])
            eye = cv2.cvtColor(eye, cv2.COLOR_GRAY2BGR)
            cv2.circle(eye, (x, y), 1, (0, 0, 255), -1)
            cv2.drawContours(eye, contours[-2], -1, (255, 0, 0), 1)
            cv2.imshow("eye", eye)
            cv2.waitKey(0)
            points = np.vstack((points, [bbox[0][0] + x, bbox[0][1] + y]))  # absolute coordinates
        except (IndexError, ZeroDivisionError):
            points = np.vstack((points, [int((bbox[0][0] + bbox[3][0])/2), int((bbox[0][1] + bbox[3][1])/2)]))"""

        """iris = cv2.HoughCircles(eye, cv2.HOUGH_GRADIENT, 1, 100,
                                     param1=20,
                                     param2=10,
                                     minRadius=5,
                                     maxRadius=20)
        if iris is not None:
            iris = np.uint16(np.around(iris))
            eye = cv2.cvtColor(eye, cv2.COLOR_GRAY2BGR)
            cv2.circle(eye, (iris[0][0][0], iris[0][0][1]), iris[0][0][2], (0, 255, 0), 1)
            cv2.circle(eye, (iris[0][0][0], iris[0][0][1]), 1, (0, 0, 255), -1)
            cv2.imshow("eye", eye)
            cv2.waitKey(0)
            points = np.vstack((points, [bbox[0][0] + iris[0][0][0], bbox[0][1] + iris[0][0][1]])) #add eye centerpoint
        else:
            points = np.vstack((points, [int((bbox[0][0] + bbox[3][0])/2), int((bbox[0][1] + bbox[3][1])/2)]))
            #add center of bbox"""

    return points

def warp(src_img, src_points, src_bbox, input_img, input_points):
    result_img = np.zeros(src_img.shape, dtype=src_img.dtype)
    #result_img = src_img.copy()

    src_delaunay = Delaunay(src_points)  # create Delaunay triangles to warp to

    triangle_affines = np.array(list(get_affine_transform(src_delaunay.simplices, input_points, src_points)))
    # create transform matrices to warp input points to source triangles

    src_bbox_points = np.array([(x, y) for x in range(src_bbox[0][0], src_bbox[3][0] + 1)
                                for y in range(src_bbox[0][1], src_bbox[3][1] + 1)])
    # create an array of all coordinates in source face area

    src_indicies = src_delaunay.find_simplex(src_bbox_points)  # returns triangle index for each point, -1 for none

    """lefteye_points = src_points[36:41]
    lefteye_Delaunay = Delaunay(lefteye_points)
    lefteye_indicies = lefteye_Delaunay.find_simplex(src_bbox_points)
    righteye_points = src_points[42:47]
    righteye_Delaunay = Delaunay(righteye_points)
    righteye_indicies = righteye_Delaunay.find_simplex(src_bbox_points)"""
    mouth_points = src_points[60:67]
    mouth_Delaunay = Delaunay(mouth_points)
    mouth_indicies = mouth_Delaunay.find_simplex(src_bbox_points)

    for index in range(len(src_indicies)):
        if (mouth_indicies[index] != -1):  # (lefteye_indicies[index] != -1) or (righteye_indicies[index] != -1) or
            src_indicies[index] = -1

    for triangle_index in range(len(src_delaunay.simplices)):  # for each triangle
        triangle_points = src_bbox_points[src_indicies == triangle_index]  # for the points in the triangle
        num_points = len(triangle_points)  # get the number of points
        out_points = np.dot(triangle_affines[triangle_index], np.vstack((triangle_points.T, np.ones(num_points))))
        # perform affine transform T = M.[x,y,1]^T to create triangle of input in the source

        x, y = triangle_points.T  # transpose [[x1,y1], [x2,y2], ...] to [x1, x2, ...], [y1, y2, ...] src_triangle coord
        warp_img = cv2.warpAffine(input_img, triangle_affines[triangle_index], (input_img.shape[1], input_img.shape[0]))
        #cv2.imshow("warp_img", warp_img)
        result_img[y, x] = warp_img[y, x]
        #cv2.imshow("result_img", result_img)
        #cv2.waitKey(20)  # these show the process for each section

    return result_img

def get_affine_transform(input_simplices, input_points, src_points):
    for triangle in input_simplices:  # for each triangle
        src_triangle = np.float32(src_points[triangle])
        input_triangle = np.float32(input_points[triangle])
        mat = cv2.getAffineTransform(input_triangle, src_triangle)  # get the transform matrix
        #mat = cv2.getAffineTransform(src_triangle, input_triangle)  # get the transform matrix
        yield mat

def draw(img):  # draws facial points, Delaunay triangles and bounding box

    points, bbox, _ = find(img)

    if points is None or bbox is None:
        print("no face to draw")
        return img

    left_bbox = [[min(points[36:41, 0]), min(points[36:41, 1])], [max(points[36:41, 0]), min(points[36:41, 1])],
                 [min(points[36:41, 0]), max(points[36:41, 1])], [max(points[36:41, 0]), max(points[36:41, 1])]]
    right_bbox = [[min(points[42:47, 0]), min(points[42:47, 1])], [max(points[42:47, 0]), min(points[42:47, 1])],
                  [min(points[42:47, 0]), max(points[42:47, 1])], [max(points[42:47, 0]), max(points[42:47, 1])]]

    bboxs = [bbox, left_bbox, right_bbox]

    for bbox in bboxs:
        cv2.line(img, tuple(bbox[0]), tuple(bbox[1]), (0, 0, 255), 2)
        cv2.line(img, tuple(bbox[0]), tuple(bbox[2]), (0, 0, 255), 2)
        cv2.line(img, tuple(bbox[2]), tuple(bbox[3]), (0, 0, 255), 2)
        cv2.line(img, tuple(bbox[1]), tuple(bbox[3]), (0, 0, 255), 2)

    for point in points:
        cv2.circle(img, (point[0], point[1]), 2, (0, 255, 0), -1)

    triangles = Delaunay(points)

    for triangle in points[triangles.simplices]:
        cv2.line(img, tuple(triangle[0]), tuple(triangle[1]), (255, 255, 255), 1)
        cv2.line(img, tuple(triangle[1]), tuple(triangle[2]), (255, 255, 255), 1)
        cv2.line(img, tuple(triangle[2]), tuple(triangle[0]), (255, 255, 255), 1)

    return img