from settings import ANOMALIES_TOPIC, TRANSACTIONS_CONSUMER_GROUP, TRANSACTIONS_TOPIC
from streaming.utils import create_consumer, create_producer
import pandas as pd
import numpy as np
import logging
import json
from joblib import load

clf = load("models/isolation_forest.joblib")

consumer = create_consumer(topic=TRANSACTIONS_TOPIC, group_id=TRANSACTIONS_CONSUMER_GROUP)

producer = create_producer()

while True:
    message = consumer.poll(timeout=50)
    if message is None:
        continue
    if message.error():
        logging.error("Consumer error: {}".format(message.error()))
        continue

    # Message that came from producer
    record = json.loads(message.value().decode('utf-8'))
    print(record)

    test_data = np.delete(record, 0)
    data = np.array(test_data, dtype=float).reshape(1, -1)

    prediction = clf.predict(data)
    if prediction == -1:
        print("prediction")
        print(prediction)

        score = clf.score_samples(data)
        print("score")
        print(score)

        roundes_score = np.round(score, 3).tolist()
        print("roundes_score")
        print(roundes_score)

        record_bytes = json.dumps(record).encode("utf-8")
        producer.produce(topic=ANOMALIES_TOPIC,
                            value=record_bytes)
        producer.flush()
consumer.flush()


import torch
import torch.nn as nn
import torch.nn.functional as F

import utils
import argparse
import os
import time
from cp_dataset import CPDataset, CPDataLoader
from networks import GMM, UnetGenerator, VGGLoss, load_checkpoint, save_checkpoint
from resnet import Embedder
from unet import UNet, VGGExtractor, Discriminator, AccDiscriminator
from torch.utils.tensorboard import SummaryWriter
from tryon_net import G
import numpy as np
from scipy.spatial import distance
from tqdm import tqdm
from distributed import (
    get_rank,
    synchronize,
    reduce_loss_dict,
    reduce_sum,
    get_world_size,
)

mse_criterion = nn.MSELoss()

def single_gpu_flag(args):
    return not args.distributed or (args.distributed and args.local_rank % torch.cuda.device_count() == 0)

def get_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument("--name", default="residual_resnet")
    parser.add_argument("--gpu_ids", default="")
    parser.add_argument('-j', '--workers', type=int, default=16)
    parser.add_argument('-b', '--batch-size', type=int, default=4)

    parser.add_argument('--local_rank', type=int, default=0, help="gpu to use, used for distributed training")

    parser.add_argument("--use_gan", action='store_true', default=True)
    parser.add_argument("--no_consist", action='store_true')

    parser.add_argument("--dataroot", default="data")
    parser.add_argument("--datamode", default="test")
    parser.add_argument("--stage", default="residual_old")
    parser.add_argument("--data_list", default="train_pairs.txt")
    parser.add_argument("--fine_width", type=int, default=192)
    parser.add_argument("--fine_height", type=int, default=256)
    parser.add_argument("--radius", type=int, default=5)
    parser.add_argument("--grid_size", type=int, default=5)
    parser.add_argument('--lr', type=float, default=0.0001, help='initial learning rate for adam')
    parser.add_argument('--tensorboard_dir', type=str, default='tensorboard', help='save tensorboard infos')
    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints', help='save checkpoint infos')
    parser.add_argument('--checkpoint', type=str, default='', help='model checkpoint for initialization')
    parser.add_argument("--display_count", type=int, default=100)
    parser.add_argument("--save_count", type=int, default=5000)
    parser.add_argument("--keep_step", type=int, default=100000)
    parser.add_argument("--decay_step", type=int, default=100000)
    parser.add_argument("--shuffle", action='store_true', help='shuffle input data')

    opt = parser.parse_args()
    return opt

opt = get_opt()
print(opt)
print("Start to train stage: %s, named: %s!" % (opt.stage, opt.name))

n_gpu = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
opt.distributed = n_gpu > 1
local_rank = opt.local_rank

# create dataset
train_dataset = CPDataset(opt)

# create dataloader
train_loader = CPDataLoader(opt, train_dataset)
data_loader = torch.utils.data.DataLoader(
                train_dataset, batch_size=opt.batch_size, shuffle=False,
                num_workers=opt.workers, pin_memory=True)

gmm_model = GMM(opt)
load_checkpoint(gmm_model, "checkpoints/gmm_train_new/step_020000.pth")
gmm_model.cuda()

generator = UnetGenerator(25, 4, 6, ngf=64, norm_layer=nn.InstanceNorm2d)
load_checkpoint(generator, "checkpoints/tom_train_new_2/step_070000.pth")
generator.cuda()

embedder_model = Embedder()
load_checkpoint(embedder_model, "checkpoints/identity_embedding_for_test/step_045000.pth")
image_embedder = embedder_model.embedder_b.cuda()
prod_embedder = embedder_model.embedder_a.cuda()

model = G()
if not opt.checkpoint == '' and os.path.exists(opt.checkpoint):
    load_checkpoint(model, opt.checkpoint)
model.cuda()

model.eval()
gmm_model.eval()
image_embedder.eval()
generator.eval()
prod_embedder.eval()

pbar = tqdm(enumerate(data_loader), total=len(data_loader))

product_embeddings = []
outfit_embeddings = []
transfer_embeddings = []

product_embeddings_gt = []
outfit_embeddings_gt = []

for i, (inputs, inputs_2) in pbar:

    im = inputs['image'].cuda()
    im_pose = inputs['pose_image']
    im_h = inputs['head']
    shape = inputs['shape']

    agnostic = inputs['agnostic'].cuda()

    c = inputs['cloth'].cuda()
    cm = inputs['cloth_mask'].cuda()
    c_2 = inputs_2['cloth'].cuda()
    cm_2 = inputs_2['cloth_mask'].cuda()

    with torch.no_grad():
        grid, theta = gmm_model(agnostic, c)
        c = F.grid_sample(c, grid, padding_mode='border')
        cm = F.grid_sample(cm, grid, padding_mode='zeros')

        grid_2, theta_2 = gmm_model(agnostic, c_2)
        c_2 = F.grid_sample(c_2, grid_2, padding_mode='border')
        cm_2 = F.grid_sample(cm_2, grid_2, padding_mode='zeros')

        outputs = generator(torch.cat([agnostic, c], 1))
        p_rendered, m_composite = torch.split(outputs, 3, 1)
        p_rendered = F.tanh(p_rendered)
        m_composite = F.sigmoid(m_composite)
        transfer_1 = c * m_composite + p_rendered * (1 - m_composite)

        outputs_2 = generator(torch.cat([agnostic, c_2], 1))
        p_rendered_2, m_composite_2 = torch.split(outputs_2, 3, 1)
        p_rendered_2 = F.tanh(p_rendered_2)
        m_composite_2 = F.sigmoid(m_composite_2)
        transfer_2 = c_2 * m_composite_2 + p_rendered_2 * (1 - m_composite_2)

        gt_residual = (torch.mean(im, dim=1) - torch.mean(transfer_1, dim=1)).unsqueeze(1)

        output_1 = model(transfer_1.detach(), gt_residual.detach())

        embedding_o = image_embedder(output_1).cpu()
        embedding_t = image_embedder(transfer_1).cpu()
        embedding_p = prod_embedder(inputs['cloth'].cuda()).cpu()

        product_embeddings.append(embedding_p)
        outfit_embeddings.append(embedding_o)
        transfer_embeddings.append(embedding_t)

        embedding_o = image_embedder(inputs['image'].cuda()).cpu()
        embedding_p = prod_embedder(inputs_2['cloth'].cuda()).cpu()
        product_embeddings_gt.append(embedding_p)
        outfit_embeddings_gt.append(embedding_o)

product_embeddings = torch.cat(product_embeddings, dim=0).numpy()
outfit_embeddings = torch.cat(outfit_embeddings, dim=0).numpy()
transfer_embeddings = torch.cat(transfer_embeddings, dim=0).numpy()

product_embeddings_gt = torch.cat(product_embeddings_gt, dim=0).numpy()
outfit_embeddings_gt = torch.cat(outfit_embeddings_gt, dim=0).numpy()

def get_correct_count(outfit_embeddings, product_embeddings, top_k=5):
    correct_count = 0

    for i in tqdm(range(outfit_embeddings.shape[0])):

        scores = []
        for j in range(product_embeddings.shape[0]):
            scores.append(distance.euclidean(outfit_embeddings[i], product_embeddings[j]))
        sort_index = np.argsort(scores).tolist()
        if i in sort_index[0:top_k]:
            correct_count += 1
    return correct_count

def get_correct_match_count(outfit_embeddings, product_embeddings, transfer_embeddings, top_k=1):
    correct_count = 0

    for i in tqdm(range(outfit_embeddings.shape[0])):

        scores_o = []
        for j in range(product_embeddings.shape[0]):
            scores_o.append(distance.euclidean(outfit_embeddings[i], product_embeddings[j]))
        sort_index_o = np.argsort(scores_o).tolist()

        scores_t = []
        for j in range(product_embeddings.shape[0]):
            scores_t.append(distance.euclidean(transfer_embeddings[i], product_embeddings[j]))
        sort_index_t = np.argsort(scores_t).tolist()
        correct_count += len(set(sort_index_o[0: top_k]).intersection(set(sort_index_t[0: top_k]))) / top_k
    return correct_count

# def build_tree(vectors, dim=64):
#     a = AnnoyIndex(dim, 'euclidean')
#     for i, v in enumerate(vectors):
#         a.add_item(i, v)
#     a.build(-1)
#     return a
#
# product_tree = build_tree(product_embeddings)

dataset_size = len(train_dataset)

correct_count = get_correct_count(transfer_embeddings, product_embeddings, top_k=1)
print("acc_transfer", correct_count , correct_count / dataset_size)

correct_count = get_correct_match_count(outfit_embeddings, product_embeddings, transfer_embeddings, top_k=1)
print("acc 1", correct_count, correct_count / dataset_size)
correct_count = get_correct_match_count(outfit_embeddings, product_embeddings, transfer_embeddings, top_k=5)
print("acc 5", correct_count, correct_count / dataset_size)


#Based off of: https://github.com/yu4u/age-gender-estimation
import cv2
import numpy as np
try:
    from wide_resnet import WideResNet
except ImportError:
    from .wide_resnet import WideResNet
from keras.backend import backend
from keras.utils.data_utils import get_file
from keras import backend as K
import gc

#https://superuser.com/questions/470664/how-to-download-dropbox-files-using-wget-command
tensorflow_model = "https://www.dropbox.com/s/ukm1dxca0vve6gn/weights.18-4.06.hdf5?dl=1"
tensorflow_file = "weights.18-4.06.hdf5"
theano_model = "https://www.dropbox.com/s/e3q3urdwv4ng2sz/weights_theano.16-3.99.hdf5?dl=1"
theano_file = "weights_theano.16-3.99.hdf5"

if backend() == "tensorflow":
    modurl = tensorflow_model
    modfile = tensorflow_file
    modhash = '89f56a39a78454e96379348bddd78c0d'
    default_mod = get_file(modfile, modurl, cache_subdir="models", file_hash=modhash)
elif backend() == "theano":
    modurl = theano_model
    modfile = theano_file
    modhash = 'd632812726dc28d30db8b1f1c99f20a4'
    default_mod = get_file(modfile, modurl, cache_subdir="models", file_hash=modhash)
else:
    default_mod = None

def run(img, mod = default_mod, img_size = 64, depth = 16, k = 8):
    K.clear_session()
    gc.collect()
    model = WideResNet(img_size, depth=depth, k=k)()
    model.load_weights(mod)
    faces = np.empty((1, img_size, img_size, 3))
    faces[0,:,:,:] = cv2.resize(img, (img_size, img_size))
    results = model.predict(faces)
    predicted_genders = results[0][0][0]
    ages = np.arange(0, 101).reshape(101, 1)
    predicted_ages = results[1].dot(ages).flatten()[0]
    return (predicted_genders, predicted_ages)

#TODO: Document the run() function and make another version
#for lists of faces


import os
import sys
import glob
import time
import copy
import random
import numpy as np
import utils
import logging
import argparse
from model import NASNetworkCIFAR
# from model_search import NASWSNetworkCIFAR, NASWSNetworkImageNet
from nao import NAO

# Basic model parameters.
parser = argparse.ArgumentParser(description='NAO CIFAR-10')
parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
parser.add_argument('--data', type=str, default='./data')
parser.add_argument('--dataset', type=str, default='cifar10', choices=['cifar10, cifar100, imagenet'])
parser.add_argument('--zip_file', action='store_true', default=False)
parser.add_argument('--lazy_load', action='store_true', default=False)
parser.add_argument('--output_dir', type=str, default='models')
parser.add_argument('--seed', type=int, default=0)
parser.add_argument('--child_batch_size', type=int, default=64)
parser.add_argument('--child_eval_batch_size', type=int, default=500)
parser.add_argument('--child_epochs', type=int, default=150)
parser.add_argument('--child_layers', type=int, default=3)
parser.add_argument('--child_nodes', type=int, default=5)
parser.add_argument('--child_channels', type=int, default=20)
parser.add_argument('--child_cutout_size', type=int, default=None)
parser.add_argument('--child_grad_bound', type=float, default=5.0)
parser.add_argument('--child_lr_max', type=float, default=0.025)
parser.add_argument('--child_lr_min', type=float, default=0.001)
parser.add_argument('--child_keep_prob', type=float, default=1.0)
parser.add_argument('--child_drop_path_keep_prob', type=float, default=0.9)
parser.add_argument('--child_l2_reg', type=float, default=3e-4)
parser.add_argument('--child_use_aux_head', action='store_true', default=False)
parser.add_argument('--child_eval_epochs', type=str, default='30')
parser.add_argument('--child_arch_pool', type=str, default=None)
parser.add_argument('--child_lr', type=float, default=0.1)
parser.add_argument('--child_label_smooth', type=float, default=0.1, help='label smoothing')
parser.add_argument('--child_gamma', type=float, default=0.97, help='learning rate decay')
parser.add_argument('--child_decay_period', type=int, default=1, help='epochs between two learning rate decays')
parser.add_argument('--controller_seed_arch', type=int, default=600)
parser.add_argument('--controller_expand', type=int, default=None)
parser.add_argument('--controller_new_arch', type=int, default=300)
parser.add_argument('--controller_encoder_layers', type=int, default=1)
parser.add_argument('--controller_encoder_hidden_size', type=int, default=96)
parser.add_argument('--controller_encoder_emb_size', type=int, default=48)
parser.add_argument('--controller_mlp_layers', type=int, default=3)
parser.add_argument('--controller_mlp_hidden_size', type=int, default=200)
parser.add_argument('--controller_decoder_layers', type=int, default=1)
parser.add_argument('--controller_decoder_hidden_size', type=int, default=96)
parser.add_argument('--controller_source_length', type=int, default=40)
parser.add_argument('--controller_encoder_length', type=int, default=20)
parser.add_argument('--controller_decoder_length', type=int, default=40)
parser.add_argument('--controller_encoder_dropout', type=float, default=0)
parser.add_argument('--controller_mlp_dropout', type=float, default=0.1)
parser.add_argument('--controller_decoder_dropout', type=float, default=0)
parser.add_argument('--controller_l2_reg', type=float, default=1e-4)
parser.add_argument('--controller_encoder_vocab_size', type=int, default=12)
parser.add_argument('--controller_decoder_vocab_size', type=int, default=12)
parser.add_argument('--controller_trade_off', type=float, default=0.8)
parser.add_argument('--controller_epochs', type=int, default=1000)
parser.add_argument('--controller_batch_size', type=int, default=100)
parser.add_argument('--controller_lr', type=float, default=0.001)
parser.add_argument('--controller_optimizer', type=str, default='adam')
parser.add_argument('--controller_grad_bound', type=float, default=5.0)
args = parser.parse_args()

def get_builder(dataset):
    if dataset == 'cifar10':
        return build_cifar10
    # elif dataset == 'cifar100':
    #     return build_cifar100
    # else:
    #     return build_imagenet

def build_cifar10(model_state_dict=None, optimizer_state_dict=None, **kwargs):
    epoch = kwargs.pop('epoch')
    ratio = kwargs.pop('ratio')
    train_ds, valid_ds = utils.load_cifar10(args.child_batch_size, args.child_cutout_size) # subset random sample

    model = NASWSNetworkCIFAR(10, args.child_layers, args.child_nodes, args.child_channels, args.child_keep_prob,
                              args.child_drop_path_keep_prob,
                              args.child_use_aux_head, args.steps)
    model = model.cuda()
    train_criterion = nn.CrossEntropyLoss().cuda()
    eval_criterion = nn.CrossEntropyLoss().cuda()
    logging.info("param size = %fMB", utils.count_parameters_in_MB(model))

    optimizer = torch.optim.SGD(
        model.parameters(),
        args.child_lr_max,
        momentum=0.9,
        weight_decay=args.child_l2_reg,
    )
    if model_state_dict is not None:
        model.load_state_dict(model_state_dict)
    if optimizer_state_dict is not None:
        optimizer.load_state_dict(optimizer_state_dict)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.child_epochs, args.child_lr_min, epoch)
    return train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler

def main():
    random.seed(args.seed)
    np.random.seed(args.seed)

    args.steps = int(np.ceil(45000 / args.child_batch_size)) * args.child_epochs

    logging.info("args = %s", args)

    if args.child_arch_pool is not None:
        logging.info('Architecture pool is provided, loading')
        with open(args.child_arch_pool) as f:
            archs = f.read().splitlines()
            archs = list(map(utils.build_arch, archs))
            child_arch_pool = archs
    elif os.path.exists(os.path.join(args.output_dir, 'arch_pool')):
        logging.info('Architecture pool is founded, loading')
        with open(os.path.join(args.output_dir, 'arch_pool')) as f:
            archs = f.read().splitlines()
            archs = list(map(utils.build_arch, archs))
            child_arch_pool = archs
    else:
        child_arch_pool = None

    child_eval_epochs = eval(args.child_eval_epochs)
    build_fn = get_builder(args.dataset)
    train_queue, valid_queue, model, train_criterion, eval_criterion, optimizer, scheduler = build_fn(ratio=0.9, epoch=-1)

    nao = NAO(
        args.controller_encoder_layers,
        args.controller_encoder_vocab_size,
        args.controller_encoder_hidden_size,
        args.controller_encoder_dropout,
        args.controller_encoder_length,
        args.controller_source_length,
        args.controller_encoder_emb_size,
        args.controller_mlp_layers,
        args.controller_mlp_hidden_size,
        args.controller_mlp_dropout,
        args.controller_decoder_layers,
        args.controller_decoder_vocab_size,
        args.controller_decoder_hidden_size,
        args.controller_decoder_dropout,
        args.controller_decoder_length,
    )
    nao = nao.cuda()
    logging.info("Encoder-Predictor-Decoder param size = %fMB", utils.count_parameters_in_MB(nao))

    # Train child model
    if child_arch_pool is None:
        logging.info('Architecture pool is not provided, randomly generating now')
        child_arch_pool = utils.generate_arch(args.controller_seed_arch, args.child_nodes, 5)  # [[[conv],[reduc]]]
    child_arch_pool_prob = None

    eval_points = utils.generate_eval_points(child_eval_epochs, 0, args.child_epochs)
    step = 0
    for epoch in range(1, args.child_epochs + 1):
        scheduler.step()
        lr = scheduler.get_lr()[0]
        logging.info('epoch %d lr %e', epoch, lr)
        # sample an arch to train
        train_acc, train_obj, step = child_train(train_queue, model, optimizer, step, child_arch_pool, child_arch_pool_prob, train_criterion)
        logging.info('train_acc %f', train_acc)

        if epoch not in eval_points:
            continue
        # Evaluate seed archs
        valid_accuracy_list = child_valid(valid_queue, model, child_arch_pool, eval_criterion)

        # Output archs and evaluated error rate
        old_archs = child_arch_pool
        old_archs_perf = valid_accuracy_list

        old_archs_sorted_indices = np.argsort(old_archs_perf)[::-1]
        old_archs = [old_archs[i] for i in old_archs_sorted_indices]
        old_archs_perf = [old_archs_perf[i] for i in old_archs_sorted_indices]
        with open(os.path.join(args.output_dir, 'arch_pool.{}'.format(epoch)), 'w') as fa:
            with open(os.path.join(args.output_dir, 'arch_pool.perf.{}'.format(epoch)), 'w') as fp:
                with open(os.path.join(args.output_dir, 'arch_pool'), 'w') as fa_latest:
                    with open(os.path.join(args.output_dir, 'arch_pool.perf'), 'w') as fp_latest:
                        for arch, perf in zip(old_archs, old_archs_perf):
                            arch = ' '.join(map(str, arch[0] + arch[1]))
                            fa.write('{}\n'.format(arch))
                            fa_latest.write('{}\n'.format(arch))
                            fp.write('{}\n'.format(perf))
                            fp_latest.write('{}\n'.format(perf))

        if epoch == args.child_epochs:
            break

        # Train Encoder-Predictor-Decoder
        logging.info('Training Encoder-Predictor-Decoder')
        encoder_input = list \
            (map(lambda x: utils.parse_arch_to_seq(x[0], 2) + utils.parse_arch_to_seq(x[1], 2), old_archs))
        # [[conv, reduc]]
        min_val = min(old_archs_perf)
        max_val = max(old_archs_perf)
        encoder_target = [(i - min_val) / (max_val - min_val) for i in old_archs_perf]

        if args.controller_expand:
            dataset = list(zip(encoder_input, encoder_target))
            n = len(dataset)
            ratio = 0.9
            split = int( n *ratio)
            np.random.shuffle(dataset)
            encoder_input, encoder_target = list(zip(*dataset))
            train_encoder_input = list(encoder_input[:split])
            train_encoder_target = list(encoder_target[:split])
            valid_encoder_input = list(encoder_input[split:])
            valid_encoder_target = list(encoder_target[split:])
            for _ in range(args.controller_expan d -1):
                for src, tgt in zip(encoder_input[:split], encoder_target[:split]):
                    a = np.random.randint(0, args.child_nodes)
                    b = np.random.randint(0, args.child_nodes)
                    src = src[:4 * a] + src[4 * a + 2:4 * a + 4] + \
                          src[4 * a:4 * a + 2] + src[4 * (a + 1):20 + 4 * b] + \
                          src[20 + 4 * b + 2:20 + 4 * b + 4] + src[20 + 4 * b:20 + 4 * b + 2] + \
                          src[20 + 4 * (b + 1):]
                    train_encoder_input.append(src)
                    train_encoder_target.append(tgt)
        else:
            train_encoder_input = encoder_input
            train_encoder_target = encoder_target
            valid_encoder_input = encoder_input
            valid_encoder_target = encoder_target
        logging.info('Train data: {}\tValid data: {}'.format(len(train_encoder_input), len(valid_encoder_input)))

        nao_train_dataset = utils.NAODataset(train_encoder_input, train_encoder_target, True, swap=True if args.controller_expand is None else False)
        nao_valid_dataset = utils.NAODataset(valid_encoder_input, valid_encoder_target, False)
        nao_train_queue = torch.utils.data.DataLoader(
            nao_train_dataset, batch_size=args.controller_batch_size, shuffle=True, pin_memory=True)
        nao_valid_queue = torch.utils.data.DataLoader(
            nao_valid_dataset, batch_size=args.controller_batch_size, shuffle=False, pin_memory=True)
        nao_optimizer = torch.optim.Adam(nao.parameters(), lr=args.controller_lr, weight_decay=args.controller_l2_reg)
        for nao_epoch in range(1, args.controller_epoch s +1):
            nao_loss, nao_mse, nao_ce = nao_train(nao_train_queue, nao, nao_optimizer)
            logging.info("epoch %04d train loss %.6f mse %.6f ce %.6f", nao_epoch, nao_loss, nao_mse, nao_ce)
            if nao_epoch % 100 == 0:
                pa, hs = nao_valid(nao_valid_queue, nao)
                logging.info("Evaluation on valid data")
                logging.info('epoch %04d pairwise accuracy %.6f hamming distance %.6f', epoch, pa, hs)

        # Generate new archs
        new_archs = []
        max_step_size = 50
        predict_step_size = 0
        top100_archs = list \
            (map(lambda x: utils.parse_arch_to_seq(x[0], 2) + utils.parse_arch_to_seq(x[1], 2), old_archs[:100]))
        nao_infer_dataset = utils.NAODataset(top100_archs, None, False)
        nao_infer_queue = torch.utils.data.DataLoader(
            nao_infer_dataset, batch_size=len(nao_infer_dataset), shuffle=False, pin_memory=True)
        while len(new_archs) < args.controller_new_arch:
            predict_step_size += 1
            logging.info('Generate new architectures with step size %d', predict_step_size)
            new_arch = nao_infer(nao_infer_queue, nao, predict_step_size, direction='+')
            for arch in new_arch:
                if arch not in encoder_input and arch not in new_archs:
                    new_archs.append(arch)
                if len(new_archs) >= args.controller_new_arch:
                    break
            logging.info('%d new archs generated now', len(new_archs))
            if predict_step_size > max_step_size:
                break
                # [[conv, reduc]]
        new_archs = list(map(lambda x: utils.parse_seq_to_arch(x, 2), new_archs))  # [[[conv],[reduc]]]
        num_new_archs = len(new_archs)
        logging.info("Generate %d new archs", num_new_archs)
        # replace bottom archs
        new_arch_pool = old_archs[:len(old_archs) - num_new_archs] + new_archs
        logging.info("Totally %d architectures now to train", len(new_arch_pool))

        child_arch_pool = new_arch_pool
        with open(os.path.join(args.output_dir, 'arch_pool'), 'w') as f:
            for arch in new_arch_pool:
                arch = ' '.join(map(str, arch[0] + arch[1]))
                f.write('{}\n'.format(arch))

        child_arch_pool_prob = None

    logging.info('Finish Searching')
    logging.info('Reranking top 5 architectures')
    # reranking top 5
    top_archs = old_archs[:5]
    if args.dataset == 'cifar10':
        top_archs_perf = train_and_evaluate_top_on_cifar10(top_archs, train_queue, valid_queue)
    elif args.dataset == 'cifar100':
        top_archs_perf = train_and_evaluate_top_on_cifar100(top_archs, train_queue, valid_queue)
    else:
        top_archs_perf = train_and_evaluate_top_on_imagenet(top_archs, train_queue, valid_queue)
    top_archs_sorted_indices = np.argsort(top_archs_perf)[::-1]
    top_archs = [top_archs[i] for i in top_archs_sorted_indices]
    top_archs_perf = [top_archs_perf[i] for i in top_archs_sorted_indices]
    with open(os.path.join(args.output_dir, 'arch_pool.final'), 'w') as fa:
        with open(os.path.join(args.output_dir, 'arch_pool.perf.final'), 'w') as fp:
            for arch, perf in zip(top_archs, top_archs_perf):
                arch = ' '.join(map(str, arch[0] + arch[1]))
                fa.write('{}\n'.format(arch))
                fp.write('{}\n'.format(perf))


import numpy as np
import glob
import random
import sys

class record_fetcher(object):
    '''
    creates feature arrays and labels from raw accelerometer/demographic data
    splits features and labels between subjects into test/train sets

    methods:
      - fetch(batch_size,binary,seed)
        separates raw accelerometer/demographic data into windows and creates
        input features and labels for conv-lstm classifier
        parameters:
          - batch_size: integer
            number of frames to use for each train/test instance
            e.g. 1000 means each test/train instance represents 10 seconds of data
          - binary: boolean (default True)
            use True to set labels for ambulatory/non-ambulatory
            use False to set labels for non-ambulatory/walking/running/upstairs/downstairs
          - seed: integer (default None)
            (optional) seed to use for random test/train splitting
        outputs:
          - numpy array representing training raw accelerometer and demographic data
            dimension 0 is the index of the batch window
            dimension 1 is the number of filter maps (1 filter map per instance)
            dimension 2 is the time frame of the accelerometer data within the batch window
            dimension 3 is the raw accelerometer and demographic data for the specified time frame
          - numpy array representing testing raw accelerometer and demographic data
            dimension 0 is the index of the batch window
            dimension 1 is the number of filter maps (1 filter map per instance)
            dimension 2 is the time frame of the accelerometer data within the batch window
            dimension 3 is the raw accelerometer and demographic data for the specified time frame
          - numpy array representing training activity label over each time window
          - numpy array representing testing activity label over each time window
    '''

    def __init__(self):

        #collect all valid subject ids
        self.subjects = [102,103,105,106,107,108,110,112,113,114,115,116,117,118,119,120,\
            121,122,123,124,125,126,127,128,129,130,131,132,133,134,136,137,138,139,140,\
            142,143,144,146,148,149,150,151,152,153,154,155,156,157,159,160,161,162,163,\
            164,165,166,169,170,171,172,173,174,175,177,178,179,180,181,182,183,184,185,\
            186,187,188,189,190,191,192]

        #categorize activity ids into ambulatory/non-ambulatory
        self.dic1 = {
        'ambulatory': [11,12,13,14,23,24,25,26,27,28,29,30,31,32,16,17,18,33,34],
        'nonambulatory': [19,20,21,22]
        }

        #categorize activity ids into non-ambulatory/walking/running/upstairs/downstairs
        self.dic2 = {
        'nonambulatory': [19,20,21,22],
        'walking': [11,12,13,14,23,24,25,26,27,28,29,30,31,32],
        'running': [16,17,18],
        'upstairs': [33],
        'downstairs': [34]
        }

        #get filenames for all activity arrays
        self.ambulatory = []
        for i in self.dic1['ambulatory']:
            self.ambulatory.extend(glob.glob('../data/arrays/*_%i_*' % i))

        self.nonambulatory = []
        for i in self.dic1['nonambulatory']:
            self.nonambulatory.extend(glob.glob('../data/arrays/*_%i_*' % i))

        self.walking = []
        for i in self.dic2['walking']:
            self.walking.extend(glob.glob('../data/arrays/*_%i_*' % i))

        self.running = []
        for i in self.dic2['running']:
            self.running.extend(glob.glob('../data/arrays/*_%i_*' % i))

        self.upstairs = []
        for i in self.dic2['upstairs']:
            self.upstairs.extend(glob.glob('../data/arrays/*_%i_*' % i))

        self.downstairs = []
        for i in self.dic2['downstairs']:
            self.downstairs.extend(glob.glob('../data/arrays/*_%i_*' % i))

    def fetch(self,batch_size,binary=True,seed=None):
        '''
        separates raw accelerometer/demographic data into windows and creates
        input features and labels for conv-lstm classifier
        parameters:
          - batch_size: integer
            number of frames to use for each train/test instance
            e.g. 1000 means each test/train instance represents 10 seconds of data
          - binary: boolean (default True)
            use True to set labels for ambulatory/non-ambulatory
            use False to set labels for non-ambulatory/walking/running/upstairs/downstairs
          - seed: integer (default None)
            (optional) seed to use for random test/train splitting
        outputs:
          - numpy array representing training raw accelerometer and demographic data
            dimension 0 is the index of the batch window
            dimension 1 is the number of filter maps (1 filter map per instance)
            dimension 2 is the time frame of the accelerometer data within the batch window
            dimension 3 is the raw accelerometer and demographic data for the specified time frame
          - numpy array representing testing raw accelerometer and demographic data
            dimension 0 is the index of the batch window
            dimension 1 is the number of filter maps (1 filter map per instance)
            dimension 2 is the time frame of the accelerometer data within the batch window
            dimension 3 is the raw accelerometer and demographic data for the specified time frame
          - numpy array representing training activity label over each time window
          - numpy array representing testing activity label over each time window
        '''
        np.random.seed(seed)
        X_test_subjects = np.random.choice(self.subjects,6)

        X_train_list = []
        y_train_list = []
        X_test_list = []
        y_test_list = []

        #for ambulatory/non-ambulatory classification
        if binary:
            for a in self.ambulatory:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([0,1]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([0,1]))

            for a in self.nonambulatory:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([1,0]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([1,0]))

        #for non-ambulatory/walking/running/upstairs/downstairs classification
        else:
            for a in self.nonambulatory:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([1,0,0,0,0]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([1,0,0,0,0]))

            for a in self.walking:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([0,1,0,0,0]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([0,1,0,0,0]))

            for a in self.running:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([0,0,1,0,0]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([0,0,1,0,0]))

            for a in self.upstairs:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([0,0,0,1,0]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([0,0,0,1,0]))

            for a in self.downstairs:
                print 'processing %s' % a
                array = np.load(a)

                #avoid arrays smaller than batch_size
                if array.shape[0] <= batch_size:
                    continue

                #separate array into batches
                seg = array.shape[0]//batch_size
                for i in range(seg):
                    batch = array[i*batch_size:(i+1)*batch_size,:]

                    #if subject id in test set, save features and labels to test set
                    if int(a[15:18]) in X_test_subjects:
                        X_test_list.append(np.expand_dims(batch,0))
                        y_test_list.append(np.array([0,0,0,0,1]))

                    #else save features and labels to train set
                    else:
                        X_train_list.append(np.expand_dims(batch,0))
                        y_train_list.append(np.array([0,0,0,0,1]))

        #pair training X/y together and shuffle
        print 'shuffling records'
        Xy = zip(X_train_list,y_train_list)
        random.shuffle(Xy)

        #separate training X from y
        X_train = np.array([record[0] for record in Xy])
        y_train = np.array([record[1] for record in Xy])

        print 'feature vector shape:', X_train.shape
        print 'label vector shape:', y_train.shape

        #pair testing X/y together and shuffle
        Xy = zip(X_test_list,y_test_list)
        random.shuffle(Xy)

        #separate testing X from y
        X_test = np.array([record[0] for record in Xy])
        y_test = np.array([record[1] for record in Xy])

        print 'feature vector shape:', X_test.shape
        print 'label vector shape:', y_test.shape

        return X_train, X_test, y_train, y_test

if __name__ == "__main__":

    # verify the required arguments are given
    if (len(sys.argv) < 2):
        print 'Usage: python record_fetcher_between_subject_raw.py <1 for 2-category labels, 0 for 5-category labels>'
        exit(1)

    if sys.argv[1] == '1':
        binary = True
    elif sys.argv[1] == '0':
        binary = False
    else:
        print 'Usage: python record_fetcher_between_subject_raw.py <1 for 2-category labels, 0 for 5-category labels>'
        exit(1)

    rf = record_fetcher()
    X_train,X_test,y_train,y_test = rf.fetch(1000,binary=binary,seed=1)
    np.save('X_train_raw',X_train)
    np.save('X_test_raw',X_test)
    np.save('y_train_raw',y_train)
    np.save('y_test_raw',y_test)


import torch
import numpy as np

print(torch.from_numpy(np.random.dirichlet(np.ones(10) * 0.01, 5)))

'''
if current_index > 5:
    print('arena')
    start = datetime.datetime.now()
    with Manager() as manager:
        output_list = manager.list()
        processes = []
        for i in range(2):
            p = Process(target=arena, args=(agent, net, current_index - 1, output_list,))
            p.start()
            processes.append(p)
        for p in processes:
            p.join()

    win, lose, draw = 0, 0, 0
    for i in output_list:
        if i == 1:
            win += 1
        elif i == -1:
            lose += 1
        else:
            draw += 1

    sum = win + lose
    win_rate = float(win) / sum
    print('win:', win, 'lose:', lose, 'draw:', draw, 'win rate:', win_rate)
    if sum > 0 and win_rate >= 0.55:
        torch.save(self.model.state_dict(), 'models/' + agent.name + '_' + str(current_index) + '.pt')
        print('save', current_index)
        current_index += 1
    else:
        model2.load_state_dict(torch.load('models/' + agent.name + '_' + str(current_index - 1) + '.pt'))
    print(datetime.datetime.now() - start)
else:
    torch.save(self.model.state_dict(), 'models/' + agent.name + '_' + str(current_index) + '.pt')
    print('save', current_index)
    current_index += 1
'''


import numpy as np
np.set_printoptions(threshold=np.inf)
import os
import re
import scipy.io as sio
import math
from CPSC2019_challenge import *
from tqdm import tqdm
import keras.backend.tensorflow_backend as KTF
from keras.models import model_from_yaml,load_model
import tensorflow as tf

from net import build_model
from keras.layers import Input
from keras.models import Model

os.environ["CUDA_VISIBLE_DEVICES"] = '0'   #指定第一块GPU可用
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.6  # 程序最多只能占用指定gpu50%的显存
config.gpu_options.allow_growth = True      #程序按需申请内存
sess = tf.Session(config = config)
# 设置session
KTF.set_session(sess)

model1 = model_from_yaml(open('model_architecture_1.yaml').read())
model1.load_weights('model_weights_1.h5')

model2 = model_from_yaml(open('model_architecture_2.yaml').read())
model2.load_weights('model_weights_2.h5')

input_layer = Input((1280, 1))
output_layer = build_model(input_layer=input_layer,start_neurons=16)
model3 = Model(input_layer, output_layer)
model3.load_weights('model_weights_3.h5')

model4 = model_from_yaml(open('model_architecture_4.yaml').read())
model4.load_weights('model_weights_4.h5')

model = []
model.append(model1)
model.append(model2)
model.append(model3)
model.append(model4)

def load_ans(data_path_, rpos_path_, fs_):
    '''
    Please modify this function when you have to load model or any other parameters in CPSC2019_challenge()
    '''
    def is_mat(l):
        return l.endswith('.mat')
    ecg_files = list(filter(is_mat, os.listdir(data_path_)))
    rpos_files = list(filter(is_mat, os.listdir(rpos_path_)))
    HR_ref = []
    R_ref = []
    HR_ans = []
    R_ans = []
    for rpos_file in tqdm(rpos_files):
        index = re.split('[_.]', rpos_file)[1]
        ecg_file = 'data_' + index + '.mat'

        ref_path = os.path.join(rpos_path_, rpos_file)
        ecg_path = os.path.join(data_path_, ecg_file)

        ecg_data = np.transpose(sio.loadmat(ecg_path)['ecg'])[0]
        ecg_data = signal.resample(ecg_data, fs_*10)
        ecg_data = med_filt_1d(ecg_data)

        r_ref = sio.loadmat(ref_path)['R_peak'].flatten()
        r_ref = (r_ref/500*fs_).astype(int)
        r_ref = r_ref[(r_ref >= 0.5*fs_) & (r_ref <= 9.5*fs_)]

        r_hr = np.array([loc for loc in r_ref if (loc > 5.5 * fs_ and loc < len(ecg_data) - 0.5 * fs_)])
        hr_ans, r_ans = CPSC2019_challenge(model,ecg_data,fs_)

        HR_ref.append(round( 60 * fs_ / np.mean(np.diff(r_hr))))
        R_ref.append(r_ref)
        HR_ans.append(hr_ans)
        R_ans.append(r_ans)

    return R_ref, HR_ref, R_ans, HR_ans

def score(r_ref, hr_ref, r_ans, hr_ans, fs_, thr_):
    HR_score = 0
    record_flags = np.ones(len(r_ref))
    for i in range(len(r_ref)):
        FN = 0
        FP = 0
        TP = 0

        if math.isnan(hr_ans[i]):
            hr_ans[i] = 0
        hr_der = abs(int(hr_ans[i]) - int(hr_ref[i]))
        if hr_der <= 0.02 * hr_ref[i]:
            HR_score = HR_score + 1
        elif hr_der <= 0.05 * hr_ref[i]:
            HR_score = HR_score + 0.75
        elif hr_der <= 0.1 * hr_ref[i]:
            HR_score = HR_score + 0.5
        elif hr_der <= 0.2 * hr_ref[i]:
            HR_score = HR_score + 0.25

        for j in range(len(r_ref[i])):
            loc = np.where(np.abs(r_ans[i] - r_ref[i][j]) <= thr_*fs_)[0]
            if j == 0:
                err = np.where((r_ans[i] >= 0.5*fs_ + thr_*fs_) & (r_ans[i] <= r_ref[i][j] - thr_*fs_))[0]
            elif j == len(r_ref[i])-1:
                err = np.where((r_ans[i] >= r_ref[i][j]+thr_*fs_) & (r_ans[i] <= 9.5*fs_ - thr_*fs_))[0]
            else:
                err = np.where((r_ans[i] >= r_ref[i][j]+thr_*fs_) & (r_ans[i] <= r_ref[i][j+1]-thr_*fs_))[0]

            FP = FP + len(err)
            if len(loc) >= 1:
                TP += 1
                FP = FP + len(loc) - 1
            elif len(loc) == 0:
                FN += 1

        if FN + FP > 1:
            record_flags[i] = 0
        elif FN == 1 and FP == 0:
            record_flags[i] = 0.3
        elif FN == 0 and FP == 1:
            record_flags[i] = 0.7

    rec_acc = round(np.sum(record_flags) / len(r_ref), 4)
    hr_acc = round(HR_score / len(r_ref), 4)

    print( 'QRS_acc: {}'.format(rec_acc))
    print('HR_acc: {}'.format(hr_acc))
    print('Scoring complete.')

    # print("Score 0   :",len(record_flags[record_flags==0.0]))
    # print("Score 0.3 :",len(record_flags[record_flags==0.3]))
    # print("Score 0.7 :",len(record_flags[record_flags==0.7]))
    # print("Score 1   :",len(record_flags[record_flags==1]))

    return rec_acc, hr_acc

if __name__ == '__main__':
    FS1 = 500
    FS2 = 128
    THR = 0.075

    DATA_PATH = './data/'
    RPOS_PATH = './ref/'

    DATA_PATH='..\\train\\data\\'
    RPOS_PATH='..\\train\\ref\\'

    R_ref, HR_ref, R_ans, HR_ans = load_ans(DATA_PATH, RPOS_PATH, FS2)
    rec_acc, hr_acc = score(R_ref, HR_ref, R_ans, HR_ans, FS2, THR)

    with open('score.txt', 'w') as score_file:
        print('Total File Number: %d\n' %(np.shape(HR_ans)[0]), file=score_file)
        print('R Detection Acc: %0.4f' %rec_acc, file=score_file)
        print('HR Detection Acc: %0.4f' %hr_acc, file=score_file)

        score_file.close()


# -*- coding: utf-8 -*-
import datetime
import logging
import os
import random
import time
from tkinter import END

import cv2
import numpy
import numpy as np
import pyautogui
from PIL import ImageGrab
from matplotlib import pyplot as plt

pyautogui.FAILSAFE = False
logging.basicConfig(format="%(asctime)s :%(levelname)s:%(message)s", datefmt="%d-%M-%Y %H:%M:%S", level=logging.DEBUG)
# 初始化SIFT探测器
SIFT = cv2.xfeatures2d.SIFT_create()

def ComputeScreenShot(screenShot):
    """
    由于屏幕分辨率高，计算耗时，这里优化一下
    :return:
    """
    kp2, des2 = SIFT.detectAndCompute(screenShot, None)
    return kp2, des2

def GetLocation(target, kp2, des2):
    """
    获取目标图像在截图中的位置
    :param target:
    :param screenShot:
    :return: 返回坐标(x,y) 与opencv坐标系对应
    """
    MIN_MATCH_COUNT = 10
    img1 = target  # cv2.cvtColor(target,cv2.COLOR_BGR2GRAY)# 查询图片
    # img2 = screenShot
    # img2 = cv2.cvtColor(screenShot, cv2.COLOR_BGR2GRAY)  # 训练图片
    # img2 = cv2.resize(img2, dsize=None, fx=0.5, fy=0.5, interpolation=cv2.INTER_NEAREST)
    # 用SIFT找到关键点和描述符

    kp1, des1 = SIFT.detectAndCompute(img1, None)

    FLANN_INDEX_KDTREE = 0
    index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=4)
    search_params = dict(checks=50)

    flann = cv2.FlannBasedMatcher(index_params, search_params)
    matches = flann.knnMatch(des1, des2, k=2)
    good = []
    for m, n in matches:
        if m.distance < 0.7 * n.distance:
            good.append(m)
    if len(good) > MIN_MATCH_COUNT:
        src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
        M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
        matchesMask = mask.ravel().tolist()
        h, w = img1.shape
        pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
        if M is not None:
            dst = cv2.perspectiveTransform(pts, M)
            arr = np.int32(dst)  #
            midPosArr = arr[0] + (arr[2] - arr[0]) // 2
            midPos = (midPosArr[0][0], midPosArr[0][1])
            # show=cv2.circle(img2,midPos,30,(255,255,255),thickness=5)
            # cv2.imshow('s',show)
            # cv2.waitKey()
            # cv2.destroyAllWindows()
            return midPos
        else:
            return None
    else:
        return None

def CheatPos(originPos, factor=5):
    """
    对原始点击坐标进行随机偏移，防止封号
    :param originPos:原始坐标
    :return:
    """
    x, y = random.randint(-factor, factor), random.randint(-factor, factor)
    newPos = (originPos[0] + x, originPos[1] + y)
    return newPos

def Click(targetPosition):
    """
    点击屏幕上的某个点
    :param targetPosition:
    :return:
    """
    if targetPosition is None:
        print('未检测到目标')
    else:

        pyautogui.moveTo(targetPosition, duration=0.20)
        pyautogui.click()
        time.sleep(random.randint(500, 1000) / 1000)

        # time.sleep(random.randint(100, 150) / 1000)

def loadImgs():
    """
    加载所有需要检测的目标图像
    :return:
    """
    obj = {}
    path = os.getcwd() + '/img'
    file_list = os.listdir(path)

    for file in file_list:
        name = file.split('.')[0]
        file_path = path + '/' + file
        a = cv2.imread(file_path, 0)
        obj[name] = a

    return obj

def GetScreenShot():
    """
    获取屏幕截图
    :return:
    """
    screen = ImageGrab.grab()
    # screen.save('screen.jpg')
    # screen = cv2.imread('screen.jpg')
    screen = cv2.cvtColor(numpy.asarray(screen), cv2.COLOR_RGB2BGR)
    logging.info('截屏成功')
    return screen

class YuHun():
    def __init__(self):
        self._flag = False
        self.NeedCloseGame=False
        self.NeedCloseSystem=False

    def Run(self, LogUI, NeedCloseGame, NeedCloseSystem):
        imgs = loadImgs()
        LogUI.insert(END,
                     time.strftime('%Y-%m-%d %H:%M:%S ',
                                   time.localtime(time.time())) + '开始挑战\n')
        Count = 1
        while self._flag is not True:
            logging.debug('开始挑战')
            screen = GetScreenShot()
            WindowShape = screen.shape
            result = []

            # 为了优化速度，把计算屏幕截图的特征提取出来，避免重复运算
            kp2, des2 = ComputeScreenShot(screen)
            for i in ['tili60', 'tili80', 'auto', 'jieshou2', 'jieshou1', 'end1', 'end2', 'reject', 'queding',
                      'tiaozhan']:
                obj = imgs[i]
                # begin = time.clock()
                pos = GetLocation(obj, kp2, des2)
                # logging.debug('检测结算目标图像')
                # print(time.clock()-begin)
                if pos is not None:
                    if i == 'tili60' or i == 'tili80':
                        print('window.py', NeedCloseSystem)
                        if self.NeedCloseSystem:
                            print('log')
                            os.system('shutdown -s -t 60')
                            return
                        if not self.NeedCloseGame:
                            # 需要手动关闭游戏
                            LogUI.insert(END,
                                         time.strftime('%Y-%m-%d %H:%M:%S ',
                                                       time.localtime(time.time())) + '体力用完，需要手动关闭加成或游戏\n')
                            return
                        # 结束进程
                        hasProcess = True
                        while hasProcess:
                            if 'onmyoji' in os.popen('tasklist /FI "IMAGENAME eq onmyoji.exe"').read():
                                os.system('TASKKILL /F /IM onmyoji.exe')
                                hasProcess = True
                            else:
                                hasProcess = False
                        # 线程结束返回
                        return
                    elif i == 'end1':
                        time.sleep(random.randint(300, 800) / 1000)
                        pos = CheatPos(pos, 50)
                    elif i == 'end2':
                        newPos = (pos[0] + 80, pos[1] + 80)
                        pos = CheatPos(newPos, 5)
                    elif i == 'tiaozhan':
                        LogUI.insert(END,
                                     time.strftime('%Y-%m-%d %H:%M:%S ',
                                                   time.localtime(time.time())) + '第' + str(Count) + '轮开始\n')
                        Count += 1
                    elif i == 'reject':
                        pos = CheatPos(pos, 3)
                    else:
                        pos = CheatPos(pos, 10)
                    result.append(pos)

                    LogUI.see(END)
                else:
                    result.append(None)
            # 开始检查结果
            for i in result:
                if i is not None:
                    print(WindowShape[1] * 0.06)
                    print(WindowShape[0] * 0.96)
                    if i[0] < WindowShape[1] * 0.06 or i[1] > WindowShape[0] * 0.96:
                        continue
                    else:
                        Click(i)
                    if len(LogUI.get('1.0', 'end-1c')) > 6000:
                        LogUI.delete(1.0, END)  # 使用 delete
                        LogUI.insert(END, ' 清空日志\n')
                        LogUI.see(END)

    def Terminate(self):
        self._flag = True

# def YuHunTwoWindow(LogUI, NeedCloseGame, NeedCloseSystem):
#     """
#     自动御魂,双开模式
#     """
#     imgs = loadImgs()
#     LogUI.insert(END,
#                  time.strftime('%Y-%m-%d %H:%M:%S ',
#                                time.localtime(time.time())) + '开始挑战\n')
#     Count = 1
#     while True:
#
#         logging.debug('开始挑战')
#         screen = GetScreenShot()
#         WindowShape = screen.shape
#         result = []
#
#         # 为了优化速度，把计算屏幕截图的特征提取出来，避免重复运算
#         kp2, des2 = ComputeScreenShot(screen)
#         for i in ['tili60', 'tili80', 'auto', 'jieshou2', 'jieshou1', 'end1', 'end2', 'reject', 'queding', 'tiaozhan']:
#             obj = imgs[i]
#             # begin = time.clock()
#             pos = GetLocation(obj, kp2, des2)
#             # logging.debug('检测结算目标图像')
#             # print(time.clock()-begin)
#             if pos is not None:
#                 if i == 'tili60' or i == 'tili80':
#                     print('window.py', NeedCloseSystem)
#                     if NeedCloseSystem:
#                         print('log')
#                         os.system('shutdown -s -t 60')
#                         return
#                     if not NeedCloseGame:
#                         # print('体力用完，需要手动关闭加成或游戏')
#                         LogUI.insert(END,
#                                      time.strftime('%Y-%m-%d %H:%M:%S ',
#                                                    time.localtime(time.time())) + '体力用完，需要手动关闭加成或游戏\n')
#                         return
#                     # 结束进程
#                     hasProcess = True
#                     while hasProcess:
#                         if 'onmyoji' in os.popen('tasklist /FI "IMAGENAME eq onmyoji.exe"').read():
#                             os.system('TASKKILL /F /IM onmyoji.exe')
#                             hasProcess = True
#                         else:
#                             hasProcess = False
#                     # 线程结束返回
#                     return
#                 elif i == 'end1':
#                     time.sleep(random.randint(300, 800) / 1000)
#                     pos = CheatPos(pos, 50)
#                 elif i == 'end2':
#                     newPos = (pos[0] + 80, pos[1] + 80)
#                     pos = CheatPos(newPos, 5)
#                 elif i == 'tiaozhan':
#                     LogUI.insert(END,
#                                  time.strftime('%Y-%m-%d %H:%M:%S ',
#                                                time.localtime(time.time())) + '第' + str(Count) + '轮开始\n')
#                     Count += 1
#                 elif i == 'reject':
#                     pos = CheatPos(pos, 3)
#                 else:
#                     pos = CheatPos(pos, 10)
#                 result.append(pos)
#
#                 LogUI.see(END)
#             else:
#                 result.append(None)
#         # 开始检查结果
#         for i in result:
#             if i is not None:
#                 print(WindowShape[1] * 0.06)
#                 print(WindowShape[0] * 0.96)
#                 if i[0] < WindowShape[1] * 0.06 or i[1] > WindowShape[0] * 0.96:
#                     continue
#                 else:
#                     Click(i)
#                 if len(LogUI.get('1.0', 'end-1c')) > 6000:
#                     LogUI.delete(1.0, END)  # 使用 delete
#                     LogUI.insert(END, ' 清空日志\n')
#                     LogUI.see(END)

if __name__ == '__main__':
    pass


# -*- coding: utf-8 -*-
# Copyright 2018 the HERA Project
# Licensed under the MIT License

import numpy as np
import os
from pyuvdata import UVCal, UVData

class HERACal(UVCal):
    '''
       Class that loads in hera omnical data into a pyuvdata calfits object.
       This can then be saved to a file, plotted, etc.
    '''

    def __init__(self, meta, gains, flags=None, DELAY=False, ex_ants=[], appendhist='', **optional):
        '''Initialize a UVCal object.
            Args:
                meta: meta information dictionary. As returned by from_fits or from_npz.
                gains: dictionary of complex gain solutions or delays.
                flags (optional): Optional input flags for gains.
                DELAY (optional): toggle if calibration solutions in gains are delays.
                ex_ants (optional): antennas that are excluded from gains.
                appendhist (optional): string to append to history
                optional (optional): dictionary of optional parameters to be passed to UVCal object.

            Notes:
                The UVCal parameter cal_style is by default set to 'redundant'. It can be changed
                by feeding cal_style as a keyword argument to HERACal. If cal_style is set to 'sky',
                then there are additional required parameters. See pyuvdata.UVCal doc for details.
                If cal_style is neither redundant or sky, it is set to unknown.
        '''

        super(HERACal, self).__init__()

        # helpful dictionaries for antenna polarization of gains
        str2pol = {'x': -5, 'y': -6}
        pol2str = {-5: 'x', -6: 'y'}

        chisqdict = {}
        datadict = {}
        flagdict = {}

        # drop antennas that are not solved for. Since we are feeding in omnical/firstcal solutions into this,
        # if we provided an ex_ants those antennas will not have a key in gains. Need to provide ex_ants list
        # to HERACal object.
        # create set to get unique antennas from both pol
        ants = np.array(list(set([ant for pol in gains.keys() for ant in gains[pol].keys()])))
        allants = np.sort(np.concatenate([ants, np.array(ex_ants)])).astype(np.int)  # total number of antennas
        ants = np.sort(ants)
        # antenna names for all antennas
        antnames = np.array(['ant' + str(int(ant)) for ant in allants])
        time = meta['times']
        freq = meta['freqs']  # this is in Hz (should be anyways)
        pols = [str2pol[p] for p in gains.keys()]  # all of the polarizations

        # get sizes of things
        nspw = 1  # This is by default 1. No support for > 1 in pyuvdata.
        npol = len(pols)
        ntimes = time.shape[0]
        nfreqs = freq.shape[0]

        datarray = np.array([[gains[pol2str[pol]][ant] for ant in ants]
                             for pol in pols]).swapaxes(0, 3).swapaxes(0, 1)
        if flags:
            flgarray = np.array([[flags[pol2str[pol]][ant] for ant in ants]
                                 for pol in pols]).swapaxes(0, 3).swapaxes(0, 1)
        else:
            if DELAY:
                flgarray = np.zeros((len(ants), nfreqs, ntimes, npol), dtype=bool)
            else:
                # dont need to swap since datarray alread same shape
                flgarray = np.zeros(datarray.shape, dtype=bool)
        # do the same for the chisquare, which is the same shape as the data
        try:
            chisqarray = np.array([[meta['chisq' + str(ant) + pol2str[pol]]
                                    for ant in ants] for pol in pols]).swapaxes(0, 3).swapaxes(0, 1)
        except BaseException:
            # XXX EXCEPT WHAT?
            chisqarray = np.ones(datarray.shape, dtype=float)
        # get the array-wide chisq, which does not have separate axes for
        # antennas or polarization
        try:
            totchisqarray = np.array(meta['chisq']).swapaxes(0, 1)
            # add a polarization axis until this is fixed properly
            totchisqarray = totchisqarray[:, :, np.newaxis]
            # repeat polarization axis npol times for proper shape
            totchisqarray = np.repeat(totchisqarray, npol, axis=-1)
        except BaseException:
            # XXX EXCEPT WHAT?
            # leave it empty
            totchisqarray = None

        pols = np.array(pols)
        freq = np.array(freq)
        antarray = np.array(list(map(int, ants)))
        numarray = np.array(list(map(int, allants)))

        # set UVCal attributes
        self.telescope_name = 'HERA'
        self.Nfreqs = nfreqs
        self.Njones = len(pols)
        self.Ntimes = ntimes
        try:
            self.history = meta['history'].replace('\n', ' ') + appendhist
        except KeyError:
            self.history = appendhist
        self.Nants_data = len(ants)  # only ants with data
        self.Nants_telescope = len(allants)  # all antennas in telescope
        self.antenna_names = antnames
        self.antenna_numbers = numarray
        self.ant_array = np.array(antarray[:self.Nants_data])

        self.Nspws = nspw  # XXX: needs to change when we support more than 1 spw!
        self.spw_array = np.array([0])

        self.freq_array = freq.reshape(self.Nspws, -1)
        self.channel_width = np.diff(self.freq_array)[0][0]
        self.jones_array = pols
        self.time_array = time
        self.integration_time = meta['inttime']
        self.gain_convention = 'divide'
        self.set_redundant()
        self.x_orientation = 'east'
        self.time_range = [self.time_array[0], self.time_array[-1]]
        self.freq_range = [self.freq_array[0][0], self.freq_array[0][-1]]

        # set the optional attributes to UVCal class, which may overwrite previous
        # parameters set by default
        for key in optional:
            setattr(self, key, optional[key])

        # look for cal_style in optional
        if 'cal_style' in optional:
            if optional['cal_style'] == 'redundant':
                pass
            elif optional['cal_style'] == 'sky':
                # set cal style to sky, and set required parameters
                self.set_sky()
                reqs = ['ref_antenna_name', 'sky_catalog', 'sky_field']
                for r in reqs:
                    if r not in optional:
                        raise AttributeError("if cal_style=='sky', then {} must be fed".format(r))
                    setattr(self, r, optional[r])
            else:
                self.set_unknown_cal_type()

        # adding new axis for the spectral window axis. This is default to 1.
        # This needs to change when support for Nspws>1 in pyuvdata.
        self.quality_array = chisqarray[:, np.newaxis, :, :, :]
        self.flag_array = flgarray.astype(np.bool)[:, np.newaxis, :, :, :]
        if DELAY:
            self.set_delay()
            self.delay_array = datarray[:, np.newaxis, :, :, :]  # units of seconds
        else:
            self.set_gain()
            self.gain_array = datarray[:, np.newaxis, :, :, :]
        if totchisqarray is not None:
            self.total_quality_array = totchisqarray[np.newaxis, :, :, :]

        # run a check
        self.check()


import tensorflow as tf
import numpy as np
# import gym
# from gym import wrappers
import tflearn
import argparse
import pprint as pp
import sys

from replay_buffer import ReplayBuffer

# ===========================
#   Actor 和 Critic 网络
# ===========================

class ActorNetwork(object):
    """
    将状态输入到Actor网络中，输出是在确定策略（deterministic policy）下的动作。
    输出层激活函数是tanh，使得输出的动作在[-action_bound,action_bound]范围内。
    """

    def __init__(self, sess, state_dim, action_dim, action_bound, learning_rate, tau, batch_size, noise_option):
        self.sess = sess
        self.s_dim = state_dim
        self.a_dim = action_dim
        self.action_bound = action_bound
        self.learning_rate = learning_rate
        self.tau = tau
        self.batch_size = batch_size

        self.noise_option = noise_option
        self.act_noise = tf.placeholder(tf.bool)

        # Actor Network
        self.inputs, self.out = self.create_actor_network()

        self.network_params = tf.trainable_variables()

        # Target Network
        self.target_inputs, self.target_out = self.create_actor_network()

        self.target_network_params = tf.trainable_variables()[
                                     len(self.network_params):]

        # 固定周期来使用在线网络（online network）来更新目标网络（target network）的权重
        self.update_target_network_params = \
            [self.target_network_params[i].assign(tf.multiply(self.network_params[i], self.tau) +
                                                  tf.multiply(self.target_network_params[i], 1. - self.tau))
             for i in range(len(self.target_network_params))]

        # 由Critic Network计算得到梯度
        self.action_gradient = tf.placeholder(tf.float32, [None, self.a_dim])

        # 结合梯度
        self.unnormalized_actor_gradients = tf.gradients(
            self.out, self.network_params, -self.action_gradient)
        self.actor_gradients = list(map(lambda x: tf.div(x, self.batch_size), self.unnormalized_actor_gradients))

        # 优化过程
        self.optimize = tf.train.AdamOptimizer(self.learning_rate). \
            apply_gradients(zip(self.actor_gradients, self.network_params))

        self.num_trainable_vars = len(
            self.network_params) + len(self.target_network_params)

    # 创建Actor Network网络结构
    def create_actor_network(self):
        inputs = tflearn.input_data(shape=[None, self.s_dim])
        # 400*400 全连接层
        net = tflearn.fully_connected(inputs, 400, weights_init='xavier', bias_init='zeros')
        # Batch Normalization
        net = tflearn.layers.normalization.batch_normalization(net)
        # ReLU激活层
        net = tflearn.activations.relu(net)
        # 300*300 全连接层
        net = tflearn.fully_connected(net, 300, weights_init='xavier', bias_init='zeros')
        # Batch Normalization
        net = tflearn.layers.normalization.batch_normalization(net)
        # ReLU激活层
        net = tflearn.activations.relu(net)

        # 添加高斯噪声来实现探索（Exploration）
        # 也可以尝试使用其他噪声

        tf_cond = tf.cond(self.act_noise, true_fn=lambda: tf.constant(1.0), false_fn=lambda: tf.constant(0.0))

        w_init = tflearn.initializations.uniform(minval=-0.003, maxval=0.003)
        steering = tflearn.fully_connected(net, 1, activation='tanh', weights_init=w_init)
        acceleration = tflearn.fully_connected(net, 1, activation='sigmoid', weights_init=w_init)
        brake = tflearn.fully_connected(net, 1, activation='sigmoid', weights_init=w_init)

        out = tflearn.layers.merge_ops.merge([steering, acceleration, brake], mode='concat', axis=1)

        return inputs, out

    def train(self, inputs, a_gradient):
        self.sess.run(self.optimize, feed_dict={
            self.inputs: inputs,
            self.action_gradient: a_gradient,
            self.act_noise: False
        })

    def predict(self, inputs, noise=False):
        return self.sess.run(self.out, feed_dict={
            self.inputs: inputs, self.act_noise: noise
        })

    def predict_target(self, inputs):
        return self.sess.run(self.target_out, feed_dict={
            self.target_inputs: inputs, self.act_noise: False
        })

    def update_target_network(self):
        self.sess.run(self.update_target_network_params)

    def get_num_trainable_vars(self):
        return self.num_trainable_vars

class CriticNetwork(object):

    def __init__(self, sess, state_dim, action_dim, learning_rate, tau, gamma, num_actor_vars):
        self.sess = sess
        self.s_dim = state_dim
        self.a_dim = action_dim
        self.learning_rate = learning_rate
        self.tau = tau
        self.gamma = gamma

        self.inputs, self.action, self.out = self.create_critic_network()

        self.network_params = tf.trainable_variables()[num_actor_vars:]

        self.target_inputs, self.target_action, self.target_out = self.create_critic_network()

        self.target_network_params = tf.trainable_variables()[(len(self.network_params) + num_actor_vars):]

        self.update_target_network_params = \
            [self.target_network_params[i].assign(tf.multiply(self.network_params[i], self.tau) \
                                                  + tf.multiply(self.target_network_params[i], 1. - self.tau))
             for i in range(len(self.target_network_params))]

        self.predicted_q_value = tf.placeholder(tf.float32, [None, 1])

        self.loss = tflearn.mean_square(self.predicted_q_value, self.out)
        self.optimize = tf.train.AdamOptimizer(
            self.learning_rate).minimize(self.loss)

        self.action_grads = tf.gradients(self.out, self.action)

    def create_critic_network(self):
        inputs = tflearn.input_data(shape=[None, self.s_dim])
        action = tflearn.input_data(shape=[None, self.a_dim])
        net = tflearn.fully_connected(inputs, 400, weights_init='xavier', bias_init='zeros')
        tflearn.add_weights_regularizer(net, 'L2', weight_decay=0.001)
        net = tflearn.layers.normalization.batch_normalization(net)
        net = tflearn.activations.relu(net)

        t1 = tflearn.fully_connected(net, 300, weights_init='xavier', bias_init='zeros')
        t2 = tflearn.fully_connected(action, 300, weights_init='xavier', bias_init='zeros')

        net = tflearn.activation(tf.matmul(net, t1.W) + tf.matmul(action, t2.W) + t2.b, activation='relu')

        w_init = tflearn.initializations.uniform(minval=-0.003, maxval=0.003)

        out = tflearn.fully_connected(net, 1, weights_init=w_init)
        return inputs, action, out

    def train(self, inputs, action, predicted_q_value):
        return self.sess.run([self.out, self.optimize], feed_dict={
            self.inputs: inputs,
            self.action: action,
            self.predicted_q_value: predicted_q_value
        })

    def predict(self, inputs, action):
        return self.sess.run(self.out, feed_dict={
            self.inputs: inputs,
            self.action: action
        })

    def predict_target(self, inputs, action):
        return self.sess.run(self.target_out, feed_dict={
            self.target_inputs: inputs,
            self.target_action: action
        })

    def action_gradients(self, inputs, actions):
        return self.sess.run(self.action_grads, feed_dict={
            self.inputs: inputs,
            self.action: actions
        })

    def update_target_network(self):
        self.sess.run(self.update_target_network_params)


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

# Synthesize waveform using converted features.
# By Wen-Chin Huang 2019.06

import json
import os

import tensorflow as tf
import numpy as np

from datetime import datetime
from importlib import import_module

import pysptk
import pyworld as pw
from scipy.io import loadmat, savemat
from scipy.io import wavfile

import argparse
import logging
import multiprocessing as mp
import glob

import sys
from preprocessing.vcc2018.feature_reader import Whole_feature_reader
from util.synthesizer import world_synthesis
from util.misc import read_hdf5
from util.postfilter import fast_MLGV
from util.f0transformation import log_linear_transformation

def read_and_synthesize(file_list, arch, stats, input_feat, output_feat):

    for i, (bin_path, feat_path) in enumerate(file_list):
        input_feat_dim = arch['feat_param']['dim'][input_feat]

        # define paths
        output_dir = os.path.dirname(bin_path).replace('converted-' + output_feat, 'converted-wav')
        basename = os.path.splitext(os.path.split(bin_path)[-1])[0]
        wav_name = os.path.join(output_dir, basename + '.wav')
        gv_wav_name = os.path.join(output_dir, basename + '-gv.wav')

        # read source features and converted spectral features
        src_data = Whole_feature_reader(feat_path, arch['feat_param'])
        cvt = np.fromfile(bin_path, dtype = np.float32).reshape([-1, input_feat_dim])

        # f0 conversion
        lf0 = log_linear_transformation(src_data['f0'], stats)

        # apply gv post filtering to converted
        cvt_gv = fast_MLGV(cvt, stats['gv_t'])

        # energy compensation
        if output_feat == 'mcc':
            en_cvt = np.c_[src_data['en_mcc'], cvt]
            en_cvt_gv = np.c_[src_data['en_mcc'], cvt_gv]
        elif output_feat == 'sp':
            cvt = np.power(10., cvt)
            en_cvt = np.expand_dims(src_data['en_sp'], 1) * cvt
            cvt_gv = np.power(10., cvt_gv)
            en_cvt_gv = np.expand_dims(src_data['en_sp'], 1) * cvt_gv

        # synthesis
        world_synthesis(wav_name, arch['feat_param'],
                        lf0.astype(np.float64).copy(order='C'),
                        src_data['ap'].astype(np.float64).copy(order='C'),
                        en_cvt.astype(np.float64).copy(order='C'),
                        output_feat)
        world_synthesis(gv_wav_name, arch['feat_param'],
                        lf0.astype(np.float64).copy(order='C'),
                        src_data['ap'].astype(np.float64).copy(order='C'),
                        en_cvt_gv.astype(np.float64).copy(order='C'),
                        output_feat)

def main():

    parser = argparse.ArgumentParser(
        description="synthesize waveforms using converted files.")
    parser.add_argument(
        "--logdir", required=True, type=str,
        help="path of log directory")
    parser.add_argument(
        "--type", default='test', type=str,
        help="test or valid (default is test)")
    parser.add_argument(
        "--input_feat", required=True,
        type=str, help="input feature type")
    parser.add_argument(
        "--output_feat", required=True,
        type=str, help="output feature type")
    parser.add_argument(
        "--n_jobs", default=12,
        type=int, help="number of parallel jobs")
    args = parser.parse_args()

    # set log level
    fmt = '%(asctime)s %(message)s'
    datefmt = '%m/%d/%Y %I:%M:%S'
    logFormatter = logging.Formatter(fmt, datefmt=datefmt)
    logging.basicConfig(
        level=logging.INFO,
        filename=os.path.join(args.logdir, 'exp.log'),
        format=fmt,
        datefmt=datefmt,
        )
    consoleHandler = logging.StreamHandler()
    consoleHandler.setFormatter(logFormatter)
    logging.getLogger().addHandler(consoleHandler)
    logging.info('====================')
    logging.info('Synthesize start')
    logging.info(args)

    train_dir = os.sep.join(os.path.normpath(args.logdir).split(os.sep)[:-1])
    output_dir = os.path.basename(os.path.normpath(args.logdir))
    src, trg = output_dir.split('-')[-2:]

    # Load architecture
    # arch = tf.gfile.Glob(os.path.join(train_dir, 'architecture*.json'))[0]  # should only be 1 file
    arch = glob.glob(os.path.join(train_dir, 'architecture*.json'))[0]  # should only be 1 file
    with open(arch) as fp:
        arch = json.load(fp)

    input_feat = args.input_feat
    output_feat = args.output_feat

    # Load statistics
    stats = {
        'mu_s' : read_hdf5(arch['stats'], '/f0/' + src + '/mean'),
        'std_s' : read_hdf5(arch['stats'], '/f0/' + src + '/std'),
        'mu_t' : read_hdf5(arch['stats'], '/f0/' + trg + '/mean'),
        'std_t' : read_hdf5(arch['stats'], '/f0/' + trg + '/std'),
        'gv_t' : read_hdf5(arch['stats'], '/gv_{}/'.format(output_feat) + trg),
    }

    # Make directory
    tf.gfile.MakeDirs(os.path.join(args.logdir, 'converted-wav'))

    # Get and divide list
    # bin_list = sorted(tf.gfile.Glob(os.path.join(args.logdir, 'converted-{}'.format(output_feat), '*.bin')))
    bin_list = sorted(glob.glob(os.path.join(args.logdir, 'converted-{}'.format(output_feat), '*.bin')))
    if args.type == 'test':
        # feat_list = sorted(tf.gfile.Glob(arch['conversion']['test_file_pattern'].format(src)))
        feat_list = sorted(glob.glob(arch['conversion']['test_file_pattern'].format(src)))
    elif args.type == 'valid':
        feat_list = []
        for p in arch['training']['valid_file_pattern']:
            # feat_list.extend(tf.gfile.Glob(p.replace('*', src)))
            feat_list.extend(glob.glob(p.replace('*', src)))
        feat_list = sorted(feat_list)

    assert(len(bin_list) == len(feat_list))
    file_list = list(zip(bin_list, feat_list))
    logging.info("number of utterances = %d" % len(file_list))
    file_lists = np.array_split(file_list, args.n_jobs)
    file_lists = [f_list.tolist() for f_list in file_lists]

    # multi processing
    processes = []
    for f in file_lists:
        p = mp.Process(target=read_and_synthesize, args=(f, arch, stats, input_feat, output_feat))
        p.start()
        processes.append(p)

    # wait for all process
    for p in processes:
        p.join()

if __name__ == '__main__':
    main()


#! /usr/bin/env python
import argparse, sys, os, errno
import logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)s [%(levelname)s] : %(message)s')
logger = logging.getLogger('train')

def segment(args):
    from models import custom_objects
    from keras.models import load_model
    import numpy as np
    import h5py
    globals().update(locals())

    logger.info('load model from file: ' + args.model_file)
    model = load_model(args.model_file, custom_objects=custom_objects)

    logger.info('read image ids from file: ' + args.image_id_file)
    image_id = read_hdf5(args.image_id_file, args.image_id_dataset)

    logger.info('read image data from file: ' + args.input_file)
    fin = h5py.File(args.input_file, 'r')
    X = fin['X'][:]
    image_id_X = fin['image_id'][:]
    fin.close()
    if model.input.shape[3] > 1:
        logger.info('convert gray-scale images to 3-channel images')
        X = np.repeat(X[array_lookup(image_id_X, image_id)], 3, axis=3)
    else:
        X = np.take(X, array_lookup(image_id_X, image_id), axis=0)

    logger.info('predict')
    y = model.predict(X, batch_size=args.batch_size)
    if args.mix:
        logger.info('mix output masks with input images')
        y = np.clip(y + X, 0, 1)

    logger.info('create output file: ' + args.output_file)
    prepare_output_file(args.output_file)
    fout = h5py.File(args.output_file, 'w')
    fout.create_dataset('X', data=y)
    fout.create_dataset('image_id', data=image_id)
    fout.close()

def classify_diseases(args):
    from keras.models import load_model
    from models import custom_objects
    import numpy as np
    import h5py
    globals().update(locals())

    logger.info('load model from file: ' + args.model_file)
    model = load_model(args.model_file, custom_objects=custom_objects)

    logger.info('read image ids from file: ' + args.image_id_file)
    image_id = read_hdf5(args.image_id_file, args.image_id_dataset)

    logger.info('read image data from file: ' + args.input_file)
    fin = h5py.File(args.input_file, 'r')
    X = fin['X'][:]
    image_id_X = fin['image_id'][:]
    fin.close()
    if model.input.shape[3] > 1:
        logger.info('convert gray-scale images to 3-channel images')
        X = np.repeat(X[array_lookup(image_id_X, image_id)], 3, axis=3)
    else:
        X = np.take(X, array_lookup(image_id_X, image_id), axis=0)

    logger.info('predict')
    y = model.predict(X, batch_size=args.batch_size)

    logger.info('create output file: ' + args.output_file)
    prepare_output_file(args.output_file)
    fout = h5py.File(args.output_file, 'w')
    fout.create_dataset('y', data=y)
    fout.create_dataset('image_id', data=image_id)
    fout.close()

if __name__ == '__main__':
    main_parser = argparse.ArgumentParser(description='Predict using saved models')
    subparsers = main_parser.add_subparsers(dest='command')

    # command: classify_diseases
    parser = subparsers.add_parser('classify_diseases')

    # command: segment
    parser = subparsers.add_parser('segment')
    parser.add_argument('-i', '--input-file', type=str, required=True,
                        help='input HDF5 file containing images to segment')
    parser.add_argument('-m', '--model-file', type=str, required=True,
                        help='a saved keras model file')
    parser.add_argument('-y', '--mask-file', type=str,
                        help='an HDF5 file containing mask images of the same shape of input images')
    parser.add_argument('--image-id-file', type=str,
                        help='HDF5 file containing image ids')
    parser.add_argument('--image-id-dataset', type=str,
                        help='HDF5 dataset name in the file specified by --image-id-file')
    parser.add_argument('--mix', action='store_true',
                        help='mix the segmentation result with input')
    parser.add_argument('--batch-size', type=int, default=32)
    parser.add_argument('-o', '--output-file', type=str, required=True,
                        help='output directory for saving models and predictions')

    args = main_parser.parse_args()
    command_handlers = {
        'segment': segment
    }
    logger = logging.getLogger('predict.' + args.command)
    import numpy as np
    from utils import read_hdf5, prepare_output_file, array_lookup
    command_handlers[args.command](args)


import logging
import random

import numpy as np
import torch
from fastprogress.fastprogress import progress_bar
from torch.utils.data import DataLoader, SequentialSampler
from transformers import ElectraForSequenceClassification, ElectraTokenizer

logger = logging.getLogger(__name__)

class GrandChallengeTextClassifier:
    def __init__(self, args):
        self.args = args

        self.tokenizer = ElectraTokenizer.from_pretrained(
            self.args.tokenizer_dir, do_lower_case=self.args.do_lower_case
        )

        logger.info("Predict text class the following checkpoints: %s", self.args.ckpt_dir)

        self.args.device = "cuda" if torch.cuda.is_available() and not self.args.no_cuda else "cpu"

        self.model = ElectraForSequenceClassification.from_pretrained(self.args.ckpt_dir)
        self.model.to(self.args.device)
        self.model.eval()

    def predict(self, test_dataset):
        sampler = SequentialSampler(test_dataset)
        dataloader = DataLoader(test_dataset, sampler=sampler, batch_size=self.args.eval_batch_size)

        logger.info("Num examples = {}".format(len(test_dataset)))

        preds = None

        for batch in progress_bar(dataloader):
            batch = tuple(t.to(self.args.device) for t in batch)

            with torch.no_grad():
                inputs = {
                    "input_ids": batch[0],
                    "attention_mask": batch[1],
                    "labels": batch[3],
                }

                outputs = self.model(**inputs)
                logits = outputs[1]

            preds = logits.detach().cpu().numpy()

        preds = np.argmax(preds, axis=1)

        return preds


#! /usr/bin/env python

import cv2
import numpy as np
import onnxruntime as rt

camera = cv2.VideoCapture(0)

session = rt.InferenceSession("yolov3.onnx")

def get_polyp(frame):
    image_size = np.array([256, 256], dtype=np.float32).reshape(1, 2)
    image_data = np.array(frame, dtype='float32')
    image_data /= 255.0
    image_data = np.transpose(image_data, [2, 0, 1])
    image_data = np.expand_dims(image_data, 0)

    outputs_index = session.run(None, {'input_1': image_data, 'image_shape': image_size})
    output_boxes = outputs_index[0]
    output_indices = outputs_index[2]

    return [output_boxes[0][_idx[2]] for _idx in output_indices[0]]

while True:
    ret, frame = camera.read()
    width = 256
    height = 256
    frame = cv2.resize(frame, (width, height))
    boxes = get_polyp(frame)
    for box in boxes:
        frame = cv2.rectangle(frame, (box[1], box[0]), (box[3], box[2]), (0, 255, 0), 3)
    frame = cv2.resize(frame, (700, 700))
    cv2.imshow('camera', frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

camera.release()
cv2.destroyAllWindows()


# -*- coding:utf-8 -*-

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

import os, sys
import tensorflow as tf
import time
import cv2
import pickle
import numpy as np
import json
sys.path.append("../")

from data.io.image_preprocess import short_side_resize_for_inference_data
from libs.configs import cfgs
from libs.networks import build_whole_network
from libs.val_libs import voc_eval
from libs.box_utils import draw_box_in_img
import argparse
from help_utils import tools
from libs.label_name_dict.label_dict import *

from data.lib_coco.PythonAPI.pycocotools.coco import COCO
from data.lib_coco.PythonAPI.pycocotools.cocoeval import COCOeval

def test_coco(det_net, real_test_img_list, eval_data, draw_imgs=False):

    # 1. preprocess img
    img_plac = tf.placeholder(dtype=tf.uint8, shape=[None, None, 3])  # is RGB. not BGR
    img_batch = tf.cast(img_plac, tf.float32)

    img_batch = short_side_resize_for_inference_data(img_tensor=img_batch,
                                                     target_shortside_len=cfgs.IMG_SHORT_SIDE_LEN,
                                                     length_limitation=cfgs.IMG_MAX_LENGTH)
    if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']:
        img_batch = (img_batch / 255 - tf.constant(cfgs.PIXEL_MEAN_)) / tf.constant(cfgs.PIXEL_STD)
    else:
        img_batch = img_batch - tf.constant(cfgs.PIXEL_MEAN)

    # img_batch = (img_batch - tf.constant(cfgs.PIXEL_MEAN)) / (tf.constant(cfgs.PIXEL_STD)*255)
    img_batch = tf.expand_dims(img_batch, axis=0)

    detection_boxes, detection_scores, detection_category = det_net.build_whole_detection_network(
        input_img_batch=img_batch,
        gtboxes_batch=None)

    init_op = tf.group(
        tf.global_variables_initializer(),
        tf.local_variables_initializer()
    )

    restorer, restore_ckpt = det_net.get_restorer()

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    with tf.Session(config=config) as sess:
        sess.run(init_op)
        if not restorer is None:
            restorer.restore(sess, restore_ckpt)
            print('restore model')

        save_path = os.path.join('./eval_coco', cfgs.VERSION)
        tools.mkdir(save_path)
        fw_json_dt = open(os.path.join(save_path, 'coco_test-dev.json'), 'w')
        coco_det = []
        for i, a_img in enumerate(real_test_img_list):

            raw_img = cv2.imread(os.path.join(eval_data, a_img['file_name']))
            raw_h, raw_w = raw_img.shape[0], raw_img.shape[1]

            start = time.time()
            resized_img, detected_boxes, detected_scores, detected_categories = \
                sess.run(
                    [img_batch, detection_boxes, detection_scores, detection_category],
                    feed_dict={img_plac: raw_img[:, :, ::-1]}  # cv is BGR. But need RGB
                )
            end = time.time()

            eval_indices = detected_scores >= 0.01
            detected_scores = detected_scores[eval_indices]
            detected_boxes = detected_boxes[eval_indices]
            detected_categories = detected_categories[eval_indices]

            # print("{} cost time : {} ".format(img_name, (end - start)))
            if draw_imgs:
                show_indices = detected_scores >= cfgs.SHOW_SCORE_THRSHOLD
                show_scores = detected_scores[show_indices]
                show_boxes = detected_boxes[show_indices]
                show_categories = detected_categories[show_indices]

                draw_img = np.squeeze(resized_img, 0)
                if cfgs.NET_NAME in ['resnet101_v1d', 'resnet50_v1d']:
                    draw_img = (draw_img * np.array(cfgs.PIXEL_STD) + np.array(cfgs.PIXEL_MEAN_)) * 255
                else:
                    draw_img = draw_img + np.array(cfgs.PIXEL_MEAN)

                # draw_img = draw_img * (np.array(cfgs.PIXEL_STD)*255) + np.array(cfgs.PIXEL_MEAN)

                final_detections = draw_box_in_img.draw_boxes_with_label_and_scores(draw_img,
                                                                                boxes=show_boxes,
                                                                                labels=show_categories,
                                                                                scores=show_scores,
                                                                                in_graph=False)
                if not os.path.exists(cfgs.TEST_SAVE_PATH):
                    os.makedirs(cfgs.TEST_SAVE_PATH)

                cv2.imwrite(cfgs.TEST_SAVE_PATH + '/' + '{}.jpg'.format(a_img['id']),
                            final_detections[:, :, ::-1])

            xmin, ymin, xmax, ymax = detected_boxes[:, 0], detected_boxes[:, 1], \
                                     detected_boxes[:, 2], detected_boxes[:, 3]

            resized_h, resized_w = resized_img.shape[1], resized_img.shape[2]

            xmin = xmin * raw_w / resized_w
            xmax = xmax * raw_w / resized_w

            ymin = ymin * raw_h / resized_h
            ymax = ymax * raw_h / resized_h

            boxes = np.transpose(np.stack([xmin, ymin, xmax-xmin, ymax-ymin]))

            # cost much time
            for j, box in enumerate(boxes):
                coco_det.append({'bbox': [float(box[0]), float(box[1]), float(box[2]), float(box[3])],
                                 'score': float(detected_scores[j]), 'image_id': a_img['id'],
                                 'category_id': int(classes_originID[LABEl_NAME_MAP[detected_categories[j]]])})

            tools.view_bar('%s image cost %.3fs' % (a_img['id'], (end - start)), i + 1, len(real_test_img_list))

        json.dump(coco_det, fw_json_dt)
        fw_json_dt.close()

def eval(num_imgs, eval_data, json_file, showbox):

    with open(json_file) as f:
        test_img_list = json.load(f)['images']

    if num_imgs == np.inf:
        real_test_img_list = test_img_list
    else:
        real_test_img_list = test_img_list[: num_imgs]

    faster_rcnn = build_whole_network.DetectionNetwork(base_network_name=cfgs.NET_NAME,
                                                       is_training=False)
    test_coco(det_net=faster_rcnn, real_test_img_list=real_test_img_list, eval_data=eval_data, draw_imgs=showbox)

def parse_args():

    parser = argparse.ArgumentParser('evaluate the result with Pascal2007 stdand')

    parser.add_argument('--eval_data', dest='eval_data',
                        help='evaluate imgs dir ',
                        default='coco_minival2014.odgt', type=str)
    parser.add_argument('--json_file', dest='json_file',
                        help='test-dev json file',
                        default='image_info_test-dev2017.json', type=str)
    parser.add_argument('--showbox', dest='showbox',
                        help='whether show detecion results when evaluation',
                        default=False, type=bool)
    parser.add_argument('--GPU', dest='GPU',
                        help='gpu id',
                        default='0', type=str)
    parser.add_argument('--eval_num', dest='eval_num',
                        help='the num of eval imgs',
                        default=np.inf, type=int)
    args = parser.parse_args()
    return args

if __name__ == '__main__':

    args = parse_args()
    print(20*"--")
    print(args)
    print(20*"--")
    os.environ["CUDA_VISIBLE_DEVICES"] = args.GPU
    eval(np.inf,  # use np.inf to test all the imgs. use 10 to test 10 imgs.
         eval_data=args.eval_data,
         json_file=args.json_file,
         showbox=args.showbox)


from lib.audio_utils import *
from lib.collection_utils import *
from lib.io_utils import *
from lib.math_utils import *
import os
from pydub import AudioSegment
import sys

def getAudioSequenceDuration(sequence):
    if len(sequence) <= 0:
        return 0
    sequence = sorted(sequence, key=lambda s: s["ms"]+s["dur"])
    lastAudioClip = sequence[-1]
    return lastAudioClip["ms"] + lastAudioClip["dur"]

def makeTrack(duration, instructions, segments, sfx=True, sampleWidth=4, sampleRate=48000, channels=2, fxPad=3000):
    # build audio
    baseAudio = AudioSegment.silent(duration=duration, frame_rate=sampleRate)
    baseAudio = baseAudio.set_channels(channels)
    baseAudio = baseAudio.set_sample_width(sampleWidth)
    instructionCount = len(instructions)
    for index, i in enumerate(instructions):
        segment = [s for s in segments if s["id"]==(i["start"], i["dur"])].pop()
        audio = segment["audio"]
        audio = applyAudioProperties(audio, i, sfx, fxPad)
        # convert sample width
        if audio.sample_width != sampleWidth:
            # print("Warning: sample width changed to %s from %s" % (sampleWidth, audio.sample_width))
            audio = audio.set_sample_width(sampleWidth)
        baseAudio = baseAudio.overlay(audio, position=i["ms"])
        sys.stdout.write('\r')
        sys.stdout.write("%s%%" % round(1.0*(index+1)/instructionCount*100,1))
        sys.stdout.flush()
    return baseAudio

def mixAudio(instructions, duration, outfilename, sfx=True, sampleWidth=4, sampleRate=48000, channels=2, fxPad=3000, masterDb=0.0, outputTracks=False, tracksDir="output/tracks/%s.wav"):
    # remove instructions with no volume
    instructions = [i for i in instructions if "volume" not in i or i["volume"] > 0]
    audioFiles = list(set([i["filename"] for i in instructions]))
    audioFiles = [{"filename": f} for f in audioFiles]
    instructionCount = len(instructions)
    trackCount = len(audioFiles)

    # calculate db
    for i, step in enumerate(instructions):
        if "volume" in step:
            instructions[i]["db"] = volumeToDb(step["volume"])

    # create base audio
    baseAudio = AudioSegment.silent(duration=duration, frame_rate=sampleRate)
    baseAudio = baseAudio.set_channels(channels)
    baseAudio = baseAudio.set_sample_width(sampleWidth)

    # Load sounds
    print("Adding tracks...")
    for i, af in enumerate(audioFiles):
        filename = af["filename"]
        audioFiles[i]["index"] = i

        # load audio file
        audio = getAudio(filename, sampleWidth, sampleRate, channels)
        audioDurationMs = len(audio)

        # look through instructions to find unique clips
        clips = [(ii["start"], ii["dur"]) for ii in instructions if ii["filename"]==filename]
        clips = list(set(clips))

        # make segments from clips
        segments = []
        for clipStart, clipDur in clips:
            clip = getAudioClip(audio, clipStart, clipDur, audioDurationMs)
            if clip is None:
                continue
            segments.append({
                "id": (clipStart, clipDur),
                "start": clipStart,
                "dur": clipDur,
                "audio": clip
            })

        # make the track
        trackInstructions = [ii for ii in instructions if ii["filename"]==af["filename"]]
        print("Making track %s of %s with %s segments and %s instructions..." % (i+1, trackCount, len(segments), len(trackInstructions)))
        trackAudio = makeTrack(duration, trackInstructions, segments, sfx=sfx, sampleWidth=sampleWidth, sampleRate=sampleRate, channels=channels, fxPad=fxPad)
        baseAudio = baseAudio.overlay(trackAudio)
        if outputTracks:
            trackfilename = tracksDir % getBasename(filename)
            format = trackfilename.split(".")[-1]
            # adjust master volume
            if masterDb != 0.0:
                trackAudio = trackAudio.apply_gain(masterDb)
            trackAudio.export(trackfilename, format=format)
            print("Wrote to %s" % trackfilename)
        print("Track %s of %s complete." % (i+1, trackCount))

    print("Writing to file...")
    format = outfilename.split(".")[-1]
    # adjust master volume
    if masterDb != 0.0:
        baseAudio = baseAudio.apply_gain(masterDb)
    f = baseAudio.export(outfilename, format=format)
    print("Wrote to %s" % outfilename)

def plotAudioSequence(seq):
    import matplotlib.pyplot as plt
    import numpy as np

    filenames = groupList(seq, "filename")
    labels = [f["filename"] for f in filenames]

    fig, ax = plt.subplots(figsize=(24,12))
    colors = iter(plt.cm.prism(np.linspace(0,1,len(filenames))))
    for i, f in enumerate(filenames):
        data = [(step["ms"]/60000.0, step["dur"]/60000.0) for step in f["items"]]
        color = next(colors)
        h = 0.8
        margin = 0.4
        ax.broken_barh(data, (i-margin,h), color=color)

    ax.set_yticks(range(len(labels)))
    # ax.set_yticklabels(labels)
    ax.set_xlabel("time [minutes]")
    plt.tight_layout()
    plt.show()