#!/usr/bin/env python
import logging
import sys

import ROOT
import numpy

# The path to the signal region shapes file.
SIGNAL_SHAPES_PATH = 'vhbb_TH_Znn_13TeV_Signal.root'
# The name of the signal region bin in the signal region shapes file.
SIGNAL_SHAPES_BIN = 'Znn_13TeV_Signal'
# The names of the signal processes.
SIGNAL_SHAPES_PROCESSES = ['ZH_hbb', 'WH_hbb', 'ggZH_hbb']
# The path to the mlfit.root file.
MLFIT_PATH = 'mlfit.root'
# The name of the signal region bin in the mlfit.root file.
MLFIT_BIN = 'Znn_SR'
# The name of the branch containing the nominal BDT score.
BDT_BRANCH = 'BDT_Znn_HighPt.Nominal'

def get_histogram_bin_edges(histogram):
    """Return an array of bin low edges for a histogram.
    """
    bin_edges = numpy.array(
        [histogram.GetXaxis().GetBinLowEdge(i) for i in xrange(1, histogram.GetNbinsX() + 1)],
        dtype=numpy.float64,
    )
    return bin_edges

def get_total_signal_prefit(shapes_path, datacard_bin, signals):
    """Get the total prefit nominal signal shape from the input shapes file.

    Parameters
    ----------
    shapes_path : path
        The path to the shapes file.
    datacard_bin : string
        The name of the datacard bin containing the shapes.
    signals : iterable of strings
        The names of the signal processes within the datacard.

    Returns
    -------
    total_signal_prefit : ROOT.TH1F
        The total prefit nominal signal shape.
    """
    shapes_file = ROOT.TFile.Open(shapes_path)
    shapes_file.cd(datacard_bin)
    signal_shapes = [
        key.ReadObj() for key in ROOT.gDirectory.GetListOfKeys()
        if key.GetName() in signals
    ]
    for i, shape in enumerate(signal_shapes):
        if i == 0:
            total_signal_prefit = shape.Clone('total_signal_prefit')
        else:
            total_signal_prefit.Add(shape)
    total_signal_prefit.SetDirectory(0)
    shapes_file.Close()
    return total_signal_prefit

def get_total_background_postfit(mlfit_path, datacard_bin, bin_edges):
    """Get the rebinned total postfit nominal background shape from an mlfit.root file.

    Parameters
    ----------
    mlfit_path : path
        The path to the mlfit.root file.
    datacard_bin : string
        The name of the datacard bin containing the shapes.
    bin_edges : numpy.array of floats
        An array of bin low edge values used to rebin the postfit shapes.

    Returns
    -------
    total_background_postfit_rebinned : ROOT.TH1F
        The rebinned total postfit nominal background shape.
    """
    mlfit_file = ROOT.TFile.Open(mlfit_path)
    mlfit_file.cd('shapes_fit_s/{}'.format(datacard_bin))
    total_background_postfit = ROOT.gDirectory.Get('total_background')
    total_background_postfit_rebinned = ROOT.TH1F('total_background_postfit', '', len(bin_edges) - 1, bin_edges)
    for i in xrange(1, total_background_postfit_rebinned.GetNbinsX() + 1):
        total_background_postfit_rebinned.SetBinContent(i, total_background_postfit.GetBinContent(i))
    total_background_postfit_rebinned.SetDirectory(0)
    mlfit_file.Close()
    return total_background_postfit_rebinned

def add_sb_weight(src, dst, bdt_branch, total_signal_prefit, total_background_postfit):
    """Add a branch named "sb_weight" which contains the per-event S/(S+B) weight
    for the events' corresponding bin in the signal region BDT score distribution.
    This is to be applied to all MC and data.

    Parameters
    ----------
    src : path
        The path to the input ntuple.
    dst : path
        The path to the output ntuple.
    bdt_branch : string
        The name of the branch containing the nominal BDT score.
    total_signal_prefit : ROOT.TH1F
        The total prefit nominal signal shape.
    total_background_postfit : ROOT.TH1F
        The total postfit nominal background shape.
    """
    logger = logging.getLogger('add_sb_weight')
    # Copy any count and weight histograms.
    infile = ROOT.TFile.Open(src)
    outfile = ROOT.TFile.Open(dst, 'recreate')
    for key in infile.GetListOfKeys():
        if key.GetName() == 'tree':
            continue
        obj = key.ReadObj()
        obj.Write()
    tree = infile.Get('tree')
    # Reset the branch in case it already exists.
    tree.SetBranchStatus('sb_weight', 0)
    # Set the BDT branch address for faster reading, making
    # sure that Xbb-style leaflists are handled properly.
    if '.' in bdt_branch:
        branch_name, leaf_name = bdt_branch.split('.')
        branch = tree.GetBranch(branch_name)
        n_leaves = branch.GetNleaves()
        leaf_index = [leaf.GetName() for leaf in branch.GetListOfLeaves()].index(leaf_name)
        bdt_buffer = numpy.zeros(n_leaves, dtype=numpy.float32)
    else:
        branch_name = bdt_branch
        leaf_index = None
        bdt_buffer = numpy.zeros(1, dtype=numpy.float32)
    tree.SetBranchAddress(branch_name, bdt_buffer)
    # Clone the original tree and add the new branch.
    tree_new = tree.CloneTree(0)
    sb_weight = numpy.zeros(1, dtype=numpy.float64)
    tree_new.Branch('sb_weight', sb_weight, 'sb_weight/D')
    # Cache the Fill method for faster filling.
    fill_tree_new = tree_new.Fill
    for i, event in enumerate(tree, start=1):
        # Find the BDT bin containing the event. If the event is
        # found in the underflow bin, use the first bin instead.
        bdt_score = bdt_buffer[0] if leaf_index is None else bdt_buffer[leaf_index]
        bin_index = total_signal_prefit.FindBin(bdt_score) or 1
        # Calculate the S/(S+B) weight for the event.
        s = total_signal_prefit.GetBinContent(bin_index)
        b = total_background_postfit.GetBinContent(bin_index)
        sb_weight[0] = s / (s + b) if b > 0 else 0
        fill_tree_new()
        if i % 1000 == 0 or sb_weight[0] > 10:
            logger.info('Processing Entry #%s: BDT Score = %s, S/(S+B) = %s', i, bdt_score, sb_weight)
    tree_new.Write()
    outfile.Close()
    infile.Close()

def main():
    """Example usage:
    python add_sb_weight.py ZH_HToBB_ZToNuNu_M125_13TeV_powheg_pythia8.root ZH_HToBB_ZToNuNu_M125_13TeV_powheg_pythia8_new.root
    """
    logging.basicConfig(format='[%(name)s] %(levelname)s - %(message)s', level=logging.INFO)
    total_signal_prefit = get_total_signal_prefit(SIGNAL_SHAPES_PATH, SIGNAL_SHAPES_BIN, SIGNAL_SHAPES_PROCESSES)
    bin_edges = get_histogram_bin_edges(total_signal_prefit)
    total_background_postfit = get_total_background_postfit(MLFIT_PATH, MLFIT_BIN, bin_edges)
    add_sb_weight(sys.argv[1], sys.argv[2], BDT_BRANCH, total_signal_prefit, total_background_postfit)

if __name__ == '__main__':

    status = main()
    sys.exit(status)


# Creator: Tennant
# Email: Tennant_1999@outlook.com

import os
import os.path as osp

# PyTorch as the main lib for neural network
import torch
torch.backends.cudnn.benchmark = True
torch.multiprocessing.set_sharing_strategy('file_system')
import torch.nn as nn
import torchvision as tv
import numpy as np
import math
# Use visdom for moniting the training process
import visdom
from utils import Visualizer
from utils import setup_logger
from utils import rank_list_to_im

# Use yacs for training config management
# argparse for overwrite
from config import cfg
import argparse

# import losses and model
from losses import make_loss
from model import build_model, convert_model
from trainer import BaseTrainer

# dataset
from dataset import make_dataloader

from optim import make_optimizer, WarmupMultiStepLR

from evaluate import eval_func, euclidean_dist, re_rank
from utils.functions import norm_gallery
from tqdm import tqdm
from utils.functions import norm_gallery
import shutil
import os.path as osp

def parse_args():
    parser = argparse.ArgumentParser(description="ReID training")
    parser.add_argument('-c', '--config_file', type=str,
                        help='the path to the training config')
    parser.add_argument('-t', '--test', action='store_true',
                        default=False, help='Model test')
    parser.add_argument('--local_rank', default=0, type=int)
    parser.add_argument('opts', help='overwriting the training config'
                        'from commandline', default=None,
                        nargs=argparse.REMAINDER)

    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    if args.test:
        test(args)
    else:
        train(args)

def train(args):
    if args.config_file != "":
        print("The config file is", args.config_file)
        cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)

    cfg.freeze()

    output_dir = cfg.OUTPUT_DIR
    output_dir = osp.join(output_dir, cfg.EXP_NAME)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    shutil.copy(args.config_file, cfg.OUTPUT_DIR)

    num_gpus = torch.cuda.device_count()
    hash_bit = str(int(cfg.QUANT.n_book)*int(math.log(int(cfg.QUANT.intn_word), 2)))
    log_name = cfg.DATASETS.NAMES + '_' + hash_bit

    logger = setup_logger(log_name, output_dir, 0)
    logger.info('Using {} GPUS'.format(num_gpus))
    logger.info(args)
    logger.info('Running with config:\n{}'.format(cfg))

    train_dl, val_dl, num_query, num_classes = make_dataloader(cfg, num_gpus)

    model = build_model(cfg, num_classes)

    loss_func = make_loss(cfg, num_classes)

    trainer = BaseTrainer(cfg, model, train_dl, val_dl,
                          loss_func, num_query, num_gpus)

    for epoch in range(trainer.epochs):
        for batch in trainer.train_dl:

            trainer.step(batch)
            trainer.handle_new_batch()

        trainer.handle_new_epoch()

def test(args):
    if args.config_file != "":
        cfg.merge_from_file(args.config_file)
    cfg.merge_from_list(args.opts)
    cfg.freeze()
    output_dir = cfg.OUTPUT_DIR
    output_dir = osp.join(output_dir, cfg.EXP_NAME)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
    log_name = cfg.DATASETS.SOURCE_NAMES + '_' + cfg.DATASETS.NAMES
    logger = setup_logger('log_name', cfg.OUTPUT_DIR, 0, train=False)

    logger.info('Running with config:\n{}'.format(cfg))

    _, val_dl, num_query, num_classes = make_dataloader(cfg)

    model = build_model(cfg, num_classes)
    if cfg.TEST.MULTI_GPU:
        model = nn.DataParallel(model)
        model = convert_model(model)
        logger.info('Use multi gpu to inference')
    weight_dir = osp.join('./outputs', cfg.EXP_NAME, 'checkpoint')
    hash_bit = str(int(cfg.QUANT.n_book)*int(math.log(int(cfg.QUANT.intn_word), 2)))
    weight_name = cfg.DATASETS.SOURCE_NAMES + '_' + str(cfg.TEST.MAP) + '_' +hash_bit + '.pth'
    weight_path = osp.join(weight_dir, weight_name)

    para_dict = torch.load(weight_path)
    model_dict = model.state_dict()

    pretrained_dict = {k:v for k,v in para_dict.items() if not k.startswith('classifier') }
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)

    model.cuda()
    model.eval()

    feats, pids, camids, paths, descriptors = [], [], [], [], []
    with torch.no_grad():
        for batch in tqdm(val_dl, total=len(val_dl),
                         leave=False):
            data, pid, camid, path = batch
            paths.extend(list(path))
            data = data.cuda()
            feat, descriptor = model(data)[0].detach().cpu(), model(data)[1].detach().cpu()
            feats.append(feat)
            descriptors.append(descriptor)
            pids.append(pid)
            camids.append(camid)
    feats = torch.cat(feats, dim=0)
    descriptors = torch.cat(descriptors, dim=0)
    pids = torch.cat(pids, dim=0)
    camids = torch.cat(camids, dim=0)

    query_feat = feats[:num_query]
    query_pid = pids[:num_query]
    query_camid = camids[:num_query]
    query_path = np.array(paths[:num_query])

    gallery_feat = descriptors[num_query:]
    gallery_pid = pids[num_query:]
    gallery_camid = camids[num_query:]
    gallery_path = np.array(paths[num_query:])

    distmat = euclidean_dist(query_feat, gallery_feat)

    cmc, mAP, all_AP = eval_func(distmat.numpy(), query_pid.numpy(), gallery_pid.numpy(),
                         query_camid.numpy(), gallery_camid.numpy(),
                          dataset_name = cfg.DATASETS.NAMES)

    if cfg.TEST.VIS:
        worst_q = np.argsort(all_AP)[:cfg.TEST.VIS_Q_NUM]
        qid = query_pid[worst_q]
        q_im = query_path[worst_q]

        ind = np.argsort(distmat, axis=1)
        gid = gallery_pid[ind[worst_q]][..., :cfg.TEST.VIS_G_NUM]
        g_im = gallery_path[ind[worst_q]][..., :cfg.TEST.VIS_G_NUM]

        for idx in range(cfg.TEST.VIS_Q_NUM):
            sid = qid[idx] == gid[idx]
            im = rank_list_to_im(range(len(g_im[idx])), sid, q_im[idx], g_im[idx])

            im.save(osp.join(cfg.OUTPUT_DIR,
                    'worst_query_{}.jpg'.format(str(idx).zfill(2))))

    logger.info('Validation Result:')
    for r in cfg.TEST.CMC:
        logger.info('CMC Rank-{}: {:.2%}'.format(r, cmc[r-1]))
    logger.info('mAP: {:.2%}'.format(mAP))
    logger.info('-' * 20)

    if not cfg.TEST.RERANK:
        return

    distmat = re_rank(query_feat, gallery_feat)
    cmc, mAP, all_AP = eval_func(distmat, query_pid.numpy(), gallery_pid.numpy(),
                         query_camid.numpy(), gallery_camid.numpy(),
                          dataset_name = cfg.DATASETS.NAMES)

    logger.info('ReRanking Result:')
    for r in cfg.TEST.CMC:
        logger.info('CMC Rank-{}: {:.2%}'.format(r, cmc[r-1]))
    logger.info('mAP: {:.2%}'.format(mAP))
    logger.info('-' * 20)

if __name__ == '__main__':
    # from torchvision import models
    # resnet50 = models.resnet50(pretrained=True)
    main()


#!/usr/bin/env python3

"""CLIP guided sampling from a diffusion model."""

import argparse
from functools import partial
from pathlib import Path
import sys

from einops import repeat
import jax
import jax.numpy as jnp
from PIL import Image
from tqdm import tqdm, trange

from diffusion import get_model, get_models, load_params, sampling, utils

MODULE_DIR = Path(__file__).resolve().parent
sys.path.append(str(MODULE_DIR / 'CLIP_JAX'))

import clip_jax

def make_normalize(mean, std):
    mean = jnp.array(mean).reshape([3, 1, 1])
    std = jnp.array(std).reshape([3, 1, 1])

    def inner(image):
        return (image - mean) / std
    return inner

def norm2(x):
    """Normalizes a batch of vectors to the unit sphere."""
    return x / jnp.sqrt(jnp.sum(jnp.square(x), axis=-1, keepdims=True))

def spherical_dist_loss(x, y):
    """Computes 1/2 the squared spherical distance between the two arguments."""
    return jnp.square(jnp.arccos(jnp.sum(norm2(x) * norm2(y), axis=-1))) / 2

def main():
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    p.add_argument('prompt', type=str,
                   help='the text prompt')
    p.add_argument('--batch-size', '-bs', type=int, default=1,
                   help='the number of images per batch')
    p.add_argument('--checkpoint', type=str,
                   help='the checkpoint to use')
    p.add_argument('--clip-guidance-scale', '-cs', type=float, default=1000.,
                   help='the CLIP guidance scale')
    p.add_argument('--eta', type=float, default=1.,
                   help='the amount of noise to add during sampling (0-1)')
    p.add_argument('--init', type=str,
                   help='the init image')
    p.add_argument('--model', type=str, choices=get_models(), required=True,
                   help='the model to use')
    p.add_argument('-n', type=int, default=1,
                   help='the number of images to sample')
    p.add_argument('--seed', type=int, default=0,
                   help='the random seed')
    p.add_argument('--starting-timestep', '-st', type=float, default=0.9,
                   help='the timestep to start at (used with init images)')
    p.add_argument('--steps', type=int, default=1000,
                   help='the number of timesteps')
    args = p.parse_args()

    model = get_model(args.model)
    checkpoint = args.checkpoint
    if not checkpoint:
        checkpoint = MODULE_DIR / f'checkpoints/{args.model}.pkl'
    params = load_params(checkpoint)

    image_fn, text_fn, clip_params, _ = clip_jax.load('ViT-B/16')
    clip_patch_size = 16
    clip_size = 224
    normalize = make_normalize(mean=[0.48145466, 0.4578275, 0.40821073],
                               std=[0.26862954, 0.26130258, 0.27577711])

    target_embed = text_fn(clip_params, clip_jax.tokenize([args.prompt]))

    if args.init:
        _, y, x = model.shape
        init = Image.open(args.init).convert('RGB').resize((x, y), Image.LANCZOS)
        init = utils.from_pil_image(init)[None]

    key = jax.random.PRNGKey(args.seed)

    def clip_cond_fn_loss(x, key, params, clip_params, t, extra_args):
        dummy_key = jax.random.PRNGKey(0)
        v = model.apply(params, dummy_key, x, repeat(t, '-> n', n=x.shape[0]), extra_args)
        alpha, sigma = utils.t_to_alpha_sigma(t)
        pred = x * alpha - v * sigma
        clip_in = jax.image.resize(pred, (*pred.shape[:2], clip_size, clip_size), 'cubic')
        extent = clip_patch_size // 2
        clip_in = jnp.pad(clip_in, [(0, 0), (0, 0), (extent, extent), (extent, extent)], 'edge')
        sat_vmap = jax.vmap(partial(jax.image.scale_and_translate, method='cubic'),
                            in_axes=(0, None, None, 0, 0))
        scales = jnp.ones([pred.shape[0], 2])
        translates = jax.random.uniform(key, [pred.shape[0], 2], minval=-extent, maxval=extent)
        clip_in = sat_vmap(clip_in, (3, clip_size, clip_size), (1, 2), scales, translates)
        image_embeds = image_fn(clip_params, normalize((clip_in + 1) / 2))
        return jnp.sum(spherical_dist_loss(image_embeds, target_embed))

    def clip_cond_fn(x, key, t, extra_args, params, clip_params):
        grad_fn = jax.grad(clip_cond_fn_loss)
        grad = grad_fn(x, key, params, clip_params, t, extra_args)
        return grad * -args.clip_guidance_scale

    def run(key, n):
        tqdm.write('Sampling...')
        key, subkey = jax.random.split(key)
        noise = jax.random.normal(subkey, [n, *model.shape])
        key, subkey = jax.random.split(key)
        cond_params = {'params': params, 'clip_params': clip_params}
        sample_step = partial(sampling.jit_cond_sample_step,
                              extra_args={},
                              cond_fn=clip_cond_fn,
                              cond_params=cond_params)
        steps = utils.get_ddpm_schedule(jnp.linspace(1, 0, args.steps + 1)[:-1])
        if args.init:
            steps = steps[steps < args.starting_timestep]
            alpha, sigma = utils.t_to_alpha_sigma(steps[0])
            noise = init * alpha + noise * sigma
        return sampling.sample_loop(model, params, subkey, noise, steps, args.eta, sample_step)

    def run_all(key, n, batch_size):
        for i in trange(0, n, batch_size):
            key, subkey = jax.random.split(key)
            cur_batch_size = min(n - i, batch_size)
            outs = run(key, cur_batch_size)
            for j, out in enumerate(outs):
                utils.to_pil_image(out).save(f'out_{i + j:05}.png')

    try:
        run_all(key, args.n, args.batch_size)
    except KeyboardInterrupt:
        pass

if __name__ == '__main__':
    main()


# -*- coding: utf-8 -*-
"""
Created on Mon Oct 29 08:39:57 2018

@author: Marcus
"""

from psychopy import core
import numpy as np

def get_defaults(et_name):

    settings = Settings(et_name)

    if et_name == 'IS4_Large_Peripheral':
        settings.SAMPLING_RATE = 90
    elif et_name == 'Tobii Pro Spectrum':
        settings.SAMPLING_RATE = 600
        settings.TRACKING_MODE = 'human'
    elif et_name == 'Tobii Pro Nano':
        settings.SAMPLING_RATE = 60
    elif et_name == 'Tobii TX300':
        settings.SAMPLING_RATE = 300
    elif et_name == 'Tobii T60 XL':
        settings.SAMPLING_RATE = 60
    elif et_name == 'Tobii Pro X3-120 EPU':
        settings.SAMPLING_RATE = 120
    elif et_name == 'Tobii Pro X3-120':
        settings.SAMPLING_RATE = 120
    elif et_name == 'X2-60_Compact':
        settings.SAMPLING_RATE = 60
    elif et_name == 'X2-30_Compact':
        settings.SAMPLING_RATE = 40
    elif et_name == 'Tobii X60':
        settings.SAMPLING_RATE = 60
    elif et_name == 'Tobii X120':
        settings.SAMPLING_RATE = 120
    elif et_name == 'Tobii T60':
        settings.SAMPLING_RATE = 60
    elif et_name == 'Tobii T120':
        settings.SAMPLING_RATE = 120
    elif et_name == 'Tobii Pro Fusion':
        settings.SAMPLING_RATE = 120
    else:
        print('eye tracker type not supported')
        core.quit()
    return settings

class Connect(object):
    def __init__(self, in_arg='dummy'):
        '''  in_arg can be either string with eye tracker name
        or 'settings' generated by calling (and optionally modifying)
        the output from get_defaults()
        '''

        if isinstance(in_arg, str):
            if 'dummy' in in_arg:
                from titta import Tobii_dummy
                self.__class__ = Tobii_dummy.Connect
                self.__class__.__init__(self)
            else:
                from titta import Tobii
                self.__class__ = Tobii.myTobii
                self.__class__.__init__(self, in_arg)
        else:
            from titta import Tobii
            self.__class__ = Tobii.myTobii
            self.__class__.__init__(self, in_arg)

class Settings(object):

    def __init__(self, et_name):
        ''' Default settings for eye tracker
        '''

        self.graphics = Graphics()

        # Default name of et-data file
        self.FILENAME                    = 'test.tsv'

        # Tracking parameters
        self.TRACKER_ADDRESS  = ''           # If none is given, find one on the network
        self.SAMPLING_RATE = 600             # Set sampling rate of tracker
        self.eye_tracker_name = et_name
        self.TRACKING_MODE = 'Default'

        # Parameters for calibration
        self.PACING_INTERVAL = 1.0           # How long to present the dot until samples are collected
        self.AUTO_PACE = 1                   # autoaccept (2), semi autoaccept (1, accept first point, default)
                                        # of accept with space bar (0)

        self.ANIMATE_CALIBRATION = True      # Static or animated calibration dots
        self.RECORD_EYE_IMAGES_DURING_CALIBRATION = False
        self.RECORD_EXTERNAL_SIGNAL_DURING_CALIBRATION = False
        self.N_CAL_TARGETS = 5               # Valid: 0, 1, 5, 9, 13

        # List all possible calibration points (in Tobii's coordinate system)
        # (0.0, 0.0) is the upper left corner and (1.0, 1.0) is the lower right corner.

        # Define the 13 point array (reading order)
        self.CAL_TARGETS = np.array([[0.1, 0.1], [0.5, 0.1], [.9,.1],
                           [.3,.3], [.7,.3],
                           [.1,.5], [.5,.5], [.9,.5],
                           [.3,.7], [.7,.7],
                           [.1,.9], [.5,.9], [.9,.9]])

        self.VAL_POS_TOBII = np.array([[0.2, 0.5], [0.5, 0.8], [0.8, 0.5], [0.5, 0.2]])

        # CAL_POS_TOBII = np.array([[0.5, 0.5], [0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0],
        #                           [0.5, 1.0], [0.5, 0.0], [0.0, 0.5], [1.0, 0.5]])

        # if N_CAL_TARGETS == 0:
        #     CAL_POS_TOBII = []
        # elif N_CAL_TARGETS == 1:
        #     CAL_POS_TOBII = CAL_POS_TOBII[0, :]
        # elif N_CAL_TARGETS == 5:
        #     CAL_POS_TOBII = CAL_POS_TOBII[[0, 1, 2, 3, 4], :]

        # # VAL_POS_TOBII = np.array([[0.2, 0.2], [0.2, 0.8], [0.8, 0.2], [0.8, 0.8]])
        # VAL_POS_TOBII = np.array([[0.2, 0.5], [0.5, 0.8], [0.8, 0.5], [0.5, 0.2]])

        # # Scale the positions so they look good on the screen
        # scaling = 0.7
        # corr = 0.5 - (scaling * 0.5)
        # self.CAL_POS_TOBII = CAL_POS_TOBII * scaling + corr
        # self.VAL_POS_TOBII = VAL_POS_TOBII * scaling + corr

class Graphics(object):

    def __init__(self):
        ''' Default settings for graphics
        '''

        blue = tuple(np.array([37, 97, 163]) / 255.0 * 2 - 1)
        blue_active = tuple(np.array([11, 122, 244]) / 255.0 * 2 - 1)
        green = tuple(np.array([0, 120, 0]) / 255.0 * 2 - 1)
        red = tuple(np.array([150, 0, 0]) / 255.0 * 2 - 1)
        yellow = tuple(np.array([255, 255, 0]) / 255.0 * 2 - 1)
        yellow_linecolor = tuple(np.array([255, 255, 0]) / 255.0 * 2 - 1)

        self.blue = blue
        self.blue_active = blue_active

        self.TEXT_SIZE = 0.04 # Size of text
        self.TEXT_COLOR = 'white'

        self.ET_SAMPLE_RADIUS = 0.1 # in deg

        # SIze of calibration dots
        self.TARGET_SIZE=0.6 # in deg
        self.TARGET_SIZE_INNER=self.TARGET_SIZE / float(5)  # inner diameter of dot

        # Theses parameters are changed directly in the EThead class
        # self.HEAD_POS_CIRCLE_FIXED_COLOR = blue
        # self.HEAD_POS_CIRCLE_FIXED_RADIUS = 0.25

        # self.HEAD_POS_CIRCLE_MOVING_COLOR = yellow
        # self.HEAD_POS_CIRCLE_MOVING_FILLCOLOR = yellow
        # self.HEAD_POS_CIRCLE_MOVING_RADIUS = 0.25
        # self.HEAD_POS_CIRCLE_MOVING_MIN_RADIUS = 0.05

        self.POS_CAL_BUTTON = (0.5, -0.8)
        self.COLOR_CAL_BUTTON =  green
        self.WIDTH_CAL_BUTTON = 0.30
        self.HEIGHT_CAL_BUTTON = 0.08
        self.CAL_BUTTON = 'space'
        self.CAL_BUTTON_TEXT = 'calibrate (spacebar)'

        self.POS_RECAL_BUTTON = (-0.5, -0.8)
        self.COLOR_RECAL_BUTTON =  red
        self.WIDTH_RECAL_BUTTON = 0.30
        self.HEIGHT_RECAL_BUTTON = 0.08
        self.RECAL_BUTTON = 'c'
        self.RECAL_BUTTON_TEXT = 're-calibrate (c)'

        self.POS_REVAL_BUTTON = (-0.21, -0.8)
        self.COLOR_REVAL_BUTTON =  red
        self.WIDTH_REVAL_BUTTON = 0.30
        self.HEIGHT_REVAL_BUTTON = 0.08
        self.REVAL_BUTTON = 'v'
        self.REVAL_BUTTON_TEXT = 're-validate (v)'

        # Button for showing eye images
        self.POS_SETUP_BUTTON = (-0.5, -0.8)
        self.COLOR_SETUP_BUTTON = blue
        self.WIDTH_SETUP_BUTTON = 0.30
        self.HEIGHT_SETUP_BUTTON = 0.08
        self.SETUP_BUTTON = 'e'
        self.SETUP_BUTTON_TEXT = 'eye images (e)'

        self.POS_ACCEPT_BUTTON = (0.5, -0.8)
        self.COLOR_ACCEPT_BUTTON = green
        self.WIDTH_ACCEPT_BUTTON = 0.30
        self.HEIGHT_ACCEPT_BUTTON = 0.08
        self.ACCEPT_BUTTON = 'space'
        self.ACCEPT_BUTTON_TEXT = 'accept (spacebar)'

        self.POS_BACK_BUTTON = (-0.5, -0.8)
        self.COLOR_BACK_BUTTON = blue
        self.WIDTH_BACK_BUTTON = 0.30
        self.HEIGHT_BACK_BUTTON = 0.08
        self.BACK_BUTTON = 'b'
        self.BACK_BUTTON_TEXT = 'basic (b)'

        self.POS_GAZE_BUTTON = (0.8, 0.8)
        self.COLOR_GAZE_BUTTON = blue
        self.WIDTH_GAZE_BUTTON = 0.25
        self.HEIGHT_GAZE_BUTTON = 0.08
        self.GAZE_BUTTON = 'g'
        self.GAZE_BUTTON_TEXT = 'show gaze (g)'

        self.POS_CAL_IMAGE_BUTTON = (-0.8, 0.8)
        self.COLOR_CAL_IMAGE_BUTTON = (0.2, 0.2, 0.2)
        self.WIDTH_CAL_IMAGE_BUTTON = 0.25
        self.HEIGHT_CAL_IMAGE_BUTTON = 0.08
        self.CAL_IMAGE_BUTTON = 's'
        self.CAL_IMAGE_BUTTON_TEXT = 'Show calibration (s)'

        self.SETUP_DOT_OUTER_DIAMETER = 0.03 # Height unit
        self.SETUP_DOT_INNER_DIAMETER = 0.005

        # Parameters for eye images (default values are for Spectrum)
        self.EYE_IMAGE_SIZE = (0.5, 0.25)
        # self.EYE_IMAGE_SIZE_PIX = (175, 496)
        # self.EYE_IMAGE_SIZE_PIX_FULL_FRAME = (512, 640)
        self.EYE_IMAGE_POS_L = (-0.5, -0.4)
        self.EYE_IMAGE_POS_R = (0.5, -0.4)
        self.EYE_IMAGE_POS_L_1 = (-0.5, -0.4) # Used for the two additional fusion images
        self.EYE_IMAGE_POS_R_1 = (0.5, -0.4)

        # Parameters for tracking monitor (norm units)
        self.EYE_SIZE = 0.03
        self.EYE_COLOR_VALID = green
        self.EYE_COLOR_INVALID = red
        self.TRACKING_MONITOR_SIZE = [0.5, 0.5]
        self.TRACKING_MONITOR_POS = [0, 0.4]
        self.TRACKING_MONITOR_COLOR = [0.2, 0.2, 0.2]


# 获取烟遮挡数据的手部区域和对应烟坐标用于训练烟头检测
import os
import csv
from shutil import copyfile
import cv2
import numpy as np
import random
from tqdm import tqdm
import json
from PIL import Image
import math
from random import shuffle

def csv2yolov5():
    # 读取csv为yolov5标签文件
    imgPath = "E:\\pycharm-projects\\dataset\\检测数据集\\keruisite\\sigurate"
    csvfile = "E:\\pycharm-projects\\dataset\\检测数据集\\keruisite\\sigurate_smoke.csv"
    ImgsavePath = "E:\\pycharm-projects\\dataset\\DSMhand&smoke2\\temp\\keruisite\\sigurate"
    TxtsavePath = "E:\\pycharm-projects\\dataset\\DSMhand&smoke2\\temp\\keruisite\\sigurate"
    with open(csvfile, 'r') as f:
        reader = csv.reader(f)
        for row in reader:
            print(row)
            imgfile = os.path.join(imgPath, row[0])
            img = cv2.imdecode(np.fromfile(imgfile, dtype=np.uint8), 1)
            width = img.shape[1]
            height = img.shape[0]
            target_file = os.path.join(ImgsavePath, row[0])
            if not os.path.exists(ImgsavePath):
                os.makedirs(ImgsavePath)
            copyfile(imgfile, target_file)
            txtfile = os.path.join(TxtsavePath, row[0].split(".")[0] + ".txt")
            with open(txtfile, "a") as f1:
                # hand 0  smoke 1
                label = 1
                object_x_center = (int(row[1]) + int(row[3])) / 2.0 / float(width)
                object_y_center = (int(row[2]) + int(row[4])) / 2.0 / float(height)
                object_width = (int(row[3]) - int(row[1])) / float(width)
                object_height = (int(row[4]) - int(row[2])) / float(height)
                f1.write(str(label) + " " + str(object_x_center) + " " + str(object_y_center) + " " + str(object_width) + " " + str(object_height) + "\n")

def deleteNojson():
    # 删除没有json的图片
    for imgname in os.listdir("F:\\阜康测试视频\\frame-16\\labelme\\test\\img"):
        txtfile = "F:\\阜康测试视频\\frame-16\\labelme\\test\\json\\" + imgname.split(".")[0] + ".json"
        if not os.path.exists(txtfile):
            imagefile = "F:\\阜康测试视频\\frame-16\\labelme\\test\\img\\" + imgname
            print(imagefile)
            os.remove(imagefile)
    for jsonname in os.listdir("F:\\阜康测试视频\\frame-16\\labelme\\test\\1\\json"):
        imagefile = "F:\\阜康测试视频\\frame-16\\labelme\\test\\1\\img\\" + jsonname.split(".")[0] + ".jpg"
        if not os.path.exists(imagefile):
            jsonfile = "F:\\阜康测试视频\\frame-16\\labelme\\test\\1\\json\\" + jsonname
            print(jsonfile)
            os.remove(jsonfile)

def copyDir():
    # 将标注的图片和标注文件拷贝到新文件夹
    i = 0
    num0 = 0
    num1 = 0
    num2 = 0
    num3 = 0
    num4 = 0
    imgDir = "F:\\阜康测试视频\\28-09\\标注任务分配\\dir"
    saveDir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\"
    for root, dirs, files in os.walk(imgDir):
        for dir in dirs:
            # if (dir != "4"):
            #     continue
            imgDir1 = os.path.join(root, dir) # 一层
            for root1, dirs1, files1 in os.walk(imgDir1):
                for dir1 in dirs1:
                    imgDir2 = os.path.join(root1, dir1) # 两层
                    for root2, dirs2, files2 in os.walk(imgDir2):
                        for dir2 in dirs2:
                            imgDir3 = os.path.join(imgDir2, dir2) # 三层
                            print(imgDir3 + "\n")
                            for filename in tqdm(os.listdir(imgDir3)):
                                # if (filename.split(".")[-1] in ["json"]):
                                #     i += 1
                                #     jsonfile = os.path.join(imgDir3, filename)
                                #     imagename = filename.split(".")[0] + ".jpg"
                                #     imagefile = os.path.join(imgDir3, imagename)
                                #     if not os.path.isfile(imagefile):
                                #         continue
                                #     # print(imagefile)
                                #     newImgsaveDir = saveDir + dir + "\\image\\"
                                #     newJsonsaveDir = saveDir + dir + "\\json\\"
                                #     if not os.path.exists(newImgsaveDir):
                                #         os.makedirs(newImgsaveDir)
                                #     if not os.path.exists(newJsonsaveDir):
                                #         os.makedirs(newJsonsaveDir)
                                #     newImgfile = newImgsaveDir + imagename
                                #     newJsonfile = newJsonsaveDir + filename
                                #     copyfile(imagefile, newImgfile)
                                #     copyfile(jsonfile, newJsonfile)
                                # 统计需标注图片总数
                                if (filename.split(".")[-1] in ["jpg"]):
                                    if (dir == "0"):
                                        num0 += 1
                                    if (dir == "1"):
                                        num1 += 1
                                    if (dir == "2"):
                                        num2 += 1
                                    if (dir == "3"):
                                        num3 += 1
                                    if (dir == "4"):
                                        num4 += 1

    print("num0:" + str(num0) + "\n")
    print("num1:" + str(num1) + "\n")
    print("num2:" + str(num2) + "\n")
    print("num3:" + str(num3) + "\n")
    print("num4:" + str(num4) + "\n")

def get_hand_smoke():
    # 将28-09fukang（all文件夹）中包含手和香烟的标注拿出
    jsondir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\0-4\\4\\json\\"
    imgdir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\0-4\\4\\image\\"
    saveDir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\all_2\\"
    class_key = ['hand', 'cigarette']
    for num, jsonname in enumerate(tqdm(os.listdir(jsondir))):
        json_file = jsondir + jsonname
        flag = False
        with open(json_file, 'r', encoding='utf-8') as fp:
            data = json.load(fp)  # 加载json文件
            for shapes in data['shapes']:
                label = shapes['label']
                if label in class_key:
                    flag = True
                    break

        if flag:
            source_file = json_file
            if not os.path.exists(saveDir + "json\\"):
                os.makedirs(saveDir + "json\\")
            target_file = saveDir + "json\\" + jsonname
            copyfile(source_file, target_file)
            # os.remove(source_file)

            imgfile = imgdir + jsonname.split(".")[0] + ".jpg"
            if os.path.exists(imgfile):
                source_file = imgfile
                if not os.path.exists(saveDir + "image\\"):
                    os.makedirs(saveDir + "image\\")
                target_file = saveDir + "image\\" + jsonname.split(".")[0] + ".jpg"
                copyfile(source_file, target_file)
                # os.remove(source_file)

def get_smoke():
    # 将28-09fukang_2（all_2文件夹）训练集中包含香烟的标注拿出
    jsondir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2+DSMhand_smoke3\\val\\json\\"
    imgdir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2+DSMhand_smoke3\\val\\image\\"
    saveDir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2+DSMhand_smoke3-cigarette\\val\\"
    class_key = ['cigarette']
    for num, jsonname in enumerate(tqdm(os.listdir(jsondir))):
        json_file = jsondir + jsonname
        flag = False
        with open(json_file, 'r', encoding='utf-8') as fp:
            data = json.load(fp)  # 加载json文件
            for shapes in data['shapes']:
                label = shapes['label']
                if label in class_key:
                    flag = True
                    break

        if flag:
            source_file = json_file
            if not os.path.exists(saveDir + "json\\"):
                os.makedirs(saveDir + "json\\")
            target_file = saveDir + "json\\" + jsonname
            copyfile(source_file, target_file)
            # os.remove(source_file)

            imgfile = imgdir + jsonname.split(".")[0] + ".jpg"
            if os.path.exists(imgfile):
                source_file = imgfile
                if not os.path.exists(saveDir + "image\\"):
                    os.makedirs(saveDir + "image\\")
                target_file = saveDir + "image\\" + jsonname.split(".")[0] + ".jpg"
                copyfile(source_file, target_file)
                # os.remove(source_file)

def csv2json():
    # 读取csv生成json
    imgPath = "E:\\pycharm-projects\\dataset\\检测数据集\\数据堂叼烟\\sigurate"
    csvfile = "E:\\pycharm-projects\\dataset\\检测数据集\\数据堂叼烟\\sigurate_smoke.csv"
    ImgsavePath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2+数据堂叼烟\\train\\image"
    TxtsavePath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2+数据堂叼烟\\train\\json"
    with open(csvfile, 'r') as f:
        reader = csv.reader(f)
        for row in reader:
            print(row)
            imgfile = os.path.join(imgPath, row[0])
            img = cv2.imdecode(np.fromfile(imgfile, dtype=np.uint8), 1)
            width = img.shape[1]
            height = img.shape[0]
            target_file = os.path.join(ImgsavePath, row[0])
            if not os.path.exists(ImgsavePath):
                os.makedirs(ImgsavePath)
            # copyfile(imgfile, target_file)

            new_dict = { "version": "4.6.0", "flags": {}, "shapes": [], "imageData": None }  # label format
            new_dict["imagePath"] = row[0]
            new_dict["imageHeight"] = height
            new_dict["imageWidth"] = width
            object = {}
            object["label"] = "cigarette"
            object["group_id"] = None
            object["shape_type"] = "rectangle"
            object["flags"] = {}
            object["points"] = []
            x1y1 = [float(row[1]), float(row[2])]
            x2y2 = [float(row[3]), float(row[4])]
            object["points"].append(x1y1)
            object["points"].append(x2y2)
            new_dict["shapes"].append(object)

            json_path = os.path.join(TxtsavePath, row[0].split(".")[0])
            with open(json_path + '.json', 'a') as f:
                json.dump(new_dict, f, indent=4)

def getRandID():
    rootdir = "E:\\pycharm-projects\\dataset\\DSM_Dataset_class4_20220211_fukang\\"
    saveDir = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\all_2\\val\\"
    dict_id = {}
    for imgname in os.listdir(rootdir + "image"):
        # txtname = imgname.split(".")[0] + ".txt"
        # source_file = os.path.join("E:\\pycharm-projects\\dataset\\DSMhand&smoke1\\labels\\train", txtname)
        # target_file = os.path.join("E:\\pycharm-projects\\dataset\\DSMhand&smoke1\\labels\\val", txtname)
        # copyfile(source_file, target_file)
        # os.remove(source_file)
        ID = imgname.split("_")[1]
        if (ID in dict_id.keys()):
            dict_id[ID].append(imgname)
        else:
            dict_id[ID] = []
            dict_id[ID].append(imgname)

    # print(dict_id)
    # 随机选择1/9个ID做测试集
    val_nums = int(len(dict_id.keys()) // 9)
    val_id = random.sample(dict_id.keys(), val_nums)
    print(val_id)
    for id in tqdm(dict_id.keys()):
        if id in val_id:
            for imgname in dict_id[id]:
                source_file = rootdir + "image\\" + imgname
                if not os.path.exists(saveDir + "image\\"):
                    os.makedirs(saveDir + "image\\")
                target_file = saveDir + "image\\" + imgname
                copyfile(source_file, target_file)
                os.remove(source_file)

                txtfile = rootdir + "json\\" + imgname.split(".")[0] + ".json"
                if os.path.exists(txtfile):
                    source_file = txtfile
                    if not os.path.exists(saveDir + "json\\"):
                        os.makedirs(saveDir + "json\\")
                    target_file = saveDir + "json\\" + imgname.split(".")[0] + ".json"
                    copyfile(source_file, target_file)
                    os.remove(source_file)
        # else:
        #     for imgname in dict_id[id]:
        #         source_file = rootdir + "image\\" + imgname
        #         if not os.path.exists(saveDir + "train\\image\\"):
        #             os.makedirs(saveDir + "train\\image\\")
        #         target_file = saveDir + "train\\image\\" + imgname
        #         copyfile(source_file, target_file)
        #         # os.remove(source_file)
        #
        #         txtfile = rootdir + "json\\" + imgname.split(".")[0] + ".json"
        #         if os.path.exists(txtfile):
        #             source_file = txtfile
        #             if not os.path.exists(saveDir + "train\\json\\"):
        #                 os.makedirs(saveDir + "train\\json\\")
        #             target_file = saveDir + "train\\json\\" + imgname.split(".")[0] + ".json"
        #             copyfile(source_file, target_file)
        #             # os.remove(source_file)

def copy_plate():
    # 将香烟粘贴在阜康真实数据解决类别不均衡问题
    baseImgPath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2_copy_plate\\train\\image"
    baseJsonPath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2_copy_plate\\train\\json"
    saveImgPath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2_copy_plate\\train\\image_cp"
    saveJsonPath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2_copy_plate\\train\\json_cp"
    smokeImgPath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang_2_copy_plate\\smokeArea"
    radio = 5
    if not os.path.exists(saveImgPath):
        os.makedirs(saveImgPath)
    if not os.path.exists(saveJsonPath):
        os.makedirs(saveJsonPath)
    for jsonname in tqdm(os.listdir(baseJsonPath)):
        json_file = os.path.join(baseJsonPath, jsonname)
        hasFlag = False
        with open(json_file, 'r', encoding='utf-8') as fp:
            data = json.load(fp)  # 加载json文件
            for shapes in data['shapes']:
                label = shapes['label']
                if label == "cigarette":
                    hasFlag = True
                    break
        if hasFlag:
            continue
        imgname = jsonname.split(".")[0] + ".jpg"
        baseimgfile = os.path.join(baseImgPath, imgname)
        baseImg = Image.open(baseimgfile)
        size1 = baseImg.size

        smoke_list = os.listdir(smokeImgPath)
        select_index = random.randint(0, len(smoke_list) - 1)
        smokeimgfile = os.path.join(smokeImgPath, smoke_list[select_index])
        smoke_img = Image.open(smokeimgfile)
        size2 = smoke_img.size

        center = (int(size1[0] / 2), int(size1[1] / 2))

        nums = 0
        flag = False
        while ((size1[0] < size2[0]) or (size1[1] < size2[1]) or (size2[0] < 20 or size2[1] < 20)):
            # nums += 1
            # if (nums > 100):
            #     flag = True
            #     break
            select_index = random.randint(0, len(smoke_list) - 1)
            smokeimgfile = os.path.join(smokeImgPath, smoke_list[select_index])
            smoke_img = Image.open(smokeimgfile)
            size2 = smoke_img.size

        crop_x1 = center[0] - int(size2[0] / 2)
        crop_y1 = center[1] - int(size2[1] / 2)
        crop_x2 = center[0] + (size2[0] - int(size2[0] / 2))
        crop_y2 = center[1] + (size2[1] - int(size2[1] / 2))
        baseImg.paste(smoke_img, (crop_x1, crop_y1))

        # baseImg.show()

def randDir():
    # 将数据打乱均匀分到子文件夹
    imgDir = ""
    saveDir = ""
    imglist = os.listdir(imgDir)
    nums = len(imglist)
    val_id = random.sample(imglist, 2)

# 从文件夹中随机抽取一定数量图片，分发标注任务
def randSample():
    imgDir = "F:\\阜康测试视频\\22-10\\frame-16"
    saveDir = "F:\\阜康测试视频\\22-10\\标注任务分配\\"
    imglist = os.listdir(imgDir)
    nums = len(imglist)
    people = 4
    ave_nums = math.ceil(nums / people)
    # val_id = random.sample(imglist, 2000)
    shuffle(imglist)
    for i, imagename in enumerate(tqdm(imglist)):
        sourcefile = os.path.join(imgDir, imagename)
        sonFold = str(math.floor(i / ave_nums))
        desPath = saveDir + sonFold + "\\"
        if not os.path.exists(desPath):
            os.makedirs(desPath)
        desfile = desPath + imagename
        copyfile(sourcefile, desfile)

# 计算IOU
def bb_intersection_over_union(boxA, boxB):
    boxA = [int(x) for x in boxA]
    boxB = [int(x) for x in boxB]

    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])

    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)

    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)

    iou = interArea / float(boxAArea + boxBArea - interArea)

    return iou

def csv2json1():
    # 读取csv与已有json融合
    csvfile = "E:\\pycharm-projects\\dataset\\检测数据集\\keruisite\\keruisite_phone.csv"
    TxtsavePath = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang+DSMhand_smoke3\\train\\image\\D3_image\\phone_new_json"
    with open(csvfile, 'r') as f:
        reader = csv.reader(f)
        for row in reader:
            orgin_json = "F:\\阜康测试视频\\28-09\\整理返回标注数据\\28-09fukang+DSMhand_smoke3\\train\\image\\D3_image\\" + row[0].split(".")[0] + ".json"
            flag = False
            if os.path.exists(orgin_json):
                with open(orgin_json, 'r', encoding='utf-8') as fp:
                    data = json.load(fp)  # 加载json文件
                    for shapes in data['shapes']:
                        json_label = shapes['label']
                        if json_label == "cellphone":
                            flag = True
                            break
                    if (flag):
                        pass
                    else:
                        object = {}
                        object["label"] = "cellphone"
                        object["group_id"] = None
                        object["shape_type"] = "rectangle"
                        object["flags"] = {}
                        object["points"] = []
                        x1y1 = [float(row[1]), float(row[2])]
                        x2y2 = [float(row[3]), float(row[4])]
                        object["points"].append(x1y1)
                        object["points"].append(x2y2)
                        data["shapes"].append(object)
                newjsonfile = os.path.join(TxtsavePath, row[0].split(".")[0] + ".json")
                with open(newjsonfile, 'a') as f:
                    json.dump(data, f, indent=4)

# 合并两个json文件
def mergeJson():
    baseJsonPath1 = "F:\\阜康测试视频\\22-10\\frame-16-(0122-0210)-json"
    baseJsonPath2 = "F:\\阜康测试视频\\22-10\\返回标注文件\\all"
    TxtsavePath = "F:\\阜康测试视频\\22-10\\frame-16"
    for jsonname in tqdm(os.listdir(baseJsonPath1)):
        if jsonname.split(".")[-1] == "json":
            json_file1 = os.path.join(baseJsonPath1, jsonname)
            json_file2 = os.path.join(baseJsonPath2, jsonname)
            with open(json_file1, 'r', encoding='utf-8') as fp1:
                data1 = json.load(fp1)  # 加载json文件
                if os.path.exists(json_file2):
                    with open(json_file2, 'r', encoding='utf-8') as fp2:
                        data2 = json.load(fp2)  # 加载json文件
                        for shape in data2['shapes']:
                            data1['shapes'].append(shape)
                newjsonfile = os.path.join(TxtsavePath, jsonname)
                with open(newjsonfile, 'a') as f:
                    json.dump(data1, f, indent=4)

# 统计标注框数量
def getJsonNum():
    baseJsonPath = "F:\\阜康测试视频\\22-10\\返回标注文件\\淑茗\\0"
    nums = 0
    for jsonname in tqdm(os.listdir(baseJsonPath)):
        if jsonname.split(".")[-1] == "json":
            json_file = os.path.join(baseJsonPath, jsonname)
            with open(json_file, 'r', encoding='utf-8') as fp:
                data = json.load(fp)  # 加载json文件
                for i, shapes in enumerate(data['shapes']):
                    label = shapes['label']
                    nums += 1
    print(nums)

# 更改json文件里的错误标签
def editJson():
    baseJsonPath = "F:\\阜康测试视频\\14-21\\整理标注文件\\image\\jjc"
    TxtsavePath = "F:\\阜康测试视频\\14-21\\整理标注文件\\image\\jjc_edit"
    for jsonname in tqdm(os.listdir(baseJsonPath)):
        if jsonname.split(".")[-1] == "json":
            json_file = os.path.join(baseJsonPath, jsonname)
            hasFlag = False
            with open(json_file, 'r', encoding='utf-8') as fp:
                data = json.load(fp)  # 加载json文件
                for i, shapes in enumerate(data['shapes']):
                    label = shapes['label']
                    if label == "手机":
                        hasFlag = True
                        data['shapes'][i]['label'] = "cellphone"
            if (hasFlag):
                newjsonfile = os.path.join(TxtsavePath, jsonname)
                with open(newjsonfile, 'a') as f:
                    json.dump(data, f, indent=4)

# 随机选择1/9个ID做测试集，并保证有特定类别样本
def getRandID1():
    label_list = ['cigarette', 'cellphone']
    rootdir = "E:\\pycharm-projects\\dataset\\DSM_Dataset_class4_20220211_fukang\\train\\fukang\\28-09fukang\\"
    saveDir = "E:\\pycharm-projects\\dataset\\DSM_Dataset_class4_20220211_fukang\\val\\"
    dict_id = {}
    for jsonname in os.listdir(rootdir + "json"):
        # txtname = imgname.split(".")[0] + ".txt"
        # source_file = os.path.join("E:\\pycharm-projects\\dataset\\DSMhand&smoke1\\labels\\train", txtname)
        # target_file = os.path.join("E:\\pycharm-projects\\dataset\\DSMhand&smoke1\\labels\\val", txtname)
        # copyfile(source_file, target_file)
        # os.remove(source_file)
        imgname = jsonname.split(".")[0] + ".jpg"
        ID = jsonname.split("_")[1]
        json_file = rootdir + "json\\" + jsonname
        flag = False
        with open(json_file, 'r', encoding='utf-8') as fp:
            data = json.load(fp)  # 加载json文件
            for i, shapes in enumerate(data['shapes']):
                label = shapes['label']
                if label in label_list:
                    flag = True
        if flag:
            if (ID in dict_id.keys()):
                dict_id[ID].append(imgname)
            else:
                dict_id[ID] = []
                dict_id[ID].append(imgname)

    # print(dict_id)
    # 随机选择1/9个ID做测试集
    val_nums = int(len(dict_id.keys()) // 9)
    val_id = random.sample(dict_id.keys(), val_nums)
    print(val_id)
    for id in tqdm(dict_id.keys()):
        if id in val_id:
            for imgname in dict_id[id]:
                source_file = rootdir + "image\\" + imgname
                if not os.path.exists(saveDir + "image\\"):
                    os.makedirs(saveDir + "image\\")
                target_file = saveDir + "image\\" + imgname
                copyfile(source_file, target_file)
                os.remove(source_file)

                txtfile = rootdir + "json\\" + imgname.split(".")[0] + ".json"
                if os.path.exists(txtfile):
                    source_file = txtfile
                    if not os.path.exists(saveDir + "json\\"):
                        os.makedirs(saveDir + "json\\")
                    target_file = saveDir + "json\\" + imgname.split(".")[0] + ".json"
                    copyfile(source_file, target_file)
                    os.remove(source_file)

if __name__ == "__main__":
    # get_smoke()
    # csv2json1()
    # editJson()
    # getRandID1()
    # randSample()
    # getJsonNum()
    # randSample()
    mergeJson()


import jax
import jax.numpy as jnp

# Load models & tokenizer
from dalle_mini.model import DalleBart
from vqgan_jax.modeling_flax_vqgan import VQModel
from transformers import BartTokenizer, CLIPProcessor, FlaxCLIPModel
import wandb

from flax.jax_utils import replicate
from functools import partial
from dalle_mini.text import TextNormalizer
from dalle_mini.model.modeling import FlaxBartModule, FlaxBartForConditionalGenerationModule
from flax.training.common_utils import shard

from flax.training.common_utils import shard_prng_key
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from tqdm.notebook import trange
import time
import random

import jax
import flax.linen as nn

from transformers.models.bart.modeling_flax_bart import (
    FlaxBartDecoder,
    FlaxBartEncoder,
    FlaxBartForConditionalGeneration,
    FlaxBartForConditionalGenerationModule,
    FlaxBartModule,
)
from transformers import BartTokenizer, FlaxBartForConditionalGeneration, BartConfig

OUTPUT_VOCAB_SIZE = 16384 + 1  # encoded image token space + 1 for bos
OUTPUT_LENGTH = 256 + 1  # number of encoded tokens + 1 for bos
BOS_TOKEN_ID = 16384
BASE_MODEL = 'facebook/bart-large'

class CustomFlaxBartModule(FlaxBartModule):
    def setup(self):
        # we keep shared to easily load pre-trained weights
        self.shared = nn.Embed(
            self.config.vocab_size,
            self.config.d_model,
            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
            dtype=self.dtype,
        )
        # a separate embedding is used for the decoder
        self.decoder_embed = nn.Embed(
            OUTPUT_VOCAB_SIZE,
            self.config.d_model,
            embedding_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
            dtype=self.dtype,
        )
        self.encoder = FlaxBartEncoder(self.config, dtype=self.dtype, embed_tokens=self.shared)

        # the decoder has a different config
        decoder_config = BartConfig(self.config.to_dict())
        decoder_config.max_position_embeddings = OUTPUT_LENGTH
        decoder_config.vocab_size = OUTPUT_VOCAB_SIZE
        self.decoder = FlaxBartDecoder(decoder_config, dtype=self.dtype, embed_tokens=self.decoder_embed)

class CustomFlaxBartForConditionalGenerationModule(FlaxBartForConditionalGenerationModule):
    def setup(self):
        self.model = CustomFlaxBartModule(config=self.config, dtype=self.dtype)
        self.lm_head = nn.Dense(
            OUTPUT_VOCAB_SIZE,
            use_bias=False,
            dtype=self.dtype,
            kernel_init=jax.nn.initializers.normal(self.config.init_std, self.dtype),
        )
        self.final_logits_bias = self.param("final_logits_bias", self.bias_init, (1, OUTPUT_VOCAB_SIZE))

class CustomFlaxBartForConditionalGeneration(FlaxBartForConditionalGeneration):
    module_class = CustomFlaxBartForConditionalGenerationModule

class Text2Image():
    def __init__(self, DALLE_MODEL, VQGAN_REPO, VQGAN_COMMIT_ID, CLIP_REPO, DALLE_COMMIT_ID=None, CLIP_COMMIT_ID=None):
        tokenizer = BartTokenizer.from_pretrained(DALLE_MODEL, revision=DALLE_COMMIT_ID)
        model = CustomFlaxBartForConditionalGeneration.from_pretrained(
            DALLE_MODEL, revision=DALLE_COMMIT_ID
        )
        self.model = model
        self.tokenizer = tokenizer

        # Load VQGAN
        self.vqgan = VQModel.from_pretrained(VQGAN_REPO, revision=VQGAN_COMMIT_ID)

        # Load CLIP
        self.clip = FlaxCLIPModel.from_pretrained(CLIP_REPO, revision=CLIP_COMMIT_ID)
        self.processor = CLIPProcessor.from_pretrained(CLIP_REPO, revision=CLIP_COMMIT_ID)

        self.bart_params = replicate(model.params)
        self.vqgan_params = replicate(self.vqgan.params)
        self.clip_params = replicate(self.clip.params)

    def inference(self, prompt, num_candidates = 32, k = 8):
        prompt = [prompt] * jax.device_count()
        inputs = self.tokenizer(
            prompt,
            return_tensors="jax",
            padding="max_length",
            truncation=True,
            max_length=128,
        ).data
        inputs = shard(inputs)
        model = self.model
        vqgan = self.vqgan

        def generate(input, rng, params):
            return model.generate(
                **input,
                max_length=257,
                num_beams=1,
                do_sample=True,
                prng_key=rng,
                eos_token_id=50000,
                pad_token_id=50000,
                params=params,
            )
        p_generate = jax.pmap(generate, "batch")

        def get_images(indices, params):
            return vqgan.decode_code(indices, params=params)
        p_get_images = jax.pmap(get_images, "batch")

        def custom_to_pil(x):
            x = np.clip(x, 0.0, 1.0)
            x = (255 * x).astype(np.uint8)
            x = Image.fromarray(x)
            if not x.mode == "RGB":
                x = x.convert("RGB")
            return x

        # inference
        all_images = []
        for _ in range(num_candidates // jax.device_count()):
            key = random.randint(0, 1e7)
            rng = jax.random.PRNGKey(key)
            rngs = jax.random.split(rng, jax.local_device_count())
            indices = p_generate(inputs, rngs, self.bart_params).sequences
            indices = indices[:, :, 1:]

            images = p_get_images(indices, self.vqgan_params)
            images = np.squeeze(np.asarray(images), 1)
            for image in images:
                all_images.append(custom_to_pil(image))

        sorted_images = self.clip_top_k(all_images, prompt, k=k)
        return sorted_images

    def clip_top_k(self, images, prompt, k=8):
        inputs = self.processor(text=prompt, images=images, return_tensors="np", padding=True)
        outputs = self.clip(**inputs)
        logits = outputs.logits_per_text
        scores = np.array(logits[0]).argsort()[-k:][::-1]
        return [images[score] for score in scores]

    def captioned_strip(self, images, caption=None, rows=1):
        increased_h = 0 if caption is None else 48
        w, h = images[0].size[0], images[0].size[1]
        img = Image.new("RGB", (len(images) * w // rows, h * rows + increased_h))
        for i, img_ in enumerate(images):
            img.paste(img_, (i // rows * w, increased_h + (i % rows) * h))

        if caption is not None:
            draw = ImageDraw.Draw(img)
            font = ImageFont.truetype(
                "/usr/share/fonts/truetype/liberation2/LiberationMono-Bold.ttf", 40
            )
            draw.text((20, 3), caption, (255, 255, 255), font=font)
        return img

if __name__ == "__main__":
    # Model references

    # dalle-mini
    # DALLE_MODEL = "dalle-mini/dalle-mini/model-3bqwu04f:latest"  # can be wandb artifact or 🤗 Hub or local folder
    # DALLE_MODEL = "tools/inference/model"
    DALLE_MODEL = "flax-community/dalle-mini"
    DALLE_COMMIT_ID = "4d34126d0df8bc4a692ae933e3b902a1fa8b6114"  # used only with 🤗 hub

    # VQGAN model
    # VQGAN_REPO = "dalle-mini/vqgan_imagenet_f16_16384"
    VQGAN_REPO = "dalle-mini/vqgan_imagenet_f16_16384"
    VQGAN_COMMIT_ID = "90cc46addd2dd8f5be21586a9a23e1b95aa506a9"

    # CLIP model
    CLIP_REPO = "openai/clip-vit-base-patch32"
    CLIP_COMMIT_ID = None

    gen = Text2Image(DALLE_MODEL, VQGAN_REPO=VQGAN_REPO, VQGAN_COMMIT_ID=VQGAN_COMMIT_ID, CLIP_REPO=CLIP_REPO, DALLE_COMMIT_ID=DALLE_COMMIT_ID)
    ranked_images = gen.inference("A woman standing in a room with a computer and other items")


import numpy as np
import tensorflow as tf
import random
import os
import math
import argparse
from araml_dataloader import Gen_Data_loader, Dis_dataloader
from araml_generator import Generator
from araml_rewarder import Rewarder
from convert import convert

import os
import time
import lm.lm_train as lm_train

SEED = int(time.time()) # random seed

parser = argparse.ArgumentParser()

parser.add_argument('--task_name', default='coco', type=str, required=True,
                    help='The name of task, i.e. coco or emnlp')

args = parser.parse_args()

if args.task_name == 'coco':
    from conf_coco import gen_config, disc_config, adv_config, lm_config
    print('coco')
else:
    from conf_emnlp import gen_config, disc_config, adv_config, lm_config
    print('emnlp')

#  Generator Hyper-Parameters Initialization
EMB_DIM = gen_config.EMB_DIM # embedding dimension
HIDDEN_DIM = gen_config.HIDDEN_DIM # hidden state dimension of lstm cell
SEQ_LENGTH = gen_config.SEQ_LENGTH # sequence length
START_TOKEN = gen_config.START_TOKEN
PRE_EPOCH_NUM = gen_config.PRE_EPOCH_NUM # MLE pretraining epochs
PRE_LOG_STEP = gen_config.PRE_LOG_STEP
BATCH_SIZE = gen_config.BATCH_SIZE
R_rate_gen = gen_config.R_rate_gen # learning rate of generator

#  Discriminator Hyper-Parameters Initialization
MID_LAYER_G = disc_config.MID_LAYER_G
MID_LAYER_R = disc_config.MID_LAYER_R
re_dropout_keep_prob = disc_config.re_dropout_keep_prob # keep probability in discriminator
re_l2_reg_lambda = disc_config.re_l2_reg_lambda # regularization coefficient
R_rate = disc_config.R_rate # learning rate of discriminator
PRE_GENERATE_NUM = disc_config.PRE_GENERATE_NUM # pretraining epochs = PRE_GENERATE_NUM * PRE_EPOCH_NUM_D
PRE_EPOCH_NUM_D = disc_config.PRE_EPOCH_NUM_D

#  Language Model Hyper-Parameters Initialization
lm_fw_traindir = lm_config.lm_fw_traindir
lm_bw_traindir = lm_config.lm_bw_traindir

#  Pre-training / Adversarial Training Hyper-Parameters
ADV_LOG_STEP = adv_config.ADV_LOG_STEP
tau = adv_config.tau # temperature for sampling
log_dir = adv_config.log_dir
result_dir = adv_config.result_dir
train_dir = adv_config.train_dir
TOTAL_BATCH = adv_config.TOTAL_BATCH # Training batches in adversarial training
positive_file = adv_config.positive_file # Training set directory
negative_file = adv_config.negative_file
vocab_file = adv_config.vocab_file # Vocabulary directory
eval_file_prefix = adv_config.eval_file_prefix
result_prefix = adv_config.result_prefix
raml_file = adv_config.raml_file
generated_num = adv_config.generated_num
result_generated_num = adv_config.result_generated_num
generated_num_pre = adv_config.generated_num_pre
restore = adv_config.restore # whether to restore the checkpoint
vocab_size = adv_config.vocab_size # vocab size of coco
sample_size = adv_config.sample_size # the number of samples from Ps
gen_update_batch = adv_config.gen_update_batch
dis_update_batch = adv_config.dis_update_batch
adv_generate_num = adv_config.adv_generate_num
alpha = adv_config.alpha # rescaling coefficient
PAD_ID = adv_config.PAD_ID # Token ID of PAD

# Generate samples to file
def generate_samples(sess, trainable_model, batch_size, generated_num, output_file):
    generated_samples = []
    for _ in range(int(generated_num / batch_size)):
        samples = trainable_model.generate(sess)
        generated_samples.extend(samples)

    with open(output_file, 'w') as fout:
        for sample in generated_samples:
            buffer = ' '.join([str(x) for x in sample]) + '\n'
            fout.write(buffer)

    return generated_samples

# Training process (one epoch)
def pre_train_epoch(sess, trainable_model, data_loader):
    supervised_g_losses = []
    data_loader.reset_pointer()

    for it in xrange(data_loader.num_batch):
        batch = data_loader.next_batch()
        _, g_loss = trainable_model.pretrain_step(sess, batch, np.ones(BATCH_SIZE*sample_size))
        supervised_g_losses.append(g_loss)

    return np.mean(supervised_g_losses)

def main():
    # Set random seed
    random.seed(SEED)
    np.random.seed(SEED)

    # Construct generator / discriminator dataloader
    gen_data_loader = Gen_Data_loader(BATCH_SIZE, tau, vocab_size, sample_size)
    dis_data_loader = Dis_dataloader(BATCH_SIZE)
    gen_data_loader.create_batches(positive_file, PAD_ID, SEQ_LENGTH)

    # Create generator / discriminator model
    generator = Generator(vocab_size, BATCH_SIZE * sample_size, EMB_DIM, HIDDEN_DIM, SEQ_LENGTH, START_TOKEN, MID_LAYER_G, R_rate_gen)
    rewarder = Rewarder(vocab_size, BATCH_SIZE, EMB_DIM * 2, HIDDEN_DIM * 2, SEQ_LENGTH, l2_reg_lambda=re_l2_reg_lambda)

    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    sess.run(tf.global_variables_initializer())

    # Saver for pretraining and adversarial training
    saver_pre = tf.train.Saver(tf.global_variables(), max_to_keep=5, keep_checkpoint_every_n_hours=1.0)
    saver_adv = tf.train.Saver(tf.global_variables(), max_to_keep=5, keep_checkpoint_every_n_hours=1.0)

    # Construct the forward/backward language model for sampling
    lm_model_fw = lm_train.create_model(sess, vocab_size, EMB_DIM, HIDDEN_DIM, SEQ_LENGTH, START_TOKEN, PAD_ID, lm_fw_traindir, 'lm_model_fw')
    lm_model_bw = lm_train.create_model(sess, vocab_size, EMB_DIM, HIDDEN_DIM, SEQ_LENGTH, START_TOKEN, PAD_ID, lm_bw_traindir, 'lm_model_bw')
    lm_train.lm_train(lm_model_fw, sess, 'lm_model_fw', PRE_EPOCH_NUM, PRE_LOG_STEP, lm_fw_traindir , gen_data_loader)
    lm_train.lm_train(lm_model_bw, sess, 'lm_model_bw', PRE_EPOCH_NUM, PRE_LOG_STEP, lm_bw_traindir, gen_data_loader)
    gen_data_loader.create_batches_augment(raml_file, lm_model_fw, lm_model_bw, sess, PAD_ID)
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    log = open(os.path.join(log_dir, 'experiment-log-'+str(tau)+'.txt'), 'w')

    if restore is False:
        print 'Start pre-training generator...'
        log.write('pre-training...\n')
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        for epoch in xrange(PRE_EPOCH_NUM):
            loss = pre_train_epoch(sess, generator, gen_data_loader)
            if epoch % PRE_LOG_STEP == 0:
                print 'pre-train epoch ', epoch, 'test_loss ', loss
                buffer = 'epoch:\t' + str(epoch) + '\tnll:\t' + str(loss) + '\n'
                log.write(buffer)
                generate_samples(sess, generator, BATCH_SIZE * sample_size, result_generated_num,
                                 result_prefix + '_pre_' + str(epoch))

        print 'Start pre-training discriminator...'
        start = time.time()
        for _ in range(PRE_GENERATE_NUM):
            generate_samples(sess, generator, BATCH_SIZE, generated_num_pre, negative_file)
            dis_data_loader.load_train_data(positive_file, negative_file, SEQ_LENGTH)

            for _ in range(PRE_EPOCH_NUM_D):
                dis_data_loader.reset_pointer()
                r_losses = []
                for it in xrange(dis_data_loader.num_batch):
                    x_text, y_label = dis_data_loader.next_batch()
                    _, r_loss = rewarder.reward_train_step(sess, x_text, y_label, R_rate, re_dropout_keep_prob)
                    r_losses.append(r_loss)
                print 'discriminator_loss', np.mean(r_losses)
        speed = time.time() - start
        print 'Discriminator pre_training Speed:{:.3f}'.format(speed)

        train_dir_pre = os.path.join(train_dir, 'pre')
        if not os.path.exists(train_dir_pre):
            os.makedirs(train_dir_pre)

        checkpoint_path = os.path.join(train_dir_pre, 'exper_pre.ckpt')
        saver_pre.save(sess, checkpoint_path)
    else:
        print 'Restore pretrained model ...'
        log.write('Restore pre-trained model...\n')
        train_dir_pre = os.path.join(train_dir, 'pre')
        ckpt = tf.train.get_checkpoint_state(train_dir_pre)
        saver_pre.restore(sess, ckpt.model_checkpoint_path)

    print 'Start Adversarial Training...'
    log.write('adversarial training...\n')

    train_dir_adv = os.path.join(train_dir, 'adv')
    if not os.path.exists(train_dir_adv):
        os.makedirs(train_dir_adv)

    for total_batch in range(TOTAL_BATCH):
        if total_batch % ADV_LOG_STEP == 0 or total_batch == TOTAL_BATCH - 1:
            generate_samples(sess, generator, BATCH_SIZE * sample_size, generated_num, eval_file_prefix + '_batch_'+ str(total_batch))
            generate_samples(sess, generator, BATCH_SIZE * sample_size, result_generated_num, result_prefix + '_batch_' + str(total_batch))
            checkpoint_path = 'exper_' + str(total_batch) + '.ckpt'
            checkpoint_path = os.path.join(train_dir_adv, checkpoint_path)
            saver_adv.save(sess, checkpoint_path)

            # convert token id to token
            convert(eval_file_prefix + '_batch_'+ str(total_batch), vocab_file, log_dir)

        # Train the generator
        start = time.time()
        g_losses = []
        for it in range(gen_update_batch):
            batch_x = gen_data_loader.next_batch_augment() # (batch_size * sample_size, sequence_length)
            rewards = rewarder.get_reward(sess, batch_x)[0]
            weights = []
            for it3 in range(BATCH_SIZE):
                tmp_weights = rewards[it3*sample_size:(it3+1)*sample_size]
                exp_weights = [math.exp(tmp_weights[w] / alpha) for w in range(sample_size)]
                weights.extend(exp_weights)
            normalizer = sum(weights)
            weights = np.array([w / normalizer for w in weights])
            _, g_loss = generator.pretrain_step(sess, batch_x, weights)
            g_losses.append(g_loss)
        speed = time.time() - start
        print 'Generator training {} round, Speed:{:.3f}, Loss:{:.3f}'.format(total_batch, speed, np.mean(g_losses))

        # Train the discriminator
        start = time.time()
        r_loss_list = []
        for _ in range(adv_generate_num):
            generate_samples(sess, generator, BATCH_SIZE * sample_size, generated_num, negative_file)
            dis_data_loader.load_train_data(positive_file, negative_file, SEQ_LENGTH)
            for _ in range(dis_update_batch):
                dis_data_loader.reset_pointer()
                for it in xrange(dis_data_loader.num_batch):
                    x_text, y_label = dis_data_loader.next_batch()
                    _, r_loss = rewarder.reward_train_step(sess, x_text, y_label, R_rate, re_dropout_keep_prob)
                    r_loss_list.append(r_loss)
        avg_loss = np.mean(r_loss_list)
        speed = time.time() - start
        print 'Discriminator training {} round, Speed:{:.3f}, Loss:{:.3f}'.format(total_batch, speed, avg_loss)

    log.close()

if __name__ == '__main__':
    main()


import tensorflow as tf
import cv2
import numpy as np
IMG_SIZE=(224,224)

Model=tf.keras.models.load_model ('Ensamble_Model(Ori).h5')
path='C:/Users/mariotiara/Desktop/GUI/Sampel/ffc04fed30e6_(0)_Normal.jpeg'

img=tf.keras.preprocessing.image.load_img(path,target_size=IMG_SIZE)

x = tf.keras.preprocessing.image.img_to_array(img)
x = np.expand_dims(x, axis=0)

Predict=Model.predict(x)

print("Normal:",Predict[0][0]*100)
print("Mild:",Predict[0][1]*100)
print("Moderate:",Predict[0][2]*100)
print("Severe:",Predict[0][3]*100)


#!/usr/bin/env python
import numpy as np
# import tensorflow as tf
# from ov_utils.myutils import make_mask, template_matching, get_surfacenormal, get_bbox, get_depth, get_img
from ov_utils.config import normal_vector
# import ov_utils.file_utils as file_utils
import eVM_utils.utils as eVM_utils
from eVM_utils.eVM_model import CNN, encodeVisualMemory, encodeVisualMemoryRelatedPath
from recovery_policy import Recovery
import torch
import sys
if '/opt/ros/kinetic/lib/python2.7/dist-packages' in sys.path:
    sys.path.remove('/opt/ros/kinetic/lib/python2.7/dist-packages')
import os
import joblib
import cv2
from PIL import Image

from ove_policy import ImageNavHFNetPolicyQ, HFNet
from torchvision import transforms
from pathlib import Path
import copy

"""
Active Navigation Module:
A module that maintains visual memory and provides three kinds of active navgitation guidance;
(1) recovery guidance, (2) exploration guidance, (3) optimal viewpoint guidance.
Author
------
Yunho Choi, Obin Kwon, Nuri Kim, Hwiyeon Yoo
"""

class ActiveNavigationModule():
    """Active Navigation Module"""
    def __init__(self, path_direction='homing', cuda=True):
        # map_manager = MapManager()
        # self.map = map_manager.getMap()
        # self.args = args
        self.action_dim = 3
        self.memory_dim = 256
        self.feature_dim = 512
        self.GRU_size = 512
        self.img_size = 224
        self.path_length = 0
        self.list2encode = []
        self.vis_mem = []
        self.feats_along_path = []
        self.rel_pose = None
        self.path_direction = path_direction
        self.cuda = cuda
        self.encode_im = CNN(self.feature_dim, self.img_size)
        if self.path_direction=='following':
            self.vis_mem_encoder = encodeVisualMemory(self.action_dim, self.memory_dim, self.feature_dim)
        if self.path_direction=='homing':
            self.vis_mem_encoder = encodeVisualMemoryRelatedPath(self.action_dim, self.memory_dim, self.feature_dim)
        self.vis_mem_encoder_model = './data_exp/model/{}/best.pth'.format(self.path_direction)
        self.ove_policy = ImageNavHFNetPolicyQ()
        self.ove_policy_model = './data_exp/model/periodic_500000.q_net'
        self.ove_extractor_model = './data_exp/model/hfnet'
        self.num_keypoints = 48
        self.outputs = ['global_descriptor', 'keypoints', 'local_descriptors', 'scores','local_descriptor_map' ]

        self.enable_recovery = False
        self.recovery_policy = Recovery(self.action_dim, self.memory_dim, self.feature_dim, self.GRU_size)
        self.recovery_guidance = None

        self.enable_exploration = False
        self.exploration_policy = None
        self.exploration_guidance = None

        self.enable_ove = False
        self.optimal_viewpoint_guidance = [0, 0, 0]

        data_list = [os.path.join('./data_exp/img_trajectory/{}'.format(self.path_direction), x) for x in os.listdir('./data_exp/img_trajectory/{}'.format(self.path_direction))]
        data = joblib.load(np.random.choice(data_list))
        self.img_list = data['rgb']
        self.guidance_list = data['action']
        self.position_list = data['position']
        rotation_quat = data['rotation']
        self.rotation_list = []
        for quat in rotation_quat:
            self.rotation_list.append(2*np.arctan2(np.linalg.norm(quat[1:]), quat[0]))
        # if self.enable_ove:
        #     self.ove_data_folder = './data_exp/optimal_viewpoint/'
        #     image_list, sf_list, bbox_list, depth_list = file_utils.get_files(self.ove_data_folder)
        #     self.im_paths, self.target_pois = file_utils.get_annos(self.ove_data_folder + 'anno/')

        # self.initialize()

    def initialize(self):
        try:
            model_dict = torch.load(self.vis_mem_encoder_model, map_location='cpu')
            im_encoder_model_dict = {}
            vis_mem_encoder_model_dict = {}
            recovery_policy_model_dict = {}
            for k, v in model_dict.items():
                if 'encode_im' in k:
                    name = k.replace('encode_im.', '')
                    im_encoder_model_dict[name] = v
                if 'rel_path_weight_fc' in k or 'visual_memory_fc' in k:
                    vis_mem_encoder_model_dict[k] = v
                if 'GRU' in k:
                    recovery_policy_model_dict[k] = v
            self.encode_im.load_state_dict(im_encoder_model_dict)
            self.vis_mem_encoder.load_state_dict(vis_mem_encoder_model_dict)
            self.recovery_policy.load_state_dict(recovery_policy_model_dict)
            # self.encode_im.cuda()
            # self.vis_mem_encoder.cuda()
            # self.recovery_policy.cuda()
            self.encode_im.eval()
            self.vis_mem_encoder.eval()
            self.recovery_policy.eval()

            ove_policy_state_dict = torch.load(self.ove_policy_model, map_location='cpu')
            self.ove_policy.load_state_dict(ove_policy_state_dict)
            self.ove_policy.eval()
            self.ove_extractor = HFNet(self.ove_extractor_model, self.outputs)
            print("load pretrained encodeVisualMemory model")
        except:
            print("Cannot load pretrained encodeVisualMemory model")
            pass

        self.enable_recovery = True
        self.enable_ove = True

        return True

    def encodeVisualMemory(self, img, guidance, random_action=False, flush=False, exp_active=False):
        """
        Visual Memory Encoder Submodule:
        A module running consistently which encodes visual trajectory information from the previous node to the current location.
        When reaching new node, flush the memory
        Input:
        - topometric_pose: node_id, edge_idx, dist
        - guidance from previous node
        - img
        Output:
        - visual memory: publish with message type VisMem
        Dependency:
        - Localizer
        """
        if random_action: # For test
            test_act = np.random.randint(0, self.action_dim)
            onehot_test_act = np.zeros(self.action_dim)
            onehot_test_act[test_act] = 1
            tensor_action = torch.tensor(onehot_test_act, dtype=torch.float32)
            tensor_img = eVM_utils.img_transform(img).unsqueeze(0)
            if self.cuda:
                tensor_img = tensor_img.cuda()
                tensor_action = tensor_action.cuda()
            with torch.no_grad():
                feat = self.encode_im(tensor_img)
            self.list2encode.append([feat, tensor_action])

        else:
            # flush when reaching new node
            if flush is True: # topometric_pose.dist < 0.1:
                self.list2encode = []
                #self.vis_mem = []

            elif exp_active is False:
                action = np.zeros(self.action_dim)
                action[guidance-1] = 1
                tensor_img = eVM_utils.img_transform(img).unsqueeze(0)
                tensor_action = torch.tensor(action, dtype=torch.float32)
                if self.cuda:
                    tensor_action = tensor_action.cuda()
                    tensor_img = tensor_img.cuda()

                with torch.no_grad():
                    feat = self.encode_im(tensor_img.cuda())
                self.list2encode.append([feat, tensor_action])

        self.path_length = len(self.list2encode)

    def calcRecoveryGuidance(self, img=None):
        """
        Recovery Guide Provider Submodule:
        A module that guides a robot to return to the previous node, when StateDeterminant module
        determines that the robot is lost.
        If matching the retrieved image with the visual memory fails, call Exploration Guidance Module instead.
        Input:
        # - state: state from StateDeterminant Module
        - img: current image input (doesn't need if visual memory contains the current input image)
        Output:
        - action(s) guides to reach previous POI
        """

        if self.enable_exploration is False:
            self.enable_recovery = True

        if self.enable_recovery is True:
            # calculate the actions to return to the starting point of visual memory # doesn't need pose
            if img is None and self.path_length > 0:
                img = self.list2encode[-1][0]
            elif img is None:
                print('Nothing to encode or calculate because there was no input at all')
                raise Exception
            # encode the input image
            tensor_img = eVM_utils.img_transform(img).unsqueeze(0)
            if self.cuda:
                tensor_img = tensor_img.cuda()
            with torch.no_grad():
                img_feature = self.encode_im(tensor_img)
            #_, img_feature = self.vis_mem_encoder(tensor_img, tensor_action)

            if len(self.vis_mem) == 0:
                positions = [np.array([0., 0.])]
                rotations = [np.array([0.])]

                for ts in range(self.path_length-1):
                    guid = self.list2encode[ts][1]
                    if guid[0] == 1:
                        pos = positions[ts] + 0.4 * np.concatenate([np.cos(rotations[ts]), np.sin(rotations[ts])])
                        rot = rotations[ts]
                    if guid[1] == 1:
                        pos = positions[ts]
                        rot = rotations[ts] + np.pi/6 * (-1. if self.path_direction == 'homing' else 1.)
                    if guid[2] == 1:
                        pos = positions[ts]
                        rot = rotations[ts] - np.pi/6 * (-1. if self.path_direction == 'homing' else 1.)
                    positions.append(pos)
                    rotations.append(rot)

                positions = torch.from_numpy(np.stack(positions))
                rotations = torch.from_numpy(np.stack(rotations))

                if self.cuda:
                    positions = positions.cuda().float()
                    rotations = rotations.cuda().float()

                rel_pos = torch.unsqueeze(positions, 0) - torch.unsqueeze(positions, 1)
                rel_pos_dist = torch.norm(rel_pos, dim=-1)

                print(rel_pos.shape, rotations.shape)
                rel_pos_theta = torch.atan2(rel_pos[...,1], rel_pos[...,0]) - rotations.repeat(1,self.path_length)
                rel_pos_theta_cos = torch.cos(rel_pos_theta)
                rel_pos_theta_sin = torch.sin(rel_pos_theta)
                rel_pos_dt = torch.stack([rel_pos_dist, rel_pos_theta_cos, rel_pos_theta_sin], -1)

                rel_rot = rotations.permute(1, 0) - rotations + (np.pi if self.path_direction == 'homing' else 0.)
                rel_rot_cos = torch.cos(rel_pos_theta)
                rel_rot_sin = torch.sin(rel_pos_theta)
                rel_orn = torch.stack([rel_rot_cos, rel_rot_sin], -1)

                self.rel_pose = torch.cat([rel_pos_dt, rel_orn], -1)

                for ts in range(self.path_length):
                    vis_mem_seg, _ = self.vis_mem_encoder(self.list2encode[ts][1], [self.list2encode[t][0] for t in range(self.path_length)], self.rel_pose[ts,...])
                    self.vis_mem.append(vis_mem_seg)

            with torch.no_grad():
                action = self.recovery_policy(self.vis_mem, img_feature)
            self.recovery_guidance = action
        """
        if info is False:
            self.enable_recovery, self.enable_exploration = False, True

        if done is True:
            self.enable_recovery, self.enable_exploration = True, False
        """

    def calcExplorationGuidance(self, img):
        """
        Exploration Guidance Provider Submodule:
        A module that guides the robot to reach nearby POI using nearby visual information and visual memory.
        This module is triggered when the POI Recovery Guidance Module judges that it is hard to
        return to a node associated with the previous POI, based on the visual memory matching result.
        Input:
        # - state: state from StateDeterminant Module
        ---(topometric_pose_conf: confidence of topometric pose)
        - (tentative) POI detection result
        - img: curreunt image input
        - visual memory
        Output:
        - action(s) guides to reach nearby POI
        Dependency:
        - Localizer module
        """
        # if state == 'normal':
        #     self.enable_exploration = False

        if self.enable_exploration is True:
            try:
                actions, done, info = self.exploration_policy(img, torch.cat(self.vis_mem))
            except:
                print("NotImplementedError")
                actions, done, info = [0., 0.40, 0.], False, False # ['forward','forward','forward']

            self.exploration_guidance = actions
            if done is True:
                self.enable_recovery, self.enable_exploration = True, False
        else:
            self.exploration_guidance = None

    # def calcNeedForOptimalViewpointGuidance(self, topometric_pose, poi_conf, entrance, entrance_conf, poi_successes):
    #     """
    #     A module calculates the need for optimal viewpoint guidance.
    #     It considers the results of POI detection and current location (destination, door).

    #     Input:
    #     - topometric_pose: node_id, edge-idx, dist
    #     - poi_conf: detection results of POI (confidence score)
    #     - (tentative) entrance: If there is an building entrance near current location, True, else False (Boolean) , tentative due to the API ambiguity.
    #     - entrance_conf: detection results of building entrances
    #     - poi_successes: A sequence of POI detection result (confidence score)

    #     Output:
    #     - require the optimal viewpoint guidance optimization (Boolean)
    #     """

    #     # Cases
    #     case1 = np.mean(poi_successes) < 0.5
    #     case2 = ((self.map.getNode(topometric_pose.node_id).edges[topometric_pose.edge_idx].length - topometric_pose.dist) < 0.2) and (poi_conf < 0.5)
    #     case3 = entrance and (entrance_conf < 0.5)

    #     if case1 or case2 or case3:
    #         self.enable_ove = True
    #     else:
    #         self.enable_ove = False

    # def calcOptimalViewpointGuidance(self, img_path, target_poi):
    #     """
    #     Optimal Viewpoint Guidance Provider Submodule:
    #     A module that provides an optimal viewpoint guidance for enhancing POI detection.
    #     Input:
    #     - img: current image input
    #     - target_poi: target PoI (searching object)
    #     Output:
    #     - Action(s) guides to find an optimal viewpoint
    #     """

    #     if self.enable_ove:
    #         self.ov_guidance = self.optimal_viewpoint(img_path, target_poi)

    def calcOptimalViewpointGuidance(self, obs_rgb, targ_rgb):
        """
        Optimal Viewpoint Guidance Provider Submodule:
        A module that provides an optimal viewpoint guidance for enhancing POI detection.
        Input:
        - obs_rgb, targ_rgb: current and target image input, PIL Image
        Output:
        - Action(s) guides to find an optimal viewpoint
        """
        print(obs_rgb.size, targ_rgb.size)
        # Crop center
        obs_crop_size = min(obs_rgb.size)
        obs_crop = transforms.CenterCrop(obs_crop_size)
        obs_rgb = obs_crop(obs_rgb).resize((256,256))

        targ_crop_size = min(targ_rgb.size)
        targ_crop = transforms.CenterCrop(targ_crop_size)
        targ_rgb = targ_crop(targ_rgb).resize((256,256))

        observation={}
        observation["targ_rgb"] = np.asarray(targ_rgb)
        pred_targ = self.ove_extractor.inference(observation["targ_rgb"], num_keypoints=self.num_keypoints)
        observation["targ_desc_glob"], observation["targ_desc"], observation["targ_det"], observation["targ_score"] = \
            torch.tensor(pred_targ['global_descriptor'][None]), torch.tensor(pred_targ['local_descriptors'][None]), \
            torch.tensor(pred_targ['keypoints'].astype(np.float32)[None]), torch.tensor(pred_targ['scores'][None])
        observation["rgb"]= np.asarray(obs_rgb)
        pred_obs = self.ove_extractor.inference(observation["rgb"], num_keypoints=self.num_keypoints)
        observation["obs_desc_glob"], observation["obs_desc"], observation["obs_det"], observation["obs_score"] = \
                torch.tensor(pred_obs['global_descriptor'][None]), torch.tensor(pred_obs['local_descriptors'][None]), \
                torch.tensor(pred_obs['keypoints'].astype(np.float32)[None]), torch.tensor(pred_obs['scores'][None])
        observation["rgb"]= torch.tensor(np.asarray(obs_rgb)[None])

        logit, n_matching = self.ove_policy.act(observation, debug=False)
        # print(n_matching)
        ind = logit.argmax().numpy()
        if ind == 0:
            action = "Move forward 0.25m"
            self.ov_guidance = [0., 0.25, 0.]
        elif ind == 1:
            action = "Turn left 10 degree"
            self.ov_guidance = [-10., 0., 0.]
        else:
            action = "Turn right 10 degree"
            self.ov_guidance = [10., 0., 0.]
        if n_matching[0] >= 15:
            action = "stop"
            self.ov_guidance = [0., 0., 0.]

    # def optimal_viewpoint(self, file_path=None, target_poi=None):
    #     if file_path and target_poi is None:
    #         file_path, target_poi = self.im_paths[0], self.target_pois[0]
    #     heading1 = D1 = heading2 = 0
    #     templates, main_template, opt_ratio = file_utils.get_templates(self.ove_data_folder, targetPOI=target_poi)
    #     img = get_img(self.ove_data_folder, file_path)
    #     depth_ = get_depth(self.ove_data_folder, file_path)
    #     bbox = get_bbox(self.ove_data_folder, file_path)
    #     h, w, _ = img.shape

    #     if len(bbox) > 0:
    #         bbs = []
    #         for bb in bbox:
    #             bbs.append(bb.rstrip().split(','))
    #         bbs = np.stack([np.float32(bbs[i]) for i in range(len(bbs))])
    #         bbs = np.reshape(bbs, [-1, 4, 2])
    #         bbs = bbs.astype(np.int32)

    #         # Find the bounding box for the target POI
    #         template_matched, bbox, index, _ = template_matching(img, bbs, templates, main_template)
    #         bbox = np.reshape(bbox, [-1, 4, 2])
    #         mask = make_mask(bbox, shape=[h, w])
    #         if np.sum(mask) > 0 and template_matched:
    #             depth = depth_[mask == 1]
    #             center_depth = np.mean(depth_[110:140, 235:265], (0,1))

    #             if np.mean(depth) <= 0.01 and len(bbs) > 1:
    #                 indices = list(np.arange(len(bbs)))
    #                 indices.pop(index)
    #                 bbs = bbs[indices]
    #                 bbs = np.reshape(bbs, [-1, 4, 2])
    #                 template_matched, bbox, _, _ = template_matching(img, bbs, templates, main_template)
    #                 bbox = np.reshape(bbox, [-1, 4, 2])
    #                 mask = make_mask(bbox, shape=[h, w])
    #                 depth = depth_[mask == 1]

    #             if np.mean(depth) <= 0.01:
    #                 depth = depth + 0.01

    #             if np.mean(center_depth) <= 0.01:
    #                 center_depth += 0.01

    #             # TODO: Estimate the exact distance
    #             D = np.mean(depth) * 19.2  # d1
    #             center_D = np.mean(center_depth) * 19.2  # d2

    #             # Decide the amount of the movement using depth
    #             ratio = (abs(bbox[0, 3, 1] - bbox[0, 0, 1]) + abs(bbox[0, 1, 1] - bbox[0, 2, 1])) / 2 / h

    #             # Decide the moving direction
    #             sf = get_surfacenormal(self.ove_data_folder, file_path)
    #             sf_norm = np.mean(sf[mask == 1], 0)
    #             sf_norm = sf_norm * 2 - 1
    #             sf_norm = sf_norm / np.linalg.norm(sf_norm, 2)

    #             POI_centloc = ("left", "right")[(bbox[0, 2, 0] + bbox[0, 0, 0]) / 2 > w / 2] #Left: Logos on the left from center, Right: Logos on the right from center

    #             # Rotate the agent until the POI is locate on the center
    #             center_sf_norm = np.mean(sf[110:140, 235:265], (0, 1))
    #             center_sf_norm = center_sf_norm * 2 - 1
    #             center_sf_norm = center_sf_norm / np.linalg.norm(center_sf_norm, 2)

    #             # Check that the bounding box is on the left or right buildings (not street or sky)
    #             if abs(sf_norm[1]) < 0.8:
    #                 center_poi_theta = np.arccos(np.dot(sf_norm, center_sf_norm))

    #                 # Align the POI and the camera center
    #                 heading = 180 / np.pi * center_poi_theta
    #                 if POI_centloc == "left":
    #                     heading = -heading

    #                 if abs(heading) > 15:
    #                     return [0, 0, heading]

    #             theta_ = np.arccos(np.dot(center_sf_norm, normal_vector)) # center point
    #             theta_tilde = np.arctan(D*np.tan(theta_)/np.abs(center_D-D))
    #             cond = (center_D > D)
    #             theta = (theta_ + theta_tilde - np.pi/2) if cond else (theta_ + np.pi/2 - theta_tilde)
    #             D = D / np.sin(theta_tilde) # distance between the robot and signage, not depth of the signage
    #             D0 = D * (1 - np.maximum(ratio / opt_ratio, 0.95))
    #             thetad = np.arctan(((D - D0) * np.sin(theta)) / (D - (D - D0) * np.cos(theta)))

    #             #Rotate before going straight
    #             if cond:
    #                 heading1 = 180 / np.pi * (thetad - (np.pi/2 - theta_tilde))
    #                 heading1 = (heading1, -heading1)[POI_centloc == "right"]
    #             else:
    #                 heading1 = 180 / np.pi * (thetad + (np.pi/2 - theta_tilde))
    #                 heading1 = (-heading1, heading1)[POI_centloc == "right"]
    #             D1 = ((D - D0)*np.sin(theta)/(np.sin(thetad)), D - D0)[thetad == 0] # Sine Law

    #             # Rotate to see the POI
    #             if cond:
    #                 heading2 = 180 / np.pi * (theta + thetad)
    #                 heading2 = (-abs(heading2), abs(heading2))[POI_centloc == "right"]
    #             else:
    #                 heading2 = 180 / np.pi * (theta + thetad)
    #                 heading2 = (abs(heading2), -abs(heading2))[POI_centloc == "right"]

    #     return [heading1, D1, heading2]

    def getVisualMemory(self):
        return self.vis_mem

    def isRecoveryGuidanceEnabled(self):
        return self.enable_recovery

    def isExplorationGuidanceEnabled(self):
        return self.enable_exploration

    def isOptimalViewpointGuidanceEnabled(self):
        return self.enable_ove

    def getExplorationGuidance(self, img, guidance, flush, exp_active, ove, im_path=None, target_poi=None):
        # import pdb
        # pdb.set_trace()
        # print(flush, exp_active, guidance)
        img_orig = copy.deepcopy(img)#cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        curr_img_orig = Image.fromarray(img_orig)
        img = cv2.resize(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), (256,256))
        curr_img = Image.fromarray(img)
        # curr_img = np.array(curr_img)
        # curr_img = torch.from_numpy(curr_img).cuda().float()
        # print(curr_img.shape)
        self.encodeVisualMemory(curr_img, guidance, flush=flush, exp_active=exp_active)

        if exp_active is True:
            if self.isOptimalViewpointGuidanceEnabled() and ove:
                if os.path.isfile('./data_exp/optimal_viewpoint/target_poi.jpg'):
                    target_poi = Image.open('./data_exp/optimal_viewpoint/target_poi.jpg')
                if target_poi is None:
                    return [[0., 0., 0.]], 'OptimalViewpoint'
                else:
                    self.calcOptimalViewpointGuidance(curr_img_orig, target_poi) # (im_path, target_poi)
                    guidance = self.ov_guidance
                    # print("Optimal Guidance [theta1, d, theta2]: [%.2f degree, %.2fm, %.2f degree]" % (guidance[0], guidance[1], guidance[2]))
                    return [guidance], 'OptimalViewpoint'

            if self.isRecoveryGuidanceEnabled():
                self.calcRecoveryGuidance(img=curr_img)
                print('Recovery guidance from the last inserted visual memory : ', self.recovery_guidance)

                if self.enable_exploration is False:
                    return [self.recovery_guidance], 'Recovery'

            if self.isExplorationGuidanceEnabled():
                self.calcExplorationGuidance(img=curr_img)
                print('Exploration guidance from the last inserted visual memory : ', self.exploration_guidance)
                return [self.exploration_guidance], 'Exploration'
        else:
            return [[0., 0., 0.]], 'Normal'

    # def getRecoveryGuidance(self):
    #     if self.enable_recovery:
    #         return self.recovery_guidance
    #     else:
    #         return None

    # def getExplorationGuidance(self):
    #     if self.enable_exploration:
    #         return self.exploration_guidance
    #     else:
    #         return None

    # def getOptimalViewpointGuidance(self):
    #     if self.enable_ove:
    #         return self.ov_guidance
    #     else:
    #         return None

# if __name__ == "__main__":
#   anm = ActiveNavigationModule()