import os
import dlib
import cv2
import numpy as np
import time
import shutil

from modules import progress

cur_dir = os.path.dirname(__file__)
predictor_path = os.path.join(cur_dir, '../assets/predictor/face_landmarks.dat')
predictor = dlib.shape_predictor(predictor_path)
detector = dlib.get_frontal_face_detector()

TEMPLATE = np.float32([
    (0.0792396913815, 0.339223741112), (0.0829219487236, 0.456955367943),
    (0.0967927109165, 0.575648016728), (0.122141515615, 0.691921601066),
    (0.168687863544, 0.800341263616), (0.239789390707, 0.895732504778),
    (0.325662452515, 0.977068762493), (0.422318282013, 1.04329000149),
    (0.531777802068, 1.06080371126), (0.641296298053, 1.03981924107),
    (0.738105872266, 0.972268833998), (0.824444363295, 0.889624082279),
    (0.894792677532, 0.792494155836), (0.939395486253, 0.681546643421),
    (0.96111933829, 0.562238253072), (0.970579841181, 0.441758925744),
    (0.971193274221, 0.322118743967), (0.163846223133, 0.249151738053),
    (0.21780354657, 0.204255863861), (0.291299351124, 0.192367318323),
    (0.367460241458, 0.203582210627), (0.4392945113, 0.233135599851),
    (0.586445962425, 0.228141644834), (0.660152671635, 0.195923841854),
    (0.737466449096, 0.182360984545), (0.813236546239, 0.192828009114),
    (0.8707571886, 0.235293377042), (0.51534533827, 0.31863546193),
    (0.516221448289, 0.396200446263), (0.517118861835, 0.473797687758),
    (0.51816430343, 0.553157797772), (0.433701156035, 0.604054457668),
    (0.475501237769, 0.62076344024), (0.520712933176, 0.634268222208),
    (0.565874114041, 0.618796581487), (0.607054002672, 0.60157671656),
    (0.252418718401, 0.331052263829), (0.298663015648, 0.302646354002),
    (0.355749724218, 0.303020650651), (0.403718978315, 0.33867711083),
    (0.352507175597, 0.349987615384), (0.296791759886, 0.350478978225),
    (0.631326076346, 0.334136672344), (0.679073381078, 0.29645404267),
    (0.73597236153, 0.294721285802), (0.782865376271, 0.321305281656),
    (0.740312274764, 0.341849376713), (0.68499850091, 0.343734332172),
    (0.353167761422, 0.746189164237), (0.414587777921, 0.719053835073),
    (0.477677654595, 0.706835892494), (0.522732900812, 0.717092275768),
    (0.569832064287, 0.705414478982), (0.635195811927, 0.71565572516),
    (0.69951672331, 0.739419187253), (0.639447159575, 0.805236879972),
    (0.576410514055, 0.835436670169), (0.525398405766, 0.841706377792),
    (0.47641545769, 0.837505914975), (0.41379548902, 0.810045601727),
    (0.380084785646, 0.749979603086), (0.477955996282, 0.74513234612),
    (0.523389793327, 0.748924302636), (0.571057789237, 0.74332894691),
    (0.672409137852, 0.744177032192), (0.572539621444, 0.776609286626),
    (0.5240106503, 0.783370783245), (0.477561227414, 0.778476346951)])

INV_TEMPLATE = np.float32([
    (-0.04099179660567834, -0.008425234314031194, 2.575498465013183),
    (0.04062510634554352, -0.009678089746831375, -1.2534351452524177),
    (0.0003666902601348179, 0.01810332406086298, -0.32206331976076663)])

TPL_MIN, TPL_MAX = np.min(TEMPLATE, axis=0), np.max(TEMPLATE, axis=0)
MINMAX_TEMPLATE = (TEMPLATE - TPL_MIN) / (TPL_MAX - TPL_MIN)

#: Landmark indices corresponding to the inner eyes and bottom lip.
INNER_EYES_AND_BOTTOM_LIP = [39, 42, 57]

#: Landmark indices corresponding to the outer eyes and nose.
OUTER_EYES_AND_NOSE = [36, 45, 33]

def find_landmarks(rgb_img, bb):
    """
    Find the landmarks of a face.
    :param rgb_img: RGB image to process. Shape: (height, width, 3)
    :type rgb_img: numpy.ndarray
    :param bb: Bounding box around the face to find landmarks for.
    :type bb: dlib.rectangle
    :return: Detected landmark locations.
    :rtype: list of (x,y) tuples
    """
    assert rgb_img is not None
    assert bb is not None

    points = predictor(rgb_img, bb)
    # return list(map(lambda p: (p.x, p.y), points.parts()))
    return [(p.x, p.y) for p in points.parts()]

# pylint: disable=dangerous-default-value
def align(img_dim, rgb_img, bb=None,
          landmarks=None, landmark_indices=INNER_EYES_AND_BOTTOM_LIP,
          skip_multi=False, scale=1.0):
    r"""align(imgDim, rgbImg, bb=None, landmarks=None, landmarkIndices=INNER_EYES_AND_BOTTOM_LIP)
    Transform and align a face in an image.
    :param img_dim: The edge length in pixels of the square the image is resized to.
    :type img_dim: int
    :param rgb_img: RGB image to process. Shape: (height, width, 3)
    :type rgb_img: numpy.ndarray
    :param bb: Bounding box around the face to align. \
               Defaults to the largest face.
    :type bb: dlib.rectangle
    :param landmarks: Detected landmark locations. \
                      Landmarks found on `bb` if not provided.
    :type landmarks: list of (x,y) tuples
    :param landmark_indices: The indices to transform to.
    :type landmark_indices: list of ints
    :param skip_multi: Skip image if more than one face detected.
    :type skip_multi: bool
    :param scale: Scale image before cropping to the size given by imgDim.
    :type scale: float
    :return: The aligned RGB image. Shape: (imgDim, imgDim, 3)
    :rtype: numpy.ndarray
    """
    assert img_dim is not None
    assert rgb_img is not None
    assert landmark_indices is not None

    if bb is None:
        bb = get_largest_face_bounding_box(rgb_img, skip_multi)
        if bb is None:
            return None, None

    if landmarks is None:
        landmarks = find_landmarks(rgb_img, bb)

    np_landmarks = np.float32(landmarks)
    np_landmark_indices = np.array(landmark_indices)

    # pylint: disable=maybe-no-member
    H = cv2.getAffineTransform(np_landmarks[np_landmark_indices],
                               img_dim * MINMAX_TEMPLATE[np_landmark_indices] * scale + img_dim * (1 - scale) / 2)
    thumbnail = cv2.warpAffine(rgb_img, H, (img_dim, img_dim))

    return thumbnail, landmarks

def get_all_face_bounding_boxes(rgb_img):
    """
    Find all face bounding boxes in an image.
    :param rgb_img: RGB image to process. Shape: (height, width, 3)
    :type rgb_img: numpy.ndarray
    :return: All face bounding boxes in an image.
    :rtype: dlib.rectangles
    """
    assert rgb_img is not None

    try:
        return detector(rgb_img, 1)
    except Exception as e:  # pylint: disable=broad-except
        print("Warning: {}".format(e))
        # In rare cases, exceptions are thrown.
        return []

def get_largest_face_bounding_box(rgb_img, skip_multi=False):
    """
    Find the largest face bounding box in an image.
    :param rgb_img: RGB image to process. Shape: (height, width, 3)
    :type rgb_img: numpy.ndarray
    :param skip_multi: Skip image if more than one face detected.
    :type skip_multi: bool
    :return: The largest face bounding box in an image, or None.
    :rtype: dlib.rectangle
    """
    assert rgb_img is not None

    faces = get_all_face_bounding_boxes(rgb_img)
    if (not skip_multi and len(faces) > 0) or len(faces) == 1:
        return max(faces, key=lambda rect: rect.width() * rect.height())
    else:
        return None

def normalize_face(face_img, size, should_show, skip_multi=False):
    """
    normalize a given image
    :param face_img: RGB image to process. Shape: (height, width, 3)
    :type face_img: numpy.ndarray
    :param size: size for the output image in height & width px
    :type size: number
    :param should_show: if the normalized face should be shown in a frame
    :type should_show: boolean
    :param skip_multi: Skip image if more than one face detected.
    :type skip_multi: bool
    :return: landmarks and aligned face
    """
    bb = get_largest_face_bounding_box(face_img, skip_multi=skip_multi)

    aligned, landmarks = align(size, face_img, bb, landmark_indices=INNER_EYES_AND_BOTTOM_LIP, skip_multi=skip_multi)
    if aligned is not None:
        if should_show:
            cv2.imshow("aligned", aligned)
        return landmarks, aligned
    return None, None

def norm_image(image_root, norm_root, image_dir, image, undetected_dir):
    """
    finds a normalized image in the given source image
    :param image_root: the image root where all image folders are stored in
    :type image_root: str
    :param norm_root: the root folder of the normalized images
    :type norm_root: str
    :param image_dir: the current image directory with the  person name
    :type image_dir: str
    :param image: the image name
    :type image: str
    :param undetected_dir: the folder to save images with undetected faces to
    :type undetected_dir: str
    :return:
    """
    cur_image_path = os.path.join(image_root, image_dir, image)
    crop_dir = os.path.join(norm_root, image_dir)
    norm_file_path = os.path.join(crop_dir, image)

    if not os.path.exists(norm_file_path):
        image_data = cv2.imread(os.path.join(image_root, image_dir, image))
        if image_data is not None:
            # image dimensions --> 96 x 96
            landmarks, detected = normalize_face(image_data, 128, False, skip_multi=True)

            # move image with undetected face to separate folder
            if detected is None or len(detected) == 0:
                print('no face detected in ' + norm_file_path)
                undetected_sub_folder = os.path.join(undetected_dir, image_dir)
                if not os.path.exists(undetected_sub_folder):
                    os.mkdir(undetected_sub_folder)
                shutil.move(cur_image_path, os.path.join(undetected_sub_folder, image))
                return False
            else:
                if not os.path.exists(crop_dir):
                    os.makedirs(crop_dir)
                cv2.imwrite(norm_file_path, detected)
                return True
    else:
        print(norm_file_path + ' already exists')

def move_to(from_path, to_path):
    shutil.move(from_path, to_path)

def start_normalization():
    """
    initiates the normalization phase for your images, also initiates a logger to follow the batch process
    """
    dir_name = input('Enter image root folder:\n--> ')
    image_root = os.path.abspath(dir_name)
    norm_dir = input('Enter the normalized images output folder:\n--> ')
    norm_root = os.path.abspath(norm_dir)
    min_image_count = int(input('Enter the minimum images per folder to normalize:\n--> '))
    image_dir_list = os.listdir(dir_name)

    print('Normalizing all images... ')
    start_time = time.time()

    undetected_dir = os.path.join(image_root, '_undetected')
    if not os.path.isdir(undetected_dir):
        os.mkdir(undetected_dir)
    if not os.path.isdir(norm_root):
        os.makedirs(norm_root)

    for index, image_dir in enumerate(image_dir_list):
        progress.set_progress(0, len(image_dir_list), index + 1)
        image_list = os.listdir(os.path.join(dir_name, image_dir))

        if len(image_list) >= min_image_count:
            for image in image_list:
                if os.path.isfile(os.path.join(image_root, image_dir, image)):
                    norm_image(image_root, norm_root, image_dir, image, undetected_dir)

    print('Completed normalization in {} seconds'.format(time.time() - start_time))

if __name__ == 'main':
    start_normalization()


"""
tricks:
1.torch-optimizer:实现了最新的一些优化器.
2.numba:import numba as nb,纯python或numpy加速,加@nb.njit或@nb.jit(nopython=True)
3.swifter:df.apply()→·df.swifter.apply()，加速pandas
4.cupy:1000万以上数据更快
5.modin:import modin.pandas as mdpd,用mdpd代替pd即可，加速pandas,加载数据和查询数据更快,统计方法pandas更快
"""
import os
import sys
import argparse
import time
import random
import wandb
from tqdm import tqdm
import numpy as np
import numba as nb
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils import data
from torch.utils.tensorboard import SummaryWriter
from torchsummary import summary
from models.module import Model
from utils.util import EarlyStopping
from data.custom_dataset import MyDataset

def train(start_epoch, num_epochs):
    early_stopping = EarlyStopping(opt, verbose=True)
    for epoch in tqdm(range(start_epoch, num_epochs)):
        torch.cuda.empty_cache()
        print("Epoch: %d" % epoch)
        train_losses = []
        model.train()
        past = time.time()
        for batch_idx, (inputs, targets) in enumerate(train_data_loader):
            # models.zero_grad()
            # optimizer.zero_grad()#当optimizer=optim.Optimizer(models.parameters())时，两者等效
            # 梯度清零
            optimizer.zero_grad()

            inputs, targets = inputs.to(device), targets.to(device)

            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()
            train_losses.append(loss.item())

            # if (batch_idx + 1) % num == 0:
            # 	print(batch_idx + 1, len(dataloader), 'Loss: %.3f' % (train_loss / num))
            # 	train_loss = 0
        now = time.time()
        train_loss = np.mean(np.array(train_losses))
        print(epoch, "loss:%.3f,time:%.2fs" % (train_loss, now - past))
        writer.add_scalar("train_loss", train_loss, epoch)
        wandb.log({"train_loss": train_loss}, step=epoch)
        train_loss = 0
        # checkpoint = {
        #     "model_state_dict": models.module.state_dict(),
        #     "opt_state_dict": optimizer.state_dict(),
        #     "epoch": epoch,
        # }

        scheduler.step()
        # the end of one epoch
        model.eval()
        val_losses = []
        with torch.no_grad():
            for batch_idx, (inputs, targets) in enumerate(validation_data_loader):
                inputs, targets = inputs.to(device), targets.to(device)

                outputs = model(inputs)
                loss = criterion(outputs, targets)
                val_losses.append(loss.item())
        val_loss = np.mean(np.array(val_losses))
        wandb.log({"val_loss": val_loss}, step=epoch)

        #####some testing#####
        print("xxxxxxx".format(xxxxxxx))
        #####some logging#####
        prefix = opt.path_to_checkpoint + opt.hidden_size + "_"
        file = prefix + "xxx_xxx_xxx.pt"
        torch.save(
            {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
            },
            opt.path_to_checkpoint
            + "%d_%d_%d_%d_%d_model.onnx"
            % (
                opt.value_dim,
                opt.moment_dim,
                opt.hidden_size,
                opt.label_dim,
                opt.LSTM_num_layers,
            ),
        )
        wandb.save("mymodel.h5")
        early_stopping(val_loss, model, opt)
        if early_stopping.early_stop:
            print("Early stopping")
            break

if __name__ == "__main__":
    # 取每个 GPU 的剩余显存数，并存放到 tmp 文件中
    os.system("nvidia-smi -q -d Memory |grep -A4 GPU|grep Free >tmp")
    memory_gpu = [int(x.split()[2]) for x in open("tmp", "r").readlines()]
    torch.cuda.set_device(np.argmax(memory_gpu))
    os.system("rm tmp")  # 删除临时生成的 tmp 文件
    ##命令行执行
    # CUDA_VISIBLE_DEVICES=0,1 python train.py
    # os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    # os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = True
    # argparse for additional flags for experiment
    parser = argparse.ArgumentParser(description="Train a network for ...")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--epochs", type=int, default=1000)
    parser.add_argument("--resume", type=bool, default=False)
    parser.add_argument("--patience", type=int, default=10)
    parser.add_argument("--path_to_checkpoint", type=str, default="../checkpoint/")

    opt = parser.parse_args()
    np.random.seed(opt.seed)
    torch.manual_seed(opt.seed)
    torch.cuda.manual_seed_all(opt.seed)

    wandb.init(project="my-project")
    wandb.config.xxx = opt.xxx
    # 准备数据
    train_dataset = MyDataset("train_dataset_path")  # 定义的数据集
    validation_dataset = MyDataset("validation_dataset_path")
    train_data_loader = data.DataLoader(
        train_dataset,
        batch_size=128,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        drop_last=True,
    )
    validation_data_loader = data.DataLoader(
        validation_dataset,
        batch_size=xxxx,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
        drop_last=True,
    )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # device_ids = [0, 1]

    # 自定义的model
    model = Model(opt)
    model.to(device)
    model.initialize_weights()

    # 并行运算，如果需要的话
    model = nn.DataParallel(model, device_ids=device_ids).to(device)
    # summary(model, input_size=(channels, H, W))
    # hl.build_graph(model, torch.zeros([1, 2, 3]))

    # loss function， 比如交叉熵
    criterion = nn.CrossEntropyLoss()
    criterion.to(device)
    wandb.watch(model, criterion)

    # optimizer，比如Adam
    optimizer = optim.Adam(model.parameters())
    # 调整学习率
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.1)

    # 训练
    num_epochs = opt.epochs

    start_epoch = 0
    if opt.resume:
        ckpt = torch.load(
            opt.path_to_checkpoint
            + "%d_%d_%d_%d_%d_model.onnx"
            % (
                opt.value_dim,
                opt.moment_dim,
                opt.hidden_size,
                opt.label_dim,
                opt.LSTM_num_layers,
            )
        )  # custom method for loading last checkpoint
        model.load_state_dict(ckpt["model_state_dict"])
        start_epoch = ckpt["epoch"]
        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
        model = wandb.restore("mymodel.h5")

        print("last checkpoint restored")
        ...
    # default `log_dir`is "runs" - we'll be more specific here
    writer = SummaryWriter("runs/models")

    train(start_epoch, num_epochs)


#coding: utf-8
from ..utils import bbox_helper
from ..utils.pair_helper import pair_nms
from ..utils.pair_helper import pair_box_transform
from ..utils.pair_helper import pair_pos_score
import pdb
import ctypes

import math
import numpy as np
import torch
from torch.autograd import Variable
import logging
import json
logger = logging.getLogger('global')
history = [0, 0]

def to_np_array(x):
    if x is None:
        return x
    if isinstance(x, Variable): x = x.data
    return x.cpu().float().numpy() if torch.is_tensor(x) else x

def compute_proposal_targets2(all_rois, pair_rois, num_classes, cfg, gt_bboxes, gt_assos, roiTable, image_info, predict_type='rcnn', ignore_regions=None):
    B = len(image_info)
    C = num_classes
    N = pair_rois.shape[0]  # 267，前一步makepair生成了 267 个pair
    gt_cls = np.array([0 for _ in range(N)]).astype(np.int32)
    all_rois, gt_assos, image_info = map(to_np_array, [all_rois, gt_assos, image_info])
    pos_ix = []
    neg_ix = []
    gt_cls = [0 for __ in range(N)]
    gt_binary_cls = [[0 for __ in range(C)] for _ in range(N)]
    for b_ix in range(B):
        idx = np.where(all_rois[:,0] == b_ix)[0]
        rois = all_rois[idx, 1: 5]  # (100, 4) 该图片下的检测bbox
        gts = gt_bboxes[b_ix]  # (16, 5)
        gts = to_np_array(gts) # (1, 13) 该 图片 下的 gt 关联
        gtasso = gt_assos[b_ix].cpu().numpy()  # [idx1, idx2, item['label']] + item['bbox1'] + [item['label1']] + item['bbox2'] + [item['label2']]
        dt_pairs = pair_rois[pair_rois[:, 0] == b_ix]
        R = rois.shape[0]
        G = gts.shape[0]
        if R == 0:
            continue
        if G > 0:
            overlaps = bbox_helper.bbox_iou_overlaps(torch.from_numpy(rois), torch.from_numpy(gts)).numpy() # (100, 16)
            argmax_overlaps = overlaps.argmax(axis=1)  # 每一个检测框与gt iou重叠最大的 bbox 的索引 (100,)
            max_overlaps = overlaps.max(axis=1)  # 每一个检测框与gt iou重叠最大的值 (100, )
            if cfg.get('allow_low_quality_match', False):
                gt_argmax_overlaps = overlaps.argmax(axis=0)
                gt_max_overlaps = overlaps.max(axis=0)
                gt_pos_g_ix = np.where(gt_max_overlaps > cfg['negative_iou_thresh'])[0]
                gt_pos_r_ix = gt_argmax_overlaps[gt_pos_g_ix]
                gt_pos_g_ix, return_index = np.unique(gt_pos_g_ix, return_index=True)
                gt_pos_r_ix = gt_pos_r_ix[return_index]
                max_overlaps[gt_pos_r_ix] = 1.0

            pos_r_ix = np.where(max_overlaps > cfg['positive_iou_thresh'])[0] # (91, ) 选择该图下的检测的正例的索引
            pos_g_ix = argmax_overlaps[pos_r_ix]  # (91, ) 改图下的检测的正例的匹配到的gt的索引
            pos_r_ix, return_index = np.unique(pos_r_ix, return_index=True)
            pos_g_ix = pos_g_ix[return_index]
            GG = gtasso.shape[0]
            #for m in range(dt_pairs.shape[0]):
            #    x = dt_pairs[m, 1]
            #    y = dt_pairs[m, 2]
            #    if x in pos_r_ix and y in pos_r_ix:
            #        continue
            #    else:
            #        id = roiTable[(x,y)]
            #        neg_ix.append(id)
            #        gt_cls[id] = 0

            for k in range(GG):  # 对于每一个 gt
                if gtasso[k][2] == 0:  # 如果是背景，则略过
                    continue
                x = gtasso[k][0]  # pair id 1
                y = gtasso[k][1]  # pair id 2
                for i in range(len(pos_r_ix)):
                    for j in range(len(pos_r_ix)):
                        if x == pos_g_ix[i] and y == pos_g_ix[j]:  # 只有都匹配到了，才算 pos
                            rx = idx[pos_r_ix[i]]
                            ry = idx[pos_r_ix[j]]
                            if not (rx, ry) in roiTable:
                                continue
                            assert((rx,ry) in roiTable)
                            id = roiTable[(rx,ry)]
                            pos_ix.append(id)
                            if predict_type == 'rcnn':
                                gt_cls[id] = gtasso[k][2]
                                gt_binary_cls[id][gtasso[k][2].astype(np.int32)] = 1
                            else:
                                gt_cls[id] = 1
                                gt_binary_cls[id][0] = 1

    neg_ix = np.array(list(set([_ for _ in range(N)]) - set(pos_ix))) # 实际本batch中得到的负例的个数，237
    batch_size = cfg['batch_size']
    all_num = batch_size * B  # 128 * 4 = 512
    #pos_num = int(all_num * cfg['positive_percent'])  # 128
    pos_ix = np.array(pos_ix)  # 实际本batch中得到的正例的个数，根绝某规则选择出的正例(iou。。。。) (30, )
    #if pos_ix.shape[0] > pos_num:  # 如果实际的正例比我们想要的正例数目还要多，则进行随机算则
    #    keep_ix = np.random.choice(len(pos_ix), size = pos_num, replace = True)  # 随机选择128个正例，降采样
    #    pos_ix = pos_ix[keep_ix]
    #    if cfg.get('dynamic_batch_size', False):  # 动态batch size
    #        all_num = int(len(pos_ix) / cfg['positive_percent'])  # 增大 batch size，并不一定是 512 了

    #neg_num = all_num - len(pos_ix)  # 我们想要的负例的个数 (482, )
    #if len(neg_ix) == 0:
    #    neg_ix = np.array([0])
    #keep_ix = np.random.choice(len(neg_ix), size = neg_num, replace = True)
    #neg_ix = neg_ix[keep_ix]  # (482, )
    neg_cls = [0 for __ in range(len(neg_ix))]
    neg_binary_cls = [[0 for __ in range(C)] for _ in range(len(neg_ix))]
    gt_cls = np.array(gt_cls)
    gt_cls = gt_cls[pos_ix.astype(np.int32)]
    gt_binary_cls = np.array(gt_binary_cls)
    gt_binary_cls = gt_binary_cls[pos_ix.astype(np.int32)]
    neg_ix = np.array(neg_ix)
    neg_cls = np.array(neg_cls)
    neg_binary_cls = np.array(neg_binary_cls)

    if len(pos_ix) == 0:
        dt = neg_ix
        gt_cls = neg_cls
        gt_binary_cls = neg_binary_cls
    else:
        dt = np.hstack([pos_ix, neg_ix])  # 正负例的索引 stack 起来
        gt_cls = np.hstack([gt_cls, neg_cls])
        gt_binary_cls = np.vstack([gt_binary_cls, neg_binary_cls])

    pdb.set_trace()
    if cfg.get('cls_type', 'softmax') == 'softmax':
        return dt, gt_cls, gt_cls.shape[0]
    else:
        return dt, gt_binary_cls, gt_binary_cls.shape[0]

def compute_proposal_targets_gt(all_rois, pair_rois, num_classes, cfg, gt_bboxes, gt_assos, roiTable, image_info, predict_type='rcnn', ignore_regions=None):
    '''
    gt_assos: [N, 3] rois1_ix, rois2_ix, cls
    '''
    #pdb.set_trace()
    B = len(image_info)
    C = num_classes
    N = pair_rois.shape[0]
    gt_cls = np.array([0 for _ in range(N)]).astype(np.int32)
    all_rois, gt_assos, image_info = map(to_np_array, [all_rois, gt_assos, image_info])
    pos_ix = []
    gt_cls = [0 for __ in range(N)]
    gt_binary_cls = [[0 for __ in range(C)] for _ in range(N)]
    for b_ix in range(B):
        idx = np.where(all_rois[:,0] == b_ix)[0]
        rois = all_rois[idx, 1: 5]
        gts = gt_bboxes[b_ix]
        gts = to_np_array(gts)
        gtasso = gt_assos[b_ix].cpu().numpy()
        R = rois.shape[0]
        G = gtasso.shape[0]
        if R == 0:
            continue
        if G > 0:
            overlaps = bbox_helper.bbox_iou_overlaps(torch.from_numpy(rois), torch.from_numpy(gts)).numpy()
            argmax_overlaps = overlaps.argmax(axis=1)
            max_overlaps = overlaps.max(axis=1)
            if cfg.get('allow_low_quality_match', False):
                gt_argmax_overlaps = overlaps.argmax(axis=0)
                gt_max_overlaps = overlaps.max(axis=0)
                gt_pos_g_ix = np.where(gt_max_overlaps > cfg['negative_iou_thresh'])[0]
                gt_pos_r_ix = gt_argmax_overlaps[gt_pos_g_ix]
                gt_pos_g_ix, return_index = np.unique(gt_pos_g_ix, return_index=True)
                gt_pos_r_ix = gt_pos_r_ix[return_index]
                max_overlaps[gt_pos_r_ix] = 1.0

            pos_r_ix = np.where(max_overlaps > cfg['positive_iou_thresh'])[0]
            pos_g_ix = argmax_overlaps[pos_r_ix]
            pos_r_ix, return_index = np.unique(pos_r_ix, return_index=True)
            pos_g_ix = pos_g_ix[return_index]
            GG = gtasso.shape[0]
            added_gt = set()
            for k in range(GG):
                if gtasso[k][2] == 0:
                    continue
                x = gtasso[k][0]
                y = gtasso[k][1]
                for i in range(len(pos_r_ix)):
                    for j in range(len(pos_r_ix)):
                        if x == pos_g_ix[i] and y == pos_g_ix[j]:
                            rx = idx[pos_r_ix[i]]
                            ry = idx[pos_r_ix[j]]
                            if not (rx, ry) in roiTable:
                                continue
                            assert((rx,ry) in roiTable)
                            id = roiTable[(rx,ry)]
                            pos_ix.append(id)
                            added_gt.add(k)
                            if predict_type == 'rcnn':
                                gt_cls[id] = gtasso[k][2]
                                gt_binary_cls[id][gtasso[k][2].astype(np.int32)] = 1
                            else:
                                gt_cls[id] = 1
                                gt_binary_cls[id][0] = 1
            overlaps = overlaps.T
            argmax_overlaps = overlaps.argmax(axis=1)
            max_overlaps = overlaps.max(axis=1)

            no_added_gt = set(list(range(GG))) - added_gt
            if len(no_added_gt) <= 1:
                continue
            for k in no_added_gt: # for not added gt
                #pdb.set_trace()
                x = int(gtasso[k][0])
                y = int (gtasso[k][1])
                if max_overlaps[x] == 0 or max_overlaps[y] == 0:
                    continue
                rx = argmax_overlaps[x]
                ry = argmax_overlaps[y]
                rx = idx[rx]
                ry = idx[ry]
                if not (rx, ry) in roiTable:
                    continue
                assert((rx,ry) in roiTable)
                id = roiTable[(rx,ry)]
                pos_ix.append(id)
                if predict_type == 'rcnn':
                    gt_cls[id] = gtasso[k][2] # todo 类别没有对应好，如果要对应类别，就不能去max
                    gt_binary_cls[id][gtasso[k][2].astype(np.int32)] = 1
                else:
                    gt_cls[id] = 1
                    gt_binary_cls[id][0] = 1

    neg_ix = np.array(list(set([_ for _ in range(N)]) - set(pos_ix)))
    #pdb.set_trace()
    batch_size = cfg['batch_size']
    all_num = batch_size * B
    pos_num = int(all_num * cfg['positive_percent'])
    pos_ix = np.array(pos_ix)
    if pos_ix.shape[0] > pos_num:
        keep_ix = np.random.choice(len(pos_ix), size = pos_num, replace = True)
        pos_ix = pos_ix[keep_ix]
        if cfg.get('dynamic_batch_size', False):
            all_num = int(len(pos_ix) / cfg['positive_percent'])

    neg_num = all_num - len(pos_ix)
    if len(neg_ix) == 0:
        neg_ix = np.array([0])
    keep_ix = np.random.choice(len(neg_ix), size = neg_num, replace = True)
    neg_ix = neg_ix[keep_ix]
    neg_cls = [0 for __ in range(len(keep_ix))]
    neg_binary_cls = [[0 for __ in range(C)] for _ in range(len(keep_ix))]
    gt_cls = np.array(gt_cls)
    gt_cls = gt_cls[pos_ix.astype(np.int32)]
    gt_binary_cls = np.array(gt_binary_cls)
    gt_binary_cls = gt_binary_cls[pos_ix.astype(np.int32)]
    neg_ix = np.array(neg_ix)
    neg_cls = np.array(neg_cls)
    neg_binary_cls = np.array(neg_binary_cls)

    if len(pos_ix) == 0:
        dt = neg_ix
        gt_cls = neg_cls
        gt_binary_cls = neg_binary_cls
    else:
        dt = np.hstack([pos_ix, neg_ix])
        gt_cls = np.hstack([gt_cls, neg_cls])
        gt_binary_cls = np.vstack([gt_binary_cls, neg_binary_cls])
    idx = np.arange(len(dt))
    np.random.shuffle(idx)
    dt = dt[idx]
    gt_cls = gt_cls[idx]
    gt_binary_cls = gt_binary_cls[idx]
    #pdb.set_trace()
    if cfg.get('cls_type', 'softmax') == 'softmax':
        return dt, gt_cls, gt_cls.shape[0]
    else:
        return dt, gt_binary_cls, gt_binary_cls.shape[0]

def compute_proposal_targets(all_rois, pair_rois, num_classes, cfg, gt_bboxes, gt_assos, roiTable, image_info, predict_type='rcnn', ignore_regions=None):
    '''
    gt_assos: [N, 3] rois1_ix, rois2_ix, cls
    '''
    #pdb.set_trace()
    B = len(image_info)
    C = num_classes
    N = pair_rois.shape[0]
    gt_cls = np.array([0 for _ in range(N)]).astype(np.int32)
    all_rois, gt_assos, image_info = map(to_np_array, [all_rois, gt_assos, image_info])
    pos_ix = []
    gt_cls = [0 for __ in range(N)]
    gt_binary_cls = [[0 for __ in range(C)] for _ in range(N)]
    for b_ix in range(B):
        idx = np.where(all_rois[:,0] == b_ix)[0]
        rois = all_rois[idx, 1: 5]
        gts = gt_bboxes[b_ix]
        gts = to_np_array(gts)
        gtasso = gt_assos[b_ix].cpu().numpy()
        R = rois.shape[0]
        G = gts.shape[0]
        if R == 0:
            continue
        if G > 0:
            overlaps = bbox_helper.bbox_iou_overlaps(torch.from_numpy(rois), torch.from_numpy(gts)).numpy()
            argmax_overlaps = overlaps.argmax(axis=1)
            max_overlaps = overlaps.max(axis=1)
            if cfg.get('allow_low_quality_match', False):
                gt_argmax_overlaps = overlaps.argmax(axis=0)
                gt_max_overlaps = overlaps.max(axis=0)
                gt_pos_g_ix = np.where(gt_max_overlaps > cfg['negative_iou_thresh'])[0]
                gt_pos_r_ix = gt_argmax_overlaps[gt_pos_g_ix]
                gt_pos_g_ix, return_index = np.unique(gt_pos_g_ix, return_index=True)
                gt_pos_r_ix = gt_pos_r_ix[return_index]
                max_overlaps[gt_pos_r_ix] = 1.0

            pos_r_ix = np.where(max_overlaps > cfg['positive_iou_thresh'])[0]
            pos_g_ix = argmax_overlaps[pos_r_ix]
            pos_r_ix, return_index = np.unique(pos_r_ix, return_index=True)
            pos_g_ix = pos_g_ix[return_index]
            GG = gtasso.shape[0]
            for k in range(GG):
                if gtasso[k][2] == 0:
                    continue
                x = gtasso[k][0]
                y = gtasso[k][1]
                for i in range(len(pos_r_ix)):
                    for j in range(len(pos_r_ix)):
                        if x == pos_g_ix[i] and y == pos_g_ix[j]:
                            rx = idx[pos_r_ix[i]]
                            ry = idx[pos_r_ix[j]]
                            if not (rx, ry) in roiTable:
                                continue
                            assert((rx,ry) in roiTable)
                            id = roiTable[(rx,ry)]
                            pos_ix.append(id)
                            if predict_type == 'rcnn':
                                gt_cls[id] = 1
                                gt_binary_cls[id][gtasso[k][2].astype(np.int32)] = 1
                            else:
                                gt_cls[id] = 1
                                gt_binary_cls[id][0] = 1
    neg_ix = np.array(list(set([_ for _ in range(N)]) - set(pos_ix)))
    neg_cls = [0 for __ in range(len(neg_ix))]
    neg_binary_cls = [[0 for __ in range(C)] for _ in range(len(neg_ix))]
    gt_cls = np.array(gt_cls)
    pos_ix = np.array(pos_ix)
    gt_cls = gt_cls[pos_ix.astype(np.int32)]
    gt_binary_cls = np.array(gt_binary_cls)
    gt_binary_cls = gt_binary_cls[pos_ix.astype(np.int32)]
    neg_ix = np.array(neg_ix)
    neg_cls = np.array(neg_cls)
    neg_binary_cls = np.array(neg_binary_cls)

    if len(pos_ix) == 0:
        dt = neg_ix
        gt_cls = neg_cls
        gt_binary_cls = neg_binary_cls
    else:
        dt = np.hstack([pos_ix, neg_ix])
        gt_cls = np.hstack([gt_cls, neg_cls])
        if len(neg_binary_cls) != 0:
            gt_binary_cls = np.vstack([gt_binary_cls, neg_binary_cls])
    idx = np.arange(len(dt))
    dt = dt[idx]
    gt_cls = gt_cls[idx]
    gt_binary_cls = gt_binary_cls[idx]
    if cfg.get('cls_type', 'softmax') == 'softmax':
        return dt, gt_cls, gt_cls.shape[0]
    else:
        return dt, gt_binary_cls, gt_binary_cls.shape[0]

def predict_assos(rois, pair_rois, pred_cls, image_info, cfg, tocaffe):
    '''
    :param cfg: config
    :param rois: [N, k] k>=7, batch_ix, x1, y1, x2, y2, score, cls
    :param pred_assos:[N, num_classes * 4, 1, 1]
    :param image_info:[N, 3]
    :return: assos: [M, 13], batch_ix, ax1, ay1, ax2, ay2, acls, bx1, by1, bx2, by2, bcls, cls, score
    '''
    rois, pair_rois, pred_cls = map(to_np_array, [rois, pair_rois, pred_cls])
    N, num_classes = pred_cls.shape[0], 3
    B = max(rois[:, 0].astype(np.int32))+1
    nmsed_assos = [np.zeros((1,13))]
    rois1 = rois[pair_rois[:, 1].astype(np.int32)]
    rois2 = rois[pair_rois[:, 2].astype(np.int32)]
    #pdb.set_trace()
    for cls in range(1, num_classes):
        #if cls == 1:
            # in body-face pair, the cls of object must be face
        #    idx = np.where(rois2[:,6] == 2)
        #else:
            # in body-hand pair, the cls of object must be hand
        #    idx = np.where(rois2[:,6] == 3)
        #idx = np.where(pred_cls[:, cls] == pred_cls.max(axis=0))
        #pdb.set_trace()
        idx = np.where(rois2[:, 6] > 0)
        tmp_pred_cls = pred_cls[idx]
        tmp_rois1 = rois1[idx]
        tmp_rois2 = rois2[idx]
        tmp_pair = pair_rois[idx]
        #scores = (tmp_pred_cls[:,cls] * tmp_rois1[:, 5] * tmp_rois2[:, 5]).reshape((-1,1))

        scores = tmp_pred_cls.reshape((-1,1))

        pair_cls = tmp_rois2[:, 6].copy()
        pair_cls[pair_cls == 2] = 1
        pair_cls[pair_cls == 3] = 2
        pair_cls = pair_cls.reshape((-1, 1))

        pair_cls = np.array([cls for _ in range(scores.shape[0])]).reshape((-1,1))
        batch = tmp_pair[:,0]
        batch = batch.reshape((-1,1))
        #nmsed_assos.append(np.hstack([batch, tmp_rois1[:,1:5], tmp_rois1[:,6].reshape((-1,1)), tmp_rois2[:, 1:5], tmp_rois2[:,6].reshape((-1,1)), pair_cls, scores]))
        batch_asso = np.hstack([batch, tmp_rois1[:,1:5], tmp_rois1[:,6].reshape((-1,1)), \
                tmp_rois2[:, 1:5], tmp_rois2[:,6].reshape((-1,1)), pair_cls, scores])
        nmsed_assos.append(batch_asso)

    if tocaffe:
        nmsed_assos = np.zeros((1,13))
    else:
        nmsed_assos = np.vstack(nmsed_assos)
    if cfg['top_n'] > 0:
        top_n_assos = [[0,0.0,0.0,0.0,0.0,0,0.0,0.0,0.0,0.0,0,0,0]]
        for b_ix in range(B):
            assos = nmsed_assos[nmsed_assos[:, 0] == b_ix]
            if assos.size == 0: continue
            scores = assos[:, -1]
            order = scores.argsort()[::-1][:cfg['top_n']]
            assos = assos[order]
            if cfg.get('score_thresh', None) is not None:
                keep = assos[:, -1] > cfg['score_thresh']
                #keep = np.bitwise_and(0.5 > assos[:, -1], assos[:, -1] > cfg['score_thresh'])
                assos = assos[keep]
            top_n_assos.append(assos)
        nmsed_assos = np.vstack(top_n_assos)
    return nmsed_assos

def predict_assos_2(rois, pair_rois, pred_cls, image_info, cfg, tocaffe):
    '''
    :param cfg: config
    :param rois: [N, k] k>=7, batch_ix, x1, y1, x2, y2, score, cls
    :param pred_assos:[N, num_classes * 4, 1, 1]
    :param image_info:[N, 3]
    :return: assos: [M, 13], batch_ix, ax1, ay1, ax2, ay2, acls, bx1, by1, bx2, by2, bcls, cls, score
    '''
    rois, pair_rois, pred_cls = map(to_np_array, [rois, pair_rois, pred_cls])
    N, num_classes = pred_cls.shape[0:2]
    B = max(rois[:, 0].astype(np.int32))+1
    nmsed_assos = [np.zeros((1,13))]
    #pdb.set_trace()
    rois1 = rois[pair_rois[:, 1].astype(np.int32)]
    rois2 = rois[pair_rois[:, 2].astype(np.int32)]
    for cls in range(1, num_classes):
        #if cls == 1:
            # in body-face pair, the cls of object must be face
        #    idx = np.where(rois2[:,6] == 2)
        #else:
            # in body-hand pair, the cls of object must be hand
        #    idx = np.where(rois2[:,6] == 3)
        idx = np.where(rois2[:, 6] > 0)
        tmp_pred_cls = pred_cls[idx]
        tmp_rois1 = rois1[idx]
        tmp_rois2 = rois2[idx]
        tmp_pair = pair_rois[idx]
        scores = (tmp_pred_cls[:,cls] * tmp_rois1[:, 5] * tmp_rois2[:, 5]).reshape((-1,1))
        #scores = tmp_pred_cls[:, cls].reshape((-1,1))
        # pair_cls = np.array([cls for _ in range(scores.shape[0])]).reshape((-1,1))
        pair_cls = tmp_rois2[:, 6].copy()
        pair_cls[pair_cls == 2] = 1
        pair_cls[pair_cls == 3] = 2
        pair_cls = pair_cls.reshape((-1, 1))
        batch = tmp_pair[:,0]
        batch = batch.reshape((-1,1))
        #nmsed_assos.append(np.hstack([batch, tmp_rois1[:,1:5], tmp_rois1[:,6].reshape((-1,1)), tmp_rois2[:, 1:5], tmp_rois2[:,6].reshape((-1,1)), pair_cls, scores]))
        batch_asso = np.hstack([batch, tmp_rois1[:,1:5], tmp_rois1[:,6].reshape((-1,1)),
            tmp_rois2[:, 1:5], tmp_rois2[:,6].reshape((-1,1)), pair_cls, scores])
        if cfg.get('cls_top_n', 0) > 0:
            for b_ix in range(B):
                assos = batch_asso[batch_asso[:, 0] == b_ix]
                if assos.size == 0: continue
                scores = assos[:, -1]
                order = scores.argsort()[::-1][:cfg['cls_top_n']]
                assos = assos[order]
                if cfg.get('use_filter', None) is not None:
                    cls_triplet = cfg['asso_triplet'][cls - 1]
                    tmp = []
                    for line in assos:
                        acls = line[5].astype(int)
                        bcls = line[10].astype(int)
                        if (str(acls) in cls_triplet.keys() and bcls in cls_triplet[str(acls)]):
                           tmp.append(line)
                    if len(tmp) > 0:
                        assos = np.vstack(tmp)
                    else:
                        assos = []
                if len(assos) > 0:
                    nmsed_assos.append(assos)
        else:
            nmsed_assos.append(batch_asso)

    if tocaffe:
        nmsed_assos = np.zeros((1,13))
    else:
        nmsed_assos = np.vstack(nmsed_assos)
        if cfg.get('pair_nms', None):
            if cfg['pre_top_n'] > 0:
                top_n_assos = []
                for b_ix in range(B):
                    assos = nmsed_assos[nmsed_assos[:, 0] == b_ix]
                    if assos.size == 0: continue
                    scores = assos[:, -1]
                    order = scores.argsort()[::-1][:cfg['pre_top_n']]
                    assos = assos[order]
                    if cfg.get('pre_score_thresh', None) is not None:
                        keep = assos[:, -1] > cfg['pre_score_thresh']
                        assos = assos[keep]
                    top_n_assos.append(assos)
                nmsed_assos = np.vstack(top_n_assos)
            nmsed_assos = pair_nms(nmsed_assos, cfg['pair_nms'])
    if cfg['top_n'] > 0:
        top_n_assos = [[0,0.0,0.0,0.0,0.0,0,0.0,0.0,0.0,0.0,0,0,0]]
        for b_ix in range(B):
            assos = nmsed_assos[nmsed_assos[:, 0] == b_ix]
            if assos.size == 0: continue
            scores = assos[:, -1]
            order = scores.argsort()[::-1][:cfg['top_n']]
            assos = assos[order]
            if cfg.get('score_thresh', None) is not None:
                keep = assos[:, -1] > cfg['score_thresh']
                assos = assos[keep]
            top_n_assos.append(assos)
        nmsed_assos = np.vstack(top_n_assos)
    return nmsed_assos


# Copyright 2021 Tomoki Hayashi
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

"""Test code for MelGAN modules."""

import numpy as np
import pytest
import torch

from espnet2.gan_tts.hifigan.loss import DiscriminatorAdversarialLoss
from espnet2.gan_tts.hifigan.loss import FeatureMatchLoss
from espnet2.gan_tts.hifigan.loss import GeneratorAdversarialLoss
from espnet2.gan_tts.melgan import MelGANGenerator
from espnet2.gan_tts.melgan import MelGANMultiScaleDiscriminator

def make_melgan_generator_args(**kwargs):
    defaults = dict(
        in_channels=80,
        out_channels=1,
        kernel_size=7,
        channels=32,
        bias=True,
        upsample_scales=[4, 4],
        stack_kernel_size=3,
        stacks=2,
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.2},
        pad="ReflectionPad1d",
        pad_params={},
        use_final_nonlinear_activation=True,
        use_weight_norm=True,
    )
    defaults.update(kwargs)
    return defaults

def make_melgan_discriminator_args(**kwargs):
    defaults = dict(
        in_channels=1,
        out_channels=1,
        scales=2,
        downsample_pooling="AvgPool1d",
        # follow the official implementation setting
        downsample_pooling_params={
            "kernel_size": 4,
            "stride": 2,
            "padding": 1,
            "count_include_pad": False,
        },
        kernel_sizes=[5, 3],
        channels=16,
        max_downsample_channels=32,
        bias=True,
        downsample_scales=[2, 2],
        nonlinear_activation="LeakyReLU",
        nonlinear_activation_params={"negative_slope": 0.2},
        pad="ReflectionPad1d",
        pad_params={},
        use_weight_norm=True,
    )
    defaults.update(kwargs)
    return defaults

@pytest.mark.skipif(
    "1.6" in torch.__version__,
    reason="group conv in pytorch 1.6 has an issue. "
    "See https://github.com/pytorch/pytorch/issues/42446.",
)
@pytest.mark.parametrize(
    "dict_g, dict_d",
    [
        ({}, {}),
        ({}, {"scales": 4}),
        ({}, {"kernel_sizes": [7, 5]}),
        ({}, {"max_downsample_channels": 128}),
        ({}, {"downsample_scales": [4, 4]}),
        ({}, {"pad": "ConstantPad1d", "pad_params": {"value": 0.0}}),
        ({}, {"nonlinear_activation": "ReLU", "nonlinear_activation_params": {}}),
    ],
)
def test_melgan_generator_and_discriminator(dict_g, dict_d):
    # setup
    batch_size = 2
    batch_length = 512
    args_g = make_melgan_generator_args(**dict_g)
    args_d = make_melgan_discriminator_args(**dict_d)
    y = torch.randn(batch_size, 1, batch_length)
    c = torch.randn(
        batch_size,
        args_g["in_channels"],
        batch_length // np.prod(args_g["upsample_scales"]),
    )
    model_g = MelGANGenerator(**args_g)
    model_d = MelGANMultiScaleDiscriminator(**args_d)
    gen_adv_criterion = GeneratorAdversarialLoss()
    dis_adv_criterion = DiscriminatorAdversarialLoss()
    feat_match_criterion = FeatureMatchLoss()
    optimizer_g = torch.optim.Adam(model_g.parameters())
    optimizer_d = torch.optim.Adam(model_d.parameters())

    # check generator trainable
    y_hat = model_g(c)
    p_hat = model_d(y_hat)
    adv_loss = gen_adv_criterion(p_hat)
    with torch.no_grad():
        p = model_d(y)
    fm_loss = feat_match_criterion(p_hat, p)
    loss_g = adv_loss + fm_loss
    optimizer_g.zero_grad()
    loss_g.backward()
    optimizer_g.step()

    # check discriminator trainable
    p = model_d(y)
    p_hat = model_d(y_hat.detach())
    real_loss, fake_loss = dis_adv_criterion(p_hat, p)
    loss_d = real_loss + fake_loss
    optimizer_d.zero_grad()
    loss_d.backward()
    optimizer_d.step()

try:
    import parallel_wavegan  # NOQA

    is_parallel_wavegan_available = True
except ImportError:
    is_parallel_wavegan_available = False

@pytest.mark.skipif(
    not is_parallel_wavegan_available, reason="parallel_wavegan is not installed."
)
def test_parallel_wavegan_compatibility():
    from parallel_wavegan.models import MelGANGenerator as PWGMelGANGenerator

    model_pwg = PWGMelGANGenerator(**make_melgan_generator_args())
    model_espnet2 = MelGANGenerator(**make_melgan_generator_args())
    model_espnet2.load_state_dict(model_pwg.state_dict())
    model_pwg.eval()
    model_espnet2.eval()

    with torch.no_grad():
        c = torch.randn(5, 80)
        out_pwg = model_pwg.inference(c)
        out_espnet2 = model_espnet2.inference(c)
        np.testing.assert_array_equal(
            out_pwg.cpu().numpy(),
            out_espnet2.cpu().numpy(),
        )


from flask import Flask, jsonify, render_template, request, redirect, url_for
import os
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import load_model
from werkzeug.utils import secure_filename
import numpy as np
import shutil
import time

app = Flask(__name__)

UPLOAD_FOLDER = "static/target/img"

@app.route('/predict')
def predict(clases):

    ImgGenerator = ImageDataGenerator(rescale = 1./255)

    classificador = load_model('best_cls.h5')

    target_datagen = ImgGenerator.flow_from_directory(batch_size=1,
                                                      directory='static/target',
                                                      target_size=(256, 256),
                                                      class_mode=None)

    result = classificador.predict(next(target_datagen))

    return f'{clases[np.argmax(result[0])]}: {max(result[0]*100):0.2f} %'

@app.route("/", methods=['GET','POST'])
def index():
    if request.method == 'POST':
        img_file = request.files['image']

        if img_file:
            img_location = os.path.join(UPLOAD_FOLDER, img_file.filename)
            img_file.save(img_location)
            start = time.time()
            pred = predict(categorias)
            print(f'Tiempo de ejecucion: {time.time()-start} -- Resultado: {pred}')

            # Move a file from the directory d1 to d2
            new_location = os.path.join("static/", img_file.filename)
            shutil.move(img_location, new_location)

            return render_template('index.html', prediction=pred, image_loc=img_file.filename)

    return render_template('index.html', prediction='', image_loc=None)

# ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}

if __name__ == '__main__':
    categorias = {0: 'Ambulancia',
                  1: 'Bomberos',
                  2: 'Patrulla',
                  3: 'Otro tipo'}

    app.run(debug=True)


#!/usr/bin/env python

# Copyright (c) Facebook, Inc. and its affiliates.

# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
"""
Class definition for Predictors
"""
import copy

import cv2
import numpy as np

n_class = 18
min_patch_std = 1
angle_dependence = [0.25, 0.5, 0.25]

try:
    xrange
except NameError:
    xrange = range

# Given image, returns image point and theta to grasp
class Predictors:
    """
    This class contains functionality to sample grasp patches from the scene image.
    """

    def __init__(self, img, grasp_obj=None):
        """
        The constructor for :class:`Predictors` class.

            :param img: An image of the scene
            :param grasp_obj: An object that contains grasp model operations
            :type img: np.ndarray
            :type grasp_obj: GraspTorchObj
        """
        self.img = img
        self.img_h, self.img_w, self.img_c = self.img.shape
        self.grasp_obj = grasp_obj

    def random_grasp(self):
        """
        Samples a random grasp configuration.

        :returns: A grasp configuration (height, width, angle) of a random grasp
        :rtype: tuple
        """
        h_g = np.random.randint(self.img_h)
        w_g = np.random.randint(self.img_w)
        t_g = np.random.randint(n_class)
        return h_g, w_g, t_g

    def center_grasp(self):
        """
        Select the center grasp configuration.

        :returns: A grasp configuration (height, width, angle) of the center grasp
        :rtype: tuple
        """
        h_g = np.int(self.img_h / 2)
        w_g = np.int(self.img_w / 2)
        t_g = np.int(n_class / 2)
        return h_g, w_g, t_g

    def graspNet_grasp(self, patch_size=300, num_samples=128):
        """
        Select grasp based on the grasp model.

        :param patch_size: Size of patch to be sampled
        :param num_samples: Number of patches to run
        :type patch_size: int
        :type num_samples: int
        """
        self.patch_size = patch_size
        half_patch_size = np.int(patch_size / 2) + 1
        h_range = self.img_h - patch_size - 2
        w_range = self.img_w - patch_size - 2

        # Initialize random patch points
        patch_hs = np.random.randint(h_range,
                                     size=num_samples) + half_patch_size
        patch_ws = np.random.randint(w_range,
                                     size=num_samples) + half_patch_size

        patch_Is = np.zeros((num_samples, patch_size,
                             patch_size, self.img_c))
        patch_Is_resized = np.zeros((num_samples,
                                     self.grasp_obj.image_size,
                                     self.grasp_obj.image_size,
                                     self.img_c))
        for looper in xrange(num_samples):
            isWhiteFlag = 1
            while isWhiteFlag == 1:
                patch_hs[looper] = np.random.randint(h_range) + half_patch_size
                patch_ws[looper] = np.random.randint(w_range) + half_patch_size
                h_b = patch_hs[looper] - half_patch_size
                w_b = patch_ws[looper] - half_patch_size
                patch_Is[looper] = self.img[h_b: h_b + patch_size,
                                   w_b: w_b + patch_size]
                # Make sure that the input has a minimal
                # amount of standard deviation from mean. If
                # not resample
                if patch_Is[looper].std() > min_patch_std:
                    isWhiteFlag = 0
                else:
                    isWhiteFlag = 1
            patch_Is_resized[looper] = cv2.resize(patch_Is[looper],
                                                  (self.grasp_obj.image_size,
                                                   self.grasp_obj.image_size),
                                                  interpolation=cv2.INTER_CUBIC)
        angle_labels = [0 for _ in patch_Is_resized]
        robot_labels = [0 for _ in patch_Is_resized]
        full_x = np.array([self.img for _ in patch_Is_resized])
        param_dict = {'x': patch_Is_resized,
                      'h': patch_hs,
                      'w': patch_ws,
                      'angle_labels': angle_labels,
                      'robot_labels': robot_labels,
                      'full_x': full_x}
        self.vals = self.grasp_obj.test_one_batch(**param_dict)

        # Normalizing angle uncertainity
        wf = angle_dependence
        self.norm_vals = copy.deepcopy(self.vals)
        for looper in xrange(num_samples):
            for norm_looper in xrange(n_class):
                n_val = (wf[1] * self.norm_vals[looper, norm_looper] +
                         wf[0] * self.norm_vals[looper, (norm_looper -
                                                         1) % n_class] +
                         wf[2] * self.norm_vals[looper, (norm_looper + 1) %
                                                n_class])
                self.norm_vals[looper, norm_looper] = n_val
        self.patch_hs = patch_hs
        self.patch_ws = patch_ws
        self.patch_Is_resized = patch_Is_resized


from typing import Any, Dict, List, Tuple, Union, Optional, Sequence
from typing_extensions import Literal

import os
from enum import auto
from math import fabs
from pathlib import Path
from collections import Iterable, defaultdict

from anndata import AnnData
from cellrank import logging as logg
from cellrank._key import Key
from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum, Backend_t
from cellrank.ul._docs import d, inject_docs
from cellrank.pl._utils import (
    _fit_bulk,
    _get_backend,
    _callback_type,
    _create_models,
    _time_range_type,
    _create_callbacks,
    _input_model_type,
    _return_model_type,
    _get_categorical_colors,
)
from cellrank.tl._utils import save_fig, _min_max_scale, _unique_order_preserving
from cellrank.ul._utils import (
    _genesymbols,
    _get_n_cores,
    valuedispatch,
    _check_collection,
)

import numpy as np
import pandas as pd
from scipy.ndimage.filters import convolve

import matplotlib as mpl
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
from seaborn import clustermap
from matplotlib import cm
from matplotlib.ticker import FormatStrFormatter
from mpl_toolkits.axes_grid1.inset_locator import inset_axes

_N_XTICKS = 10

class HeatmapMode(ModeEnum):  # noqa: D101
    GENES = auto()
    LINEAGES = auto()

@d.dedent
@inject_docs(m=HeatmapMode)
@_genesymbols
def heatmap(
    adata: AnnData,
    model: _input_model_type,
    genes: Sequence[str],
    lineages: Optional[Union[str, Sequence[str]]] = None,
    backward: bool = False,
    mode: Literal["genes", "lineages"] = HeatmapMode.LINEAGES,
    time_key: str = "latent_time",
    time_range: Optional[Union[_time_range_type, List[_time_range_type]]] = None,
    callback: _callback_type = None,
    cluster_key: Optional[Union[str, Sequence[str]]] = None,
    show_absorption_probabilities: bool = False,
    cluster_genes: bool = False,
    keep_gene_order: bool = False,
    scale: bool = True,
    n_convolve: Optional[int] = 5,
    show_all_genes: bool = False,
    cbar: bool = True,
    lineage_height: float = 0.33,
    fontsize: Optional[float] = None,
    xlabel: Optional[str] = None,
    cmap: mcolors.ListedColormap = cm.viridis,
    dendrogram: bool = True,
    return_genes: bool = False,
    return_models: bool = False,
    n_jobs: Optional[int] = 1,
    backend: Backend_t = _DEFAULT_BACKEND,
    show_progress_bar: bool = True,
    figsize: Optional[Tuple[float, float]] = None,
    dpi: Optional[int] = None,
    save: Optional[Union[str, Path]] = None,
    **kwargs: Any,
) -> Optional[
    Union[Dict[str, pd.DataFrame], Tuple[_return_model_type, Dict[str, pd.DataFrame]]]
]:
    """
    Plot a heatmap of smoothed gene expression along specified lineages.

    Parameters
    ----------
    %(adata)s
    %(model)s
    %(genes)s
    lineages
        Names of the lineages for which to plot. If `None`, plot all lineages.
    %(backward)s
    mode
        Valid options are:

            - `{m.LINEAGES!r}` - group by ``genes`` for each lineage in ``lineages``.
            - `{m.GENES!r}` - group by ``lineages`` for each gene in ``genes``.
    time_key
        Key in attr:`anndata.AnnData.obs` where the pseudotime is stored.
    %(time_range)s

        This can also be specified on per-lineage basis.
    %(gene_symbols)s
    %(model_callback)s
    cluster_key
        Key(s) in :attr:`anndata.AnnData.obs` containing categorical observations to be plotted on the top of heatmap.
        Only available when ``mode = {m.LINEAGES!r}``.
    show_absorption_probabilities
        Whether to also plot absorption probabilities alongside the smoothed expression.
        Only available when ``mode = {m.LINEAGES!r}``.
    cluster_genes
        Whether to cluster genes using :func:`seaborn.clustermap` when ``mode = 'lineages'``.
    keep_gene_order
        Whether to keep the gene order for later lineages after the first was sorted.
        Only available when ``cluster_genes = False`` and ``mode = {m.LINEAGES!r}``.
    scale
        Whether to normalize the gene expression `[0, 1]` range.
    n_convolve
        Size of the convolution window when smoothing absorption probabilities.
    show_all_genes
        Whether to show all genes on y-axis.
    cbar
        Whether to show the colorbar.
    lineage_height
        Height of a bar when ``mode = {m.GENES!r}``.
    fontsize
        Size of the title's font.
    xlabel
        Label on the x-axis. If `None`, it is determined based on ``time_key``.
    cmap
        Colormap to use when visualizing the smoothed expression.
    dendrogram
        Whether to show dendrogram when ``cluster_genes = True``.
    return_genes
        Whether to return the sorted or clustered genes. Only available when ``mode = {m.LINEAGES!r}``.
    %(return_models)s
    %(parallel)s
    %(plotting)s
    kwargs
        Keyword arguments for :meth:`cellrank.ul.models.BaseModel.prepare`.

    Returns
    -------
    %(plots_or_returns_models)s

    If ``return_genes = True`` and ``mode = {m.LINEAGES!r}``, returns :class:`pandas.DataFrame`
    containing the clustered or sorted genes.
    """

    def find_indices(series: pd.Series, values) -> Tuple[Any]:
        def find_nearest(array: np.ndarray, value: float) -> int:
            ix = np.searchsorted(array, value, side="left")
            if ix > 0 and (
                ix == len(array)
                or fabs(value - array[ix - 1]) < fabs(value - array[ix])
            ):
                return int(ix - 1)
            return int(ix)

        series = series[np.argsort(series.values)]

        return tuple(series[[find_nearest(series.values, v) for v in values]].index)

    def subset_lineage(lname: str, rng: np.ndarray) -> np.ndarray:
        time_series = adata.obs[time_key]
        ixs = find_indices(time_series, rng)

        lin = adata[ixs, :].obsm[lineage_key][lname]

        lin = lin.X.copy().squeeze()
        if n_convolve is not None:
            lin = convolve(lin, np.ones(n_convolve) / n_convolve, mode="nearest")

        return lin

    def create_col_colors(
        lname: str, rng: np.ndarray
    ) -> Tuple[np.ndarray, mcolors.Colormap, mcolors.Normalize]:
        color = adata.obsm[lineage_key][lname].colors[0]
        lin = subset_lineage(lname, rng)

        h, _, v = mcolors.rgb_to_hsv(mcolors.to_rgb(color))
        end_color = mcolors.hsv_to_rgb([h, 1, v])

        lineage_cmap = mcolors.LinearSegmentedColormap.from_list(
            "lineage_cmap", ["#ffffff", end_color], N=len(rng)
        )
        norm = mcolors.Normalize(vmin=np.min(lin), vmax=np.max(lin))
        scalar_map = cm.ScalarMappable(cmap=lineage_cmap, norm=norm)

        return (
            np.array([mcolors.to_hex(c) for c in scalar_map.to_rgba(lin)]),
            lineage_cmap,
            norm,
        )

    def create_col_categorical_color(cluster_key: str, rng: np.ndarray) -> np.ndarray:
        colors, mapper = _get_categorical_colors(adata, cluster_key)
        ixs = find_indices(adata.obs[time_key], rng)

        return np.array(
            [mcolors.to_hex(mapper[v]) for v in adata[ixs, :].obs[cluster_key].values]
        )

    def create_cbar(
        ax,
        x_delta: float,
        cmap: mcolors.Colormap,
        norm: mcolors.Normalize,
        label: Optional[str] = None,
    ) -> plt.Axes:
        cax = inset_axes(
            ax,
            width="1%",
            height="100%",
            loc="lower right",
            bbox_to_anchor=(x_delta, 0, 1, 1),
            bbox_transform=ax.transAxes,
        )

        _ = mpl.colorbar.ColorbarBase(
            cax,
            cmap=cmap,
            norm=norm,
            label=label,
            ticks=np.linspace(norm.vmin, norm.vmax, 5),
        )

        return cax

    @valuedispatch
    def _plot_heatmap(mode: HeatmapMode) -> plt.Figure:
        raise NotImplementedError(mode.value)

    @_plot_heatmap.register(HeatmapMode.GENES)
    def _() -> Tuple[plt.Figure, None]:
        def color_fill_rec(ax, xs, y1, y2, colors=None, cmap=cmap, **kwargs) -> None:
            colors = colors if cmap is None else cmap(colors)

            x = 0
            for i, (color, x, y1, y2) in enumerate(zip(colors, xs, y1, y2)):
                dx = (xs[i + 1] - xs[i]) if i < len(x) else (xs[-1] - xs[-2])
                ax.add_patch(
                    plt.Rectangle((x, y1), dx, y2 - y1, color=color, ec=color, **kwargs)
                )

            ax.plot(x, y2, lw=0)

        fig, axes = plt.subplots(
            nrows=len(genes) + show_absorption_probabilities,
            figsize=(12, len(genes) + len(lineages) * lineage_height)
            if figsize is None
            else figsize,
            dpi=dpi,
            constrained_layout=True,
        )

        if not isinstance(axes, Iterable):
            axes = [axes]
        axes = np.ravel(axes)

        if show_absorption_probabilities:
            data["absorption probability"] = data[next(iter(data.keys()))]

        for ax, (gene, models) in zip(axes, data.items()):
            if scale:
                vmin, vmax = 0, 1
            else:
                c = np.array([m.y_test for m in models.values()])
                vmin, vmax = np.nanmin(c), np.nanmax(c)

            norm = mcolors.Normalize(vmin=vmin, vmax=vmax)

            ix = 0
            ys = [ix]

            if gene == "absorption probability":
                norm = mcolors.Normalize(vmin=0, vmax=1)
                for ln, x in ((ln, m.x_test) for ln, m in models.items()):
                    y = np.ones_like(x)
                    c = subset_lineage(ln, x.squeeze())

                    color_fill_rec(
                        ax, x, y * ix, y * (ix + lineage_height), colors=norm(c)
                    )

                    ix += lineage_height
                    ys.append(ix)
            else:
                for x, c in ((m.x_test, m.y_test) for m in models.values()):
                    y = np.ones_like(x)
                    c = _min_max_scale(c) if scale else c

                    color_fill_rec(
                        ax, x, y * ix, y * (ix + lineage_height), colors=norm(c)
                    )

                    ix += lineage_height
                    ys.append(ix)

            xs = np.array([m.x_test for m in models.values()])
            x_min, x_max = np.min(xs), np.max(xs)
            ax.set_xticks(np.linspace(x_min, x_max, _N_XTICKS))

            ax.set_yticks(np.array(ys[:-1]) + lineage_height / 2)
            ax.spines["left"].set_position(
                ("data", 0)
            )  # move the left spine to the rectangles to get nicer yticks
            ax.set_yticklabels(models.keys(), ha="right")

            ax.set_title(gene, fontdict={"fontsize": fontsize})
            ax.set_ylabel("lineage")

            for pos in ["top", "bottom", "left", "right"]:
                ax.spines[pos].set_visible(False)

            if cbar:
                cax, _ = mpl.colorbar.make_axes(ax)
                _ = mpl.colorbar.ColorbarBase(
                    cax,
                    ticks=np.linspace(vmin, vmax, 5),
                    norm=norm,
                    cmap=cmap,
                    label="value"
                    if gene == "absorption probability"
                    else "scaled expression"
                    if scale
                    else "expression",
                )

            ax.tick_params(
                top=False,
                bottom=False,
                left=True,
                right=False,
                labelleft=True,
                labelbottom=False,
            )

        ax.xaxis.set_major_formatter(FormatStrFormatter("%.3f"))
        ax.tick_params(
            top=False,
            bottom=True,
            left=True,
            right=False,
            labelleft=True,
            labelbottom=True,
        )
        ax.set_xlabel(xlabel)

        return fig, None

    @_plot_heatmap.register(HeatmapMode.LINEAGES)
    def _() -> Tuple[List[plt.Figure], pd.DataFrame]:
        data_t = defaultdict(dict)  # transpose
        for gene, lns in data.items():
            for ln, y in lns.items():
                data_t[ln][gene] = y

        figs = []
        gene_order = None
        sorted_genes = pd.DataFrame() if return_genes else None

        for lname, models in data_t.items():
            xs = np.array([m.x_test for m in models.values()])
            x_min, x_max = np.nanmin(xs), np.nanmax(xs)

            df = pd.DataFrame([m.y_test for m in models.values()], index=models.keys())
            df.index.name = "genes"

            if not cluster_genes:
                if gene_order is not None:
                    df = df.loc[gene_order]
                else:
                    max_sort = np.argsort(
                        np.argmax(df.apply(_min_max_scale, axis=1).values, axis=1)
                    )
                    df = df.iloc[max_sort, :]
                    if keep_gene_order:
                        gene_order = df.index

            cat_colors = None
            if cluster_key is not None:
                cat_colors = np.stack(
                    [
                        create_col_categorical_color(
                            c, np.linspace(x_min, x_max, df.shape[1])
                        )
                        for c in cluster_key
                    ],
                    axis=0,
                )

            if show_absorption_probabilities:
                col_colors, col_cmap, col_norm = create_col_colors(
                    lname, np.linspace(x_min, x_max, df.shape[1])
                )
                if cat_colors is not None:
                    col_colors = np.vstack([cat_colors, col_colors[None, :]])
            else:
                col_colors, col_cmap, col_norm = cat_colors, None, None

            row_cluster = cluster_genes and df.shape[0] > 1
            show_clust = row_cluster and dendrogram

            g = clustermap(
                data=df,
                cmap=cmap,
                figsize=(10, min(len(genes) / 8 + 1, 10))
                if figsize is None
                else figsize,
                xticklabels=False,
                row_cluster=row_cluster,
                col_colors=col_colors,
                colors_ratio=0,
                col_cluster=False,
                cbar_pos=None,
                yticklabels=show_all_genes or "auto",
                standard_scale=0 if scale else None,
            )

            if cbar:
                cax = create_cbar(
                    g.ax_heatmap,
                    0.1,
                    cmap=cmap,
                    norm=mcolors.Normalize(
                        vmin=0 if scale else np.min(df.values),
                        vmax=1 if scale else np.max(df.values),
                    ),
                    label="scaled expression" if scale else "expression",
                )
                g.fig.add_axes(cax)

                if col_cmap is not None and col_norm is not None:
                    cax = create_cbar(
                        g.ax_heatmap,
                        0.25,
                        cmap=col_cmap,
                        norm=col_norm,
                        label="absorption probability",
                    )
                    g.fig.add_axes(cax)

            if g.ax_col_colors:
                main_bbox = _get_ax_bbox(g.fig, g.ax_heatmap)
                n_bars = show_absorption_probabilities + (
                    len(cluster_key) if cluster_key is not None else 0
                )
                _set_ax_height_to_cm(
                    g.fig,
                    g.ax_col_colors,
                    height=min(
                        5, max(n_bars * main_bbox.height / len(df), 0.25 * n_bars)
                    ),
                )
                g.ax_col_colors.set_title(lname, fontdict={"fontsize": fontsize})
            else:
                g.ax_heatmap.set_title(lname, fontdict={"fontsize": fontsize})

            g.ax_col_dendrogram.set_visible(False)  # gets rid of top free space

            g.ax_heatmap.yaxis.tick_left()
            g.ax_heatmap.yaxis.set_label_position("right")

            g.ax_heatmap.set_xlabel(xlabel)
            g.ax_heatmap.set_xticks(np.linspace(0, len(df.columns), _N_XTICKS))
            g.ax_heatmap.set_xticklabels(
                list(map(lambda n: round(n, 3), np.linspace(x_min, x_max, _N_XTICKS)))
            )

            if show_clust:
                # robustly show dendrogram, because gene names can be long
                g.ax_row_dendrogram.set_visible(True)
                dendro_box = g.ax_row_dendrogram.get_position()

                pad = 0.005
                bb = g.ax_heatmap.yaxis.get_tightbbox(
                    g.fig.canvas.get_renderer()
                ).transformed(g.fig.transFigure.inverted())

                dendro_box.x0 = bb.x0 - dendro_box.width - pad
                dendro_box.x1 = bb.x0 - pad

                g.ax_row_dendrogram.set_position(dendro_box)
            else:
                g.ax_row_dendrogram.set_visible(False)

            if return_genes:
                sorted_genes[lname] = (
                    df.index[g.dendrogram_row.reordered_ind]
                    if hasattr(g, "dendrogram_row") and g.dendrogram_row is not None
                    else df.index
                )

            figs.append(g)

        return figs, sorted_genes

    mode = HeatmapMode(mode)

    lineage_key = Key.obsm.abs_probs(backward)
    if lineage_key not in adata.obsm:
        raise KeyError(f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.")

    if lineages is None:
        lineages = adata.obsm[lineage_key].names
    elif isinstance(lineages, str):
        lineages = [lineages]
    lineages = _unique_order_preserving(lineages)

    _ = adata.obsm[lineage_key][lineages]

    if cluster_key is not None:
        if isinstance(cluster_key, str):
            cluster_key = [cluster_key]
        cluster_key = _unique_order_preserving(cluster_key)

    if isinstance(genes, str):
        genes = [genes]
    genes = _unique_order_preserving(genes)
    _check_collection(adata, genes, "var_names", use_raw=kwargs.get("use_raw", False))

    kwargs["backward"] = backward
    kwargs["time_key"] = time_key
    models = _create_models(model, genes, lineages)
    all_models, data, genes, lineages = _fit_bulk(
        models,
        _create_callbacks(adata, callback, genes, lineages, **kwargs),
        genes,
        lineages,
        time_range,
        return_models=True,  # always return (better error messages)
        filter_all_failed=True,
        parallel_kwargs={
            "show_progress_bar": show_progress_bar,
            "n_jobs": _get_n_cores(n_jobs, len(genes)),
            "backend": _get_backend(models, backend),
        },
        **kwargs,
    )

    xlabel = time_key if xlabel is None else xlabel

    logg.debug(f"Plotting `{mode!r}` heatmap")
    fig, genes = _plot_heatmap(mode)

    if save is not None and fig is not None:
        if not isinstance(fig, Iterable):
            save_fig(fig, save)
        elif len(fig) == 1:
            save_fig(fig[0], save)
        else:
            for ln, f in zip(lineages, fig):
                save_fig(f, os.path.join(save, f"lineage_{ln}"))

    if return_genes and mode == HeatmapMode.LINEAGES:
        return (all_models, genes) if return_models else genes
    elif return_models:
        return all_models

def _get_ax_bbox(fig: plt.Figure, ax: plt.Axes):
    return ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())

def _set_ax_height_to_cm(fig: plt.Figure, ax: plt.Axes, height: float) -> None:
    from mpl_toolkits.axes_grid1 import Size, Divider

    height /= 2.54  # cm to inches

    bbox = _get_ax_bbox(fig, ax)

    hori = [Size.Fixed(bbox.x0), Size.Fixed(bbox.width), Size.Fixed(bbox.x1)]
    vert = [Size.Fixed(bbox.y0), Size.Fixed(height), Size.Fixed(bbox.y1)]

    divider = Divider(fig, (0.0, 0.0, 1.0, 1.0), hori, vert, aspect=False)

    ax.set_axes_locator(divider.new_locator(nx=1, ny=1))


#!/usr/bin/env python
# coding: utf-8
import os
import sys
import random

caffe_path = '../caffe/python/'
sys.path.insert(0, caffe_path)
import caffe
caffe.set_device(0)
caffe.set_mode_gpu()
from caffe.proto import caffe_pb2
from google.protobuf import text_format
import numpy as np
#import cv2
#import matplotlib.pyplot as plt
import time
from mq_write import dh_infer_label
import oss2
import infer

'''
init oss account, datahub account, caffe prototxt, caffemodel, labelmaps
'''
auth = oss2.Auth('', '')
bucket = oss2.Bucket(auth, 'http://oss-cn-beijing.aliyuncs.com', '')
#################################################################################

root_googlenet = '../model/'
#deploy_googlenet = root_googlenet + 'deploy.prototxt'
labels_filename = root_googlenet + 'labels.txt'

deploy_googlenet = root_googlenet + 'deploy-googlenet.prototxt'
#labels_filename = root_googlenet + 'labels.txt'
caffe_model_googlenet = root_googlenet + 'googlenet.caffemodel'
googlenet = caffe.Net(deploy_googlenet, caffe_model_googlenet, caffe.TEST)
#################################################################################

def infer_label(path, th = 0.5):
    object_id = path.split('/')[-1]
    object_id = object_id.split('.')[0]
    url = path
    labels_googlenet, score_googlenet, _ = infer.infer_img(googlenet, url)
#    if labels_googlenet == 'normal' and score_googlenet < 0.89:
#        labels_googlenet = 'unk'
#        score_googlenet = 0.66666666
#    else:
#        pass
    dh_infer_label(object_id, labels_googlenet, score_googlenet)
    return object_id, labels_googlenet, score_googlenet

if __name__ == "__main__":
    path = '../model/a.jpg'
    object_id, ai_label, score = infer_label(path)
    print(object_id)
    print(ai_label, score)


from argparse import ArgumentParser
from glob import glob
import os

import numpy as np
from sklearn.utils import shuffle
import matplotlib.pyplot as plt
import pandas as pd
from tqdm import tqdm
from PIL import Image
import matplotlib.pyplot as plt

def load_image(filename):
    try:
        with open(filename, "rb") as f:
            image = Image.open(f)
            return image.convert("RGB")
    except UserWarning as e:
        print(filename)
        input("Something wrong happens while loading image: {} {}".format(filename, str(e)))

# Example Model definition
class Model(object):
    def __init__(self, dirname):
        import animecv

        self.encoder = animecv.general.create_OML_ImageFolder_Encoder(dirname)
        self.encoder.to("cuda")

    # img: PIL image
    def encode(self, img):
        vecs = self.encoder.encode([img]).detach().cpu().numpy()
        return vecs[0]

if __name__=="__main__":
    parser = ArgumentParser()

    parser.add_argument("--test-pairs", help="CSV file which lists test image pairs.")
    parser.add_argument("--test-dataset-dir", help="Directory of test images.")
    parser.add_argument("--ignore-list", default=None, help="List of images which should be ignored during pair sampling.")

    parser.add_argument("--out-fn", default="adversarial.csv")

    parser.add_argument("--n-negative", type=int, default=3000)

    args = parser.parse_args()

    if not os.path.exists(args.out_fn):
        if args.ignore_list is not None:
            df = pd.read_csv(args.ignore_list, header=None)
            ignore_list = set(df.values.flatten().tolist())
        else:
            ignore_list = set()

        # Generate adversarial negative pairs.
        model = Model("0206_resnet152")

        images = glob(os.path.join(args.test_dataset_dir, "**"), recursive=True)
        images = [fn for fn in images if os.path.isfile(fn)]
        labels = [fn.split(os.path.sep)[-2] for fn in images]

        vecs = []
        for fn in tqdm(images):
            img = load_image(fn)
            vecs.append(model.encode(img).reshape((1,-1)))
        vecs = np.concatenate(vecs, axis=0)

        scores = np.sum(vecs[:,np.newaxis,:] * vecs[np.newaxis,:,:], axis=2)

        negative_pairs = []
        n_img = scores.shape[0]
        sorted_idx = np.argsort(-scores, axis=None).tolist()
        strip_len = len(args.test_dataset_dir + os.path.sep)
        while len(negative_pairs) < args.n_negative:
            idx = sorted_idx.pop(0)
            i,j = idx // n_img, idx % n_img
            if i<=j:
                continue
            if labels[i] == labels[j]:
                continue
            if os.path.basename(images[i]) in ignore_list:
                continue
            if os.path.basename(images[j]) in ignore_list:
                continue
            negative_pairs.append((images[i][strip_len:], images[j][strip_len:], 0, -1, 0))

        # Reuse positive pairs.
        positive_pairs = []
        df = pd.read_csv(args.test_pairs)
        for pathA, pathB in df[df["label"]==1][["pathA", "pathB"]].values:
            #print(pathA, pathB)
            positive_pairs.append((pathA, pathB, 1, -1, 0))

        pairs = shuffle(positive_pairs + negative_pairs)

        df = pd.DataFrame(pairs, columns=["pathA", "pathB", "label", "human_prediction", "invalid"])
        df.to_csv(args.out_fn, index=False)
    else:
        print("Reload")
        df = pd.read_csv(args.out_fn)

    for i_row in tqdm(list(range(df.values.shape[0]))):
        pathA, pathB, label, pred, invalid = df.loc[i_row].values
        #print(pathA, pathB)
        if pred >= 0:
            continue
        else:
            im1 = np.array(Image.open(os.path.join(args.test_dataset_dir, pathA)))
            im2 = np.array(Image.open(os.path.join(args.test_dataset_dir, pathB)))
            ax = plt.subplot(1,2,1)
            ax.imshow(im1)
            ax = plt.subplot(1,2,2)
            ax.imshow(im2)
            plt.draw()
            plt.pause(0.001)
            cmd = input("correct?[y/n]: ")
            if cmd=="y":
                pred = 1
            elif cmd=="n":
                pred = 0
            else:
                pred = 0
                df.loc[i_row, "invalid"] = 1
            df.loc[i_row, "human_prediction"] = pred
            df.to_csv(args.out_fn, index=False)
            plt.close()


import os
import SimpleITK as sitk
import argparse
import pandas as pd
import numpy as np
from tqdm import tqdm

from dpipe.dataset.wrappers import apply
from ssn.cc359 import CC359, Rescale3D, scale_mri

CC359_DATA_PATH = '/shared/data/cc359'

train_ids = [
"CC0184",
"CC0185",
"CC0188",
"CC0193",
"CC0194",
"CC0195",
"CC0196",
"CC0197",
"CC0201",
"CC0202",
"CC0203",
"CC0207",
"CC0209",
"CC0213",
"CC0214",
"CC0215",
"CC0219",
"CC0222",
"CC0223",
"CC0224",
"CC0225",
"CC0227",
"CC0229",
"CC0230",
"CC0231",
"CC0233",
"CC0234",
"CC0236",
"CC0237",
"CC0238"
]

valid_ids = ["CC0359"]

test_ids = [
"CC0180",
"CC0181",
"CC0182",
"CC0183",
"CC0186",
"CC0187",
"CC0189",
"CC0190",
"CC0191",
"CC0192",
"CC0198",
"CC0199",
"CC0200",
"CC0204",
"CC0205",
"CC0206",
"CC0208",
"CC0210",
"CC0211",
"CC0212",
"CC0216",
"CC0217",
"CC0218",
"CC0220",
"CC0221",
"CC0226",
"CC0228",
"CC0232",
"CC0235",
"CC0239",
"CC0001",
"CC0002",
"CC0003",
"CC0004",
"CC0005",
"CC0006",
"CC0007",
"CC0008",
"CC0009",
"CC0010",
"CC0011",
"CC0012",
"CC0013",
"CC0014",
"CC0015",
"CC0016",
"CC0017",
"CC0018",
"CC0019",
"CC0020",
"CC0021",
"CC0022",
"CC0023",
"CC0024",
"CC0025",
"CC0026",
"CC0027",
"CC0028",
"CC0029",
"CC0030",
"CC0031",
"CC0032",
"CC0033",
"CC0034",
"CC0035",
"CC0036",
"CC0037",
"CC0038",
"CC0039",
"CC0040",
"CC0041",
"CC0042",
"CC0043",
"CC0044",
"CC0045",
"CC0046",
"CC0047",
"CC0048",
"CC0049",
"CC0050",
"CC0051",
"CC0052",
"CC0053",
"CC0054",
"CC0055",
"CC0056",
"CC0057",
"CC0058",
"CC0059",
"CC0060",
"CC0061",
"CC0062",
"CC0063",
"CC0064",
"CC0065",
"CC0066",
"CC0067",
"CC0068",
"CC0069",
"CC0070",
"CC0071",
"CC0072",
"CC0073",
"CC0074",
"CC0075",
"CC0076",
"CC0077",
"CC0078",
"CC0079",
"CC0080",
"CC0081",
"CC0082",
"CC0083",
"CC0084",
"CC0085",
"CC0086",
"CC0087",
"CC0088",
"CC0089",
"CC0090",
"CC0091",
"CC0092",
"CC0093",
"CC0094",
"CC0095",
"CC0096",
"CC0097",
"CC0098",
"CC0099",
"CC0100",
"CC0101",
"CC0102",
"CC0103",
"CC0104",
"CC0105",
"CC0106",
"CC0107",
"CC0108",
"CC0109",
"CC0110",
"CC0111",
"CC0112",
"CC0113",
"CC0114",
"CC0115",
"CC0116",
"CC0117",
"CC0118",
"CC0119",
"CC0120",
"CC0121",
"CC0122",
"CC0123",
"CC0124",
"CC0125",
"CC0126",
"CC0127",
"CC0128",
"CC0129",
"CC0130",
"CC0131",
"CC0132",
"CC0133",
"CC0134",
"CC0135",
"CC0136",
"CC0137",
"CC0138",
"CC0139",
"CC0140",
"CC0141",
"CC0142",
"CC0143",
"CC0144",
"CC0145",
"CC0146",
"CC0147",
"CC0148",
"CC0149",
"CC0150",
"CC0151",
"CC0152",
"CC0153",
"CC0154",
"CC0155",
"CC0156",
"CC0157",
"CC0158",
"CC0159",
"CC0160",
"CC0161",
"CC0162",
"CC0163",
"CC0164",
"CC0165",
"CC0166",
"CC0167",
"CC0168",
"CC0169",
"CC0170",
"CC0171",
"CC0172",
"CC0173",
"CC0174",
"CC0175",
"CC0176",
"CC0177",
"CC0178",
"CC0179",
"CC0240",
"CC0241",
"CC0242",
"CC0243",
"CC0244",
"CC0245",
"CC0246",
"CC0247",
"CC0248",
"CC0249",
"CC0250",
"CC0251",
"CC0252",
"CC0253",
"CC0254",
"CC0255",
"CC0256",
"CC0257",
"CC0258",
"CC0259",
"CC0260",
"CC0261",
"CC0262",
"CC0263",
"CC0264",
"CC0265",
"CC0266",
"CC0267",
"CC0268",
"CC0269",
"CC0270",
"CC0271",
"CC0272",
"CC0273",
"CC0274",
"CC0275",
"CC0276",
"CC0277",
"CC0278",
"CC0279",
"CC0280",
"CC0281",
"CC0282",
"CC0283",
"CC0284",
"CC0285",
"CC0286",
"CC0287",
"CC0288",
"CC0289",
"CC0290",
"CC0291",
"CC0292",
"CC0293",
"CC0294",
"CC0295",
"CC0296",
"CC0297",
"CC0298",
"CC0299",
"CC0300",
"CC0301",
"CC0302",
"CC0303",
"CC0304",
"CC0305",
"CC0306",
"CC0307",
"CC0308",
"CC0309",
"CC0310",
"CC0311",
"CC0312",
"CC0313",
"CC0314",
"CC0315",
"CC0316",
"CC0317",
"CC0318",
"CC0319",
"CC0320",
"CC0321",
"CC0322",
"CC0323",
"CC0324",
"CC0325",
"CC0326",
"CC0327",
"CC0328",
"CC0329",
"CC0330",
"CC0331",
"CC0332",
"CC0333",
"CC0334",
"CC0335",
"CC0336",
"CC0337",
"CC0338",
"CC0339",
"CC0340",
"CC0341",
"CC0342",
"CC0343",
"CC0344",
"CC0345",
"CC0346",
"CC0347",
"CC0348",
"CC0349",
"CC0350",
"CC0351",
"CC0352",
"CC0353",
"CC0354",
"CC0355",
"CC0356",
"CC0357",
"CC0358",
"CC0359"
]

def get_brain_mask(t1):
    brain_mask = sitk.GetImageFromArray((sitk.GetArrayFromImage(t1) > 0).astype(np.uint8))
    brain_mask.CopyInformation(t1)
    brain_mask = sitk.Cast(brain_mask, sitk.sitkUInt8)
    return brain_mask

def z_score_normalisation(channel, brain_mask, cutoff_percentiles=(5., 95.), cutoff_below_mean=True):
    low, high = np.percentile(channel[brain_mask.astype(bool)], cutoff_percentiles)
    norm_mask = np.logical_and(brain_mask, np.logical_and(channel > low, channel < high))
    if cutoff_below_mean:
        norm_mask = np.logical_and(norm_mask, channel > np.mean(channel))
    masked_channel = channel[norm_mask]
    normalised_channel = (channel - np.mean(masked_channel)) / np.std(masked_channel)
    return normalised_channel

def fix_segmentation_labels(seg):
    array = sitk.GetArrayFromImage(seg)
    array[array == 4] = 3
    new_seg = sitk.GetImageFromArray(array)
    new_seg.CopyInformation(seg)
    return new_seg

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--input-dir',
                        required=True,
                        type=str,
                        help='Path to input directory.')
    parser.add_argument('--output-dir',
                        required=True,
                        type=str,
                        help='Path to output directory.')

    parse_args, unknown = parser.parse_known_args()

    data_path = CC359_DATA_PATH
    voxel_spacing = (1, 0.95, 0.95)

    preprocessed_dataset = apply(Rescale3D(CC359(data_path), voxel_spacing), load_image=scale_mri)
    dataset = apply(preprocessed_dataset, load_image=np.float32)

    output_dataframe = pd.DataFrame()
    for id_ in tqdm(dataset.ids):
        segm = sitk.GetImageFromArray(dataset.load_segm(id_))
        seg_path = os.path.join(parse_args.output_dir, id_) + '_seg.nii.gz'
        output_dataframe.loc[id_, 'seg'] = seg_path
        os.makedirs(os.path.dirname(seg_path), exist_ok=True)
        sitk.WriteImage(segm, seg_path)

        image = sitk.GetImageFromArray(dataset.load_image(id_))
        brain_mask = get_brain_mask(image)
        output_path = os.path.join(parse_args.output_dir, id_) + f'_brain_mask.nii.gz'
        output_dataframe.loc[id_, 'sampling_mask'] = output_path
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        sitk.WriteImage(brain_mask, output_path)

        channel_array = sitk.GetArrayFromImage(image)
        normalised_channel_array = z_score_normalisation(channel_array, sitk.GetArrayFromImage(brain_mask))
        normalised_channel = sitk.GetImageFromArray(normalised_channel_array)
        normalised_channel.CopyInformation(image)
        output_path = os.path.join(parse_args.output_dir, id_) + f'_t1.nii.gz'
        output_dataframe.loc[id_, 't1'] = output_path
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        sitk.WriteImage(normalised_channel, output_path)

    output_dataframe.index.name = 'id'
    os.makedirs('stochastic_segmentation_networks/assets/cc359_data', exist_ok=True)
    train_index = output_dataframe.loc[train_ids]
    train_index.to_csv('stochastic_segmentation_networks/assets/cc359_data/data_index_train.csv')
    valid_index = output_dataframe.loc[valid_ids]
    valid_index.to_csv('stochastic_segmentation_networks/assets/cc359_data/data_index_valid.csv')
    test_index = output_dataframe.loc[test_ids]
    test_index.to_csv('stochastic_segmentation_networks/assets/cc359_data/data_index_test.csv')


import numpy as np
import OpenGL.GL as GL
import traitlets
from yt.data_objects.data_containers import YTDataContainer
from yt.data_objects.index_subobjects.octree_subset import OctreeSubsetBlockSlice

from yt_idv.constants import aabb_triangle_strip
from yt_idv.opengl_support import Texture3D, VertexArray, VertexAttribute
from yt_idv.scene_data.base_data import SceneData

class OctreeBlockCollection(SceneData):
    name = "octree_block_collection"
    data_source = traitlets.Instance(YTDataContainer)
    data_textures = traitlets.List(value_trait=traitlets.Instance(Texture3D))
    bitmap_textures = traitlets.List(value_trait=traitlets.Instance(Texture3D))
    shapes = traitlets.List(value_trait=traitlets.CInt())

    @traitlets.default("vertex_array")
    def _default_vertex_array(self):
        return VertexArray(name="octree_block_info", each=14)

    def add_data(self, field):
        r"""Adds a source of data for the block collection.

        Given a `data_source` and a `field` to populate from, adds the data
        to the block collection so that is able to be rendered.

        Parameters
        ----------
        data_source : YTRegion
            A YTRegion object to use as a data source.
        field : string
            A field to populate from.
        no_ghost : bool (False)
            Should we speed things up by skipping ghost zone generation?
        """
        ds = self.data_source.ds
        ds.index._identify_base_chunk(self.data_source)
        left_edges = []
        right_edges = []
        dx = []
        data = []

        for obj in self.data_source._current_chunk.objs:
            bs = OctreeSubsetBlockSlice(obj, ds)
            LE = bs._fcoords[0, 0, 0, :, :].d - bs._fwidth[0, 0, 0, :, :].d * 0.5
            RE = bs._fcoords[-1, -1, -1, :, :].d + bs._fwidth[-1, -1, -1, :, :].d * 0.5
            dx.append(bs._fwidth[-1, -1, -1, :, :].d)
            left_edges.append(LE)
            right_edges.append(RE)
            d = bs.get_vertex_centered_data([field])[field]
            data.append(
                np.concatenate(
                    [d[:, :, :, i] for i in range(d.shape[-1])], axis=2
                ).copy(order="C")
            )

        # Let's reshape ...
        left_edges = np.concatenate(left_edges, axis=0).astype("f4")
        right_edges = np.concatenate(right_edges, axis=0).astype("f4")
        dx = np.concatenate(dx, axis=0).astype("f4")
        data = np.concatenate(data, axis=-1).astype("f4")
        data = data.reshape((3, 3, -1))

        self.min_val = np.nanmin(data)
        self.max_val = np.nanmax(data)

        if hasattr(self.min_val, "in_units"):
            self.min_val = self.min_val.d
        if hasattr(self.max_val, "in_units"):
            self.max_val = self.max_val.d

        if self.max_val != self.min_val:
            data = (data - self.min_val) / (
                (self.max_val - self.min_val)
            )  # * self.diagonal)

        self.vertex_array.attributes.append(
            VertexAttribute(name="model_vertex", data=aabb_triangle_strip, divisor=0)
        )
        self.vertex_array.attributes.append(
            VertexAttribute(name="in_dx", data=dx, divisor=1)
        )
        self.vertex_array.attributes.append(
            VertexAttribute(name="in_left_edge", data=left_edges, divisor=1)
        )
        self.vertex_array.attributes.append(
            VertexAttribute(name="in_right_edge", data=right_edges, divisor=1)
        )

        # Now we set up our textures; we need to break our texture up into
        # groups of MAX_3D_TEXTURE_SIZE
        tex_size = GL.glGetInteger(GL.GL_MAX_3D_TEXTURE_SIZE)
        # for now, use one texture for all the bitmaps
        bt = Texture3D(data=np.ones((3, 3, tex_size), dtype="u1") * 255)
        for start_index in np.mgrid[0 : data.shape[-1] : tex_size]:
            d = data[:, :, start_index : start_index + tex_size]
            self.data_textures.append(Texture3D(data=d))
            self.bitmap_textures.append(bt)
            self.shapes.append(d.shape[-1])


import numpy as np
from gym import Env, spaces
from .core.controller import RobotController
from .core.observer import Observer

#Camera settings
IMAGE_WIDTH = 320
IMAGE_HEIGHT = 240
IMAGE_SIZE = (IMAGE_WIDTH, IMAGE_HEIGHT)

#Actuator settings
MIN_STEERING = 1.0
MAX_STEERING = -1.0
MIN_THROTTLE = 0.0
MAX_THROTTLE = 1.0

class JetbotEnv(Env):

    def __init__(self):
        super(JetbotEnv, self).__init__()
        self.controller = RobotController()
        self.observer = Observer(IMAGE_WIDTH, IMAGE_HEIGHT)
        self.observation_space = spaces.Box(low=np.finfo(np.float32).min,
                                            high=np.finfo(np.float32).max,
                                            shape=IMAGE_SIZE,
                                            dtype=np.float32)
        self.action_space = spaces.Box(low=np.array([MIN_STEERING, MIN_THROTTLE]),
                                       high=np.array([MAX_STEERING, MAX_THROTTLE]), dtype=np.float32)
        self.ie = {}

        self.observer.start()

    def step(self, action):
        self.controller.action(action[0], action[1])
        obs = self.observer.observation()
        reward = 1.0
        done = False
        return obs, reward, done, self.ie
        pass

    def reset(self):
        self.controller.action(0,0)
        obs = self.observer.observation()
        return obs

    def render(self, mode='human'):
        pass

    def seed(self,seed):
        pass

    def close(self):
        self.observer.stop()
        pass