# SPDX-License-Identifier: Apache-2.0

"""
The following code compares the speed of tensorflow against onnxruntime
with a model downloaded from Tensorflow Hub.
"""
import os
import sys
import time
import tarfile
import zipfile
import subprocess
import datetime
from collections import OrderedDict
import numpy
from tqdm import tqdm
import onnxruntime

def generate_random_images(shape=(1, 100, 100, 3), n=10, dtype=numpy.float32, scale=255):
    imgs = []
    for i in range(n):
        sh = shape
        img = numpy.clip(numpy.abs(numpy.random.randn(*sh)), 0, 1) * scale
        img = img.astype(dtype)
        imgs.append(img)
    return imgs

def generate_text_inputs():
    """
    preprocessor = hub.load("http://tfhub.dev/tensorflow/albert_en_preprocess/3")
    encoder = hub.load("https://tfhub.dev/tensorflow/albert_en_xlarge/3")
    sentences = tf.constant(["Hi I'm some text"])
    embedded_inputs = {k: v.numpy() for k, v in preprocessor(sentences).items()}
    """
    one = OrderedDict([
        ('input_word_ids', numpy.array([[
            2, 4148, 31, 22, 79, 109, 1854, 3, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0,0]]).reshape((1, -1))),
        ('input_type_ids', numpy.array([[
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).reshape((1, -1))),
        ('input_mask', numpy.array([[
            1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
            0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]).reshape((1, -1)))])
    return [one for i in range(10)]

def measure_time(fct, imgs, n=50, timeout=15):
    """
    Runs *n* times the same function taking one parameter
    from *imgs*. It stops if the total time overcomes *timeout*.
    It also runs once the function before measuring.
    """
    # Let's run it once first.
    fct(imgs[0])
    # The time is measured for n iterations except if the total time
    # overcomes timeout.
    results = []
    times = []
    for i in tqdm(range(0, n)):
        img = imgs[i % len(imgs)]
        begin = time.perf_counter()
        result = fct(img)
        end = time.perf_counter()
        results.append(result)
        times.append(end - begin)
        if sum(times) > timeout:
            break
    return results, times

def download_model(url, dest, verbose=True):
    """
    Downloads a model from tfhub and unzips it.
    The function assumes the format is `.tar.gz`.
    """
    if not os.path.exists(dest):
        os.makedirs(dest)
    fpath = os.path.join(dest, "model.tar.gz")
    if not os.path.exists(fpath):
        from tf2onnx import utils
        if verbose:
            print("Download %r." % fpath)
        utils.get_url(url, fpath)
    tname = os.path.join(dest, "model_path")
    if not os.path.exists(tname):
        if verbose:
            print("Untar %r." % tname)
        tar = tarfile.open(fpath)
        tar.extractall(tname)
        tar.close()
    return fpath, tname

def download_tflite(url, dest, verbose=True):
    """
    Downloads a model from tfhub.
    The function assumes the format is `.tflite`.
    """
    if not os.path.exists(dest):
        os.makedirs(dest)
    fpath = os.path.join(dest, "model.tflite")
    if not os.path.exists(fpath):
        from tf2onnx import utils
        if verbose:
            print("Download %r." % fpath)
        utils.get_url(url, fpath)
    return fpath

def convert_model(model_name, output_path, opset=13, tag=None, verbose=True):
    """
    Converts the downloaded model into ONNX.
    """
    ext = os.path.splitext(output_path)[-1]
    large_model = ext == ".zip"
    if not os.path.exists(output_path):
        begin = datetime.datetime.now()
        cmdl = ['-m', 'tf2onnx.convert', '--saved-model',
                '"%s"' % os.path.abspath(model_name).replace("\\", "/"),
                '--output', '"%s"' % os.path.abspath(output_path).replace("\\", "/"),
                '--opset', "%d" % opset]
        if tag is not None:
            cmdl.append('--tag="%s"' % tag)
        if large_model:
            cmdl.append('--large_model')
        if verbose:
            print("cmd: python %s" % " ".join(cmdl))
        pproc = subprocess.Popen(
            cmdl, shell=True, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            executable=sys.executable.replace("pythonw", "python"))
        stdoutdata, stderrdata = pproc.communicate()
        if verbose:
            print('--OUT--')
            print(stdoutdata.decode('ascii'))
            print('--ERR--')
            print(stderrdata.decode('ascii'))
            print("Duration %r." % (datetime.datetime.now() - begin))

def convert_tflite(model_name, output_path, opset=13, verbose=True):
    """
    Converts the downloaded model into ONNX.
    """
    if not os.path.exists(output_path):
        begin = datetime.datetime.now()
        cmdl = ['-m', 'tf2onnx.convert', '--tflite',
                '"%s"' % os.path.abspath(model_name).replace("\\", "/"),
                '--output', '"%s"' % os.path.abspath(output_path).replace("\\", "/"),
                '--opset', "%d" % opset]
        if verbose:
            print("cmd: python %s" % " ".join(cmdl))
        pproc = subprocess.Popen(
            cmdl, shell=True, stdin=None, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            executable=sys.executable.replace("pythonw", "python"))
        stdoutdata, stderrdata = pproc.communicate()
        if verbose:
            print('--OUT--')
            print(stdoutdata.decode('ascii'))
            print('--ERR--')
            print(stderrdata.decode('ascii'))
            print("Duration %r." % (datetime.datetime.now() - begin))

def check_discrepencies(out1, out2, threshold=1e-3):
    """
    Compares two tensors. Raises an exception if it fails.
    """
    if out1.dtype != out2.dtype:
        raise AssertionError("Type mismatch %r != %r." % (out1.dtype, out2.dtype))
    if out1.shape != out2.shape:
        raise AssertionError("Shape mismatch %r != %r." % (out1.shape, out2.shape))
    diff = numpy.abs(out1.ravel() - out2.ravel()).max()
    if diff > threshold:
        raise AssertionError("Discrependcies %r > %r." % (diff, threshold))

def benchmark(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3,
              signature=None, tag=None, output_name=None, ort_name=None,
              optimize=True):
    """
    Runs a simple benchmark.
    Goes through every steps (download, convert).
    Skips them if already done.
    """
    fpath, tname = download_model(url, dest)
    if verbose:
        print("Created %r, %r." % (fpath, tname))

    # Converts the model.
    if verbose:
        print("Convert model in %r." % dest)
    convert_model(tname, onnx_name, opset, tag=tag)
    if verbose:
        print("Created %r." % onnx_name)

    # unzip large_model
    ext = os.path.splitext(onnx_name)[-1]
    if ext == ".zip":
        onnx_name_unzipped = os.path.join(dest, "large_model", "__MODEL_PROTO.onnx")
        if not os.path.exists(onnx_name_unzipped):
            if verbose:
                print("Unzip model in %r." % os.path.join(dest, "large_model"))
            with zipfile.ZipFile(onnx_name, 'r') as z:
              z.extractall(os.path.join(dest, "large_model"))
        onnx_name = onnx_name_unzipped

    # Benchmarks both models.
    if optimize:
        ort = onnxruntime.InferenceSession(onnx_name)
    else:
        so = onnxruntime.SessionOptions()
        so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
        ort = onnxruntime.InferenceSession(onnx_name)

    if verbose:
        print("ONNX inputs:")
        for a in ort.get_inputs():
            print("  {}: {}, {}".format(a.name, a.type, a.shape))
        print("ONNX outputs:")
        for a in ort.get_outputs():
            print("  {}: {}, {}".format(a.name, a.type, a.shape))

    # onnxruntime
    if output_name is None or ort_name is None:
        index = 0
    else:
        output_names = [o.name for o in ort.get_outputs()]
        if output_name in output_names:
            index = output_names.index(output_name)
        elif ort_name in output_names:
            index = output_names.index(ort_name)
        else:
            index = 0
    if isinstance(imgs[0], dict):
        fct_ort = lambda img: ort.run(None, img)[index]
    else:
        input_name = ort.get_inputs()[0].name
        fct_ort = lambda img: ort.run(None, {input_name: img})[index]
    results_ort, duration_ort = measure_time(fct_ort, imgs)
    if verbose:
        print("ORT", len(imgs), duration_ort)

    # tensorflow
    import tensorflow_hub as hub
    from tensorflow import convert_to_tensor
    if isinstance(imgs[0], OrderedDict):
        imgs_tf = [
            OrderedDict((k, convert_to_tensor(v)) for k, v in img.items())
            for img in imgs]
    else:
        imgs_tf = [convert_to_tensor(img) for img in imgs]
    model = hub.load(url.split("?")[0])
    if signature is not None:
        model = model.signatures[signature]
    results_tf, duration_tf = measure_time(model, imgs_tf)

    if verbose:
        print("TF", len(imgs), duration_tf)
        mean_ort = sum(duration_ort) / len(duration_ort)
        mean_tf = sum(duration_tf) / len(duration_tf)
        print("ratio ORT=%r / TF=%r = %r" % (mean_ort, mean_tf, mean_ort / mean_tf))

    # checks discrepencies
    res = model(imgs_tf[0])
    if isinstance(res, dict):
        if output_name is None:
            if len(res) != 1:
                raise NotImplementedError(
                    "TF output contains more than one output=%r and output names=%r." % (
                        list(res), [o.name for o in ort.get_outputs()]))
            else:
                output_name = ort.get_outputs()[0].name
        if output_name not in res:
            raise AssertionError("Unable to find output %r in %r." % (output_name, list(sorted(res))))
        res = res[output_name]
    try:
        check_discrepencies(fct_ort(imgs[0]), res.numpy(), threshold)
    except AttributeError as e:
        raise AssertionError(
            "Unable to check discrepencies for res=%r." % res) from e
    except AssertionError as e:
        output_names = [o.name for o in ort.get_outputs()]
        res = ort.run(None, imgs[0])
        for i, r in enumerate(res):
            print("ORT %d: %s: %r: %r" % (i, output_names[i], r.dtype, r.shape))
        raise e
    return duration_ort, duration_tf

def benchmark_tflite(url, dest, onnx_name, opset, imgs, verbose=True, threshold=1e-3):
    """
    Runs a simple benchmark with a tflite model.
    Goes through every steps (download, convert).
    Skips them if already done.
    """
    tname = download_tflite(url, dest)
    if verbose:
        print("Created %r." % tname)

    # Converts the model.
    if verbose:
        print("Convert model in %r." % dest)
    convert_tflite(tname, onnx_name, opset)
    if verbose:
        print("Created %r." % onnx_name)

    # Benchmarks both models.
    ort = onnxruntime.InferenceSession(onnx_name)

    if verbose:
        print("ONNX inputs:")
        for a in ort.get_inputs():
            print("  {}: {}, {}".format(a.name, a.type, a.shape))
        print("ONNX outputs:")
        for a in ort.get_outputs():
            print("  {}: {}, {}".format(a.name, a.type, a.shape))

    # onnxruntime
    input_name = ort.get_inputs()[0].name
    fct_ort = lambda img: ort.run(None, {input_name: img})[0]
    results_ort, duration_ort = measure_time(fct_ort, imgs)
    if verbose:
        print("ORT", len(imgs), duration_ort)

    # tensorflow
    import tensorflow_hub as hub
    from tensorflow import convert_to_tensor
    if isinstance(imgs[0], OrderedDict):
        imgs_tf = [
            OrderedDict((k, convert_to_tensor(v)) for k, v in img.items())
            for img in imgs]
    else:
        imgs_tf = [convert_to_tensor(img) for img in imgs]
    model = hub.load(url.split("?")[0])
    if signature is not None:
        model = model.signatures['serving_default']
    results_tf, duration_tf = measure_time(model, imgs_tf)

    if verbose:
        print("TF", len(imgs), duration_tf)
        mean_ort = sum(duration_ort) / len(duration_ort)
        mean_tf = sum(duration_tf) / len(duration_tf)
        print("ratio ORT=%r / TF=%r = %r" % (mean_ort, mean_tf, mean_ort / mean_tf))

    # checks discrepencies
    res = model(imgs_tf[0])
    if isinstance(res, dict):
        if len(res) != 1:
            raise NotImplementedError("TF output contains more than one output: %r." % res)
        output_name = ort.get_outputs()[0].name
        if output_name not in res:
            raise AssertionError("Unable to find output %r in %r." % (output_name, list(sorted(res))))
        res = res[output_name]
    check_discrepencies(fct_ort(imgs[0]), res.numpy(), threshold)
    return duration_ort, duration_tf


from CPSC2019_challenge import CPSC2019_challenge
import numpy as np
np.set_printoptions(threshold=np.inf)
import os
import re
from scipy import io

'''packages'''
from keras.models import load_model
from plain_data_make import PlainPreprocessor
from plain_model import hyper_params
from sig_tool import normalize, diff
import time
import pickle
from plain_model import  gen_pp_data_dir_name
import matplotlib.pyplot as plt

def load_ans(data_path_, rpos_path_, fs_):
    '''
    Please modify this function when you have to load model or any other parameters in CPSC2019_challenge()
    '''
    models = []

    # model = load_model('model_9705.h5')
    # models.append(model)
    #
    # model = load_model('model_9647.h5')
    # models.append(model)
    # model = load_model('model_9657.h5')
    # models.append(model)
    #
    # model = load_model('models/rematch_ckpt_plain_rev2_40_36874_0_044_0.0602_0.0994_0.9807_0.9670.h5') #conv-lstm
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev2_40_36973_0_060_0.0676_0.1113_0.9789_0.9635.h5') #U-net FCN
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev2_40_43632_0_347_0.0892_0.0877_0.9654_0.9665.h5') #lstm-conv
    # models.append(model)

    '''ver.2.3'''
    # model = load_model('models/rematch_ckpt_plain_rev4_40_19031_0_031_0.0587_0.0711_0.9763_0.9722.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_40_92098_0_047_0.0654_0.0733_0.9740_0.9720.h5')
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_40_11085_0_037_0.0669_0.0644_0.9729_0.9741.h5')
    # models.append(model)

    '''ver.2.4'''
    # model = load_model('models/rematch_ckpt_plain_rev4_30_69622_0_002_0.0498_0.0590_0.9793_0.9770.h5') # U-net LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_52643_0_001_0.0562_0.0622_0.9772_0.9746.h5') # stacked LSTM
    # models.append(model)
    # model = load_model('models/rematch_ckpt_plain_rev4_30_17321_2_002_0.0517_0.0608_0.9799_0.9768.h5') # U-net++ LSTM
    # models.append(model)

    '''ver.3 '''
    # model = load_model('models/tune_fold10_30_netlstm_sig1_89383_4_002_0.0553_0.0452_0.9776_0.9811.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_netlstm_sig1_99576_3_007_0.0579_0.0484_0.9765_0.9801.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_netlstm_sig1_59925_7_003_0.0641_0.0546_0.9744_0.9781.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_netlstm_sig1_64082_1_007_0.0628_0.0658_0.9746_0.9732.h5')
    # models.append(model)

    '''ver.4'''
    # model = load_model('models/tune_fold10_30_transunetlstm_sig1_54232_4_001_0.0560_0.0513_0.9773_0.9795.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_transunetlstm_sig1_22794_5_005_0.0580_0.0548_0.9767_0.9777.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_transunetlstm_sig1_38928_6_001_0.0594_0.0532_0.9760_0.9784.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_transunetlstm_sig1_26234_0_002_0.0639_0.0591_0.9742_0.9761.h5')
    # models.append(model)
    # model = load_model('models/tune_fold10_30_transunetlstm_sig1_37497_1_004_0.0585_0.0562_0.9764_0.9773.h5')
    # models.append(model)

    '''ver.5'''
    model = load_model('models/shuffle_30_aconvunet_sig1_68742_0_246_0.0496_0.0784_0.9967_0.9875.h5')
    models.append(model)

    preprocessor = PlainPreprocessor(hyper_params)

    def is_mat(l):
        return l.endswith('.mat')
    ecg_files = list(filter(is_mat, os.listdir(data_path_)))
    rpos_files = list(filter(is_mat, os.listdir(rpos_path_)))
    HR_ref = []
    R_ref = []
    HR_ans = []
    R_ans = []
    cnt = 0
    for rpos_file in rpos_files:
        index = re.split('[_.]', rpos_file)[1]
        ecg_file = 'data_' + index + '.mat'

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

        ecg_data = np.transpose(io.loadmat(ecg_path)['ecg'])[0]
        r_ref = io.loadmat(ref_path)['R_peak'].flatten()
        r_ref = r_ref[(r_ref >= 0.5*fs_) & (r_ref <= 9.5*fs_)]

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

        print('estimating ', ecg_path)
        ecg_data = preprocessor.single_preprocess(ecg_data)

        hr_ans, r_ans = CPSC2019_challenge(ecg_data, models)
        r_ans = np.array(r_ans)
        print('finished ', ecg_path)

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

    return R_ref, HR_ref, R_ans, HR_ans

def score(r_ref, hr_ref, r_ans, hr_ans, fs_, thr_):
    HR_score = 0
    record_flags = np.ones(len(r_ref))
    for i in range(len(r_ref)):
        FN = 0
        FP = 0
        TP = 0
        hr_der = abs(int(hr_ans[i]) - int(hr_ref[i]))
        if hr_der <= 0.02 * hr_ref[i]:
            HR_score = HR_score + 1
        elif hr_der <= 0.05 * hr_ref[i]:
            HR_score = HR_score + 0.75
        elif hr_der <= 0.1 * hr_ref[i]:
            HR_score = HR_score + 0.5
        elif hr_der <= 0.2 * hr_ref[i]:
            HR_score = HR_score + 0.25

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

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

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

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

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

    return rec_acc, hr_acc

if __name__ == '__main__':
    FS = 500
    THR = 0.075

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

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

    print('rec_acc:', rec_acc, 'hr_acc:', hr_acc)


import logging as log
import os
import pathlib
import cv2
import numpy as np
from openvino.inference_engine import IENetwork, IECore
import torch
import torchvision
from .tracker import update_tracker
from .StatusDetector import FaceDet

def xywh2xyxy(x):

    y = torch.zeros_like(x) if isinstance(
        x, torch.Tensor) else np.zeros_like(x)
    y[:, 0] = x[:, 0] - x[:, 2] / 2  # top left x
    y[:, 1] = x[:, 1] - x[:, 3] / 2  # top left y
    y[:, 2] = x[:, 0] + x[:, 2] / 2  # bottom right x
    y[:, 3] = x[:, 1] + x[:, 3] / 2  # bottom right y

    return y

def non_max_suppression(prediction, conf_thres=0.05, iou_thres=0.4):

    prediction = torch.from_numpy(prediction)
    if prediction.dtype is torch.float16:
        prediction = prediction.float()  # to FP32

    nc = prediction[0].shape[1] - 5  # number of classes
    xc = prediction[..., 4] > conf_thres  # candidates

    min_wh, max_wh = 2, 4096
    max_det = 300  # maximum number of detections per image
    multi_label = nc > 1  # multiple labels per box (adds 0.5ms/img)

    output = [None] * prediction.shape[0]
    for xi, x in enumerate(prediction):  # image index, image inference
        # Apply constraints
        # x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0  # width-height
        x = x[xc[xi]]  # confidence

        # If none remain process next image
        if not x.shape[0]:
            continue

        # Compute conf
        x[:, 5:] *= x[:, 4:5]  # conf = obj_conf * cls_conf

        # Box (center x, center y, width, height) to (x1, y1, x2, y2)
        box = xywh2xyxy(x[:, :4])

        if multi_label:
            i, j = (x[:, 5:] > conf_thres).nonzero(as_tuple=False).T
            x = torch.cat((box[i], x[i, j + 5, None], j[:, None].float()), 1)
        else:  # best class only
            conf, j = x[:, 5:].max(1, keepdim=True)
            x = torch.cat((box, conf, j.float()), 1)[
                conf.view(-1) > conf_thres]

        n = x.shape[0]  # number of boxes
        if not n:
            continue
        # Batched NMS
        c = x[:, 5:6] * max_wh  # classes
        # boxes (offset by class), scores
        boxes, scores = x[:, :4] + c, x[:, 4]
        i = torchvision.ops.boxes.nms(boxes, scores, iou_thres)
        if i.shape[0] > max_det:  # limit detections
            i = i[:max_det]

        output[xi] = x[i]

    return output

class Detector(object):

    def __init__(self):

        self.device = 'GPU'
        self.threshold = 0.1
        self.stride = 10
        self.size = 640
        self.faceDet = FaceDet()
        log.basicConfig(level=log.DEBUG)

        # For objection detection task, replace your target labels here.
        self.label_id_map = ["face", "normal", "phone",
                             "write", "smoke", "eat", "camputer", "sleep"]
        self.model_xml = 'weights/openvo_fp16_yolov5/last.xml'
        self.net = self.init_model()
        self.input_blob = next(iter(self.net.inputs))
        self.out_blob = next(iter(self.net.outputs))

        self.build_config()

    def init_model(self):
        if not os.path.isfile(self.model_xml):
            log.error('model_xml does not exist')
            return None
        model_bin = pathlib.Path(self.model_xml).with_suffix('.bin').as_posix()
        net = IENetwork(model=self.model_xml, weights=model_bin)

        ie = IECore()
        self.exec_net = ie.load_network(network=net, device_name=self.device)
        input_blob = next(iter(net.inputs))
        n, c, h, w = net.inputs[input_blob].shape
        self.input_h, self.input_w = h, w
        self.input_c, self.input_n = c, n

        return net

    def build_config(self):

        self.faceTracker = {}
        self.faceClasses = {}
        self.faceLocation1 = {}
        self.faceLocation2 = {}
        self.frameCounter = 0
        self.currentCarID = 0
        self.recorded = []

        self.font = cv2.FONT_HERSHEY_SIMPLEX

    def preprocess(self, img):

        img0 = img.copy()
        img = cv2.resize(img, (self.size, self.size))
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)
        img = img.astype(np.float) / 255.0  # 图像归一化
        img = np.expand_dims(img, axis=0)

        return img0, img

    def detect(self, input_image):
        if not self.net or input_image is None:
            log.error('Invalid input args')
            return None
        ih, iw, _ = input_image.shape

        input_image, images = self.preprocess(input_image)
        res = self.exec_net.infer(inputs={self.input_blob: images})

        data = res[self.out_blob]

        data = non_max_suppression(data, 0.4, 0.5)
        detect_objs = []
        if data[0] == None:
            return detect_objs
        else:
            data = data[0].numpy()
            for proposal in data:
                if proposal[4] > self.threshold:
                    xmin = int(iw * (proposal[0]/self.size))
                    ymin = int(ih * (proposal[1]/self.size))
                    xmax = int(iw * (proposal[2]/self.size))
                    ymax = int(ih * (proposal[3]/self.size))
                    # xmin = int(iw * (proposal[0]))
                    # ymin = int(ih * (proposal[1]))
                    # xmax = int(iw * (proposal[2]))
                    # ymax = int(ih * (proposal[3]))
                    detect_objs.append((
                        xmin, ymin, xmax,
                        ymax, self.label_id_map[int(proposal[5])]
                    ))

            return input_image, detect_objs

    def feedCap(self, im, func_status):

        retDict = {
            'frame': None,
            'faces': None,
            'list_of_ids': None,
            'face_bboxes': []
        }
        self.frameCounter += 1
        with_headpose = func_status['headpose']

        im, faces, face_bboxes = update_tracker(self, im, with_headpose)

        retDict['frame'] = im
        retDict['faces'] = faces
        retDict['face_bboxes'] = face_bboxes

        return retDict

if __name__ == '__main__':

    predictor = Detector()
    base = 'images'

    for p in os.listdir(base):

        img = cv2.imread(os.path.join(base, p))
        result = predictor.feedCap(img)
        cv2.imshow('result', img)
        cv2.waitKey(0)


# A modification version from chainercv repository.
# (See https://github.com/chainer/chainercv/blob/master/chainercv/evaluations/eval_detection_voc.py)
from __future__ import division

import os
import numpy as np
import torch
from maskrcnn_benchmark.structures.bounding_box import BoxList
from maskrcnn_benchmark.data.datasets.evaluation.utils import evaluate_box_proposals

def do_vg_evaluation(dataset, predictions, output_folder, box_only, eval_attributes, logger, save_predictions=True):
    # TODO need to make the use_07_metric format available
    # for the user to choose
    # we use int for box_only. 0: False, 1: box for RPN, 2: box for object detection,
    if box_only:
        if box_only == 1:
            limits = [100, 1000]
        elif box_only == 2:
            limits = [36, 99]
        else:
            raise ValueError("box_only can be either 0/1/2, but get {0}".format(box_only))
        areas = {"all": "", "small": "s", "medium": "m", "large": "l"}
        result = {}
        for area, suffix in areas.items():
            for limit in limits:
                logger.info("Evaluating bbox proposals@{:d}".format(limit))
                stats = evaluate_box_proposals(
                    predictions, dataset, area=area, limit=limit
                )
                key_ar = "AR{}@{:d}".format(suffix, limit)
                key_num_pos = "num_pos{}@{:d}".format(suffix, limit)
                result[key_num_pos] = stats["num_pos"]
                result[key_ar] = stats["ar"].item()
                key_recalls = "Recalls{}@{:d}".format(suffix, limit)
                # result[key_recalls] = stats["recalls"]
                print(key_recalls, stats["recalls"])
                print(key_ar, "ar={:.4f}".format(result[key_ar]))
                print(key_num_pos, "num_pos={:d}".format(result[key_num_pos]))
        logger.info(result)
        logger.info(result)
        # check_expected_results(result, expected_results, expected_results_sigma_tol)
        if output_folder and save_predictions:
            if box_only == 1:
                torch.save(result, os.path.join(output_folder, "rpn_proposals.pth"))
            elif box_only == 2:
                torch.save(result, os.path.join(output_folder, "box_proposals.pth"))
            else:
                raise ValueError("box_only can be either 0/1/2, but get {0}".format(box_only))
        return {"box_proposal": result}

    pred_boxlists = []
    gt_boxlists = []
    for image_id, prediction in sorted(predictions.items()):
        img_info = dataset.get_img_info(image_id)
        if len(prediction) == 0:
            continue
        image_width = img_info["width"]
        image_height = img_info["height"]
        prediction = prediction.resize((image_width, image_height))
        pred_boxlists.append(prediction)

        gt_boxlist = dataset.get_groundtruth(image_id)
        gt_boxlists.append(gt_boxlist)
    if eval_attributes:
        classes = dataset.attributes
    else:
        classes = dataset.classes
    result = eval_detection_voc(
        pred_boxlists=pred_boxlists,
        gt_boxlists=gt_boxlists,
        classes=classes,
        iou_thresh=0.5,
        eval_attributes=eval_attributes,
        use_07_metric=False,
    )
    result_str = "mAP: {:.4f}\n".format(result["map"])
    for i, ap in enumerate(result["ap"]):
        # if i == 0:  # skip background
        #     continue
        # we skipped background in result['ap'], so we need to use i+1
        if eval_attributes:
            result_str += "{:<16}: {:.4f}\n".format(
                dataset.map_attribute_id_to_attribute_name(i+1), ap
            )
        else:
            result_str += "{:<16}: {:.4f}\n".format(
                dataset.map_class_id_to_class_name(i+1), ap
            )
    logger.info(result_str)
    # return mAP and weighted mAP
    if eval_attributes:
        if output_folder and save_predictions:
            with open(os.path.join(output_folder, "result_attr.txt"), "w") as fid:
                fid.write(result_str)
        return {"attr": {"map": result["map"], "weighted map": result["weighted map"]}}
    else:
        if output_folder and save_predictions:
            with open(os.path.join(output_folder, "result_obj.txt"), "w") as fid:
                fid.write(result_str)
        return {"obj": {"map": result["map"], "weighted map": result["weighted map"]}}

def eval_detection_voc(pred_boxlists, gt_boxlists, classes, iou_thresh=0.5, eval_attributes=False, use_07_metric=False):
    """Evaluate on voc dataset.
    Args:
        pred_boxlists(list[BoxList]): pred boxlist, has labels and scores fields.
        gt_boxlists(list[BoxList]): ground truth boxlist, has labels field.
        iou_thresh: iou thresh
        use_07_metric: boolean
    Returns:
        dict represents the results
    """
    assert len(gt_boxlists) == len(
        pred_boxlists
    ), "Length of gt and pred lists need to be same."

    aps = []
    nposs = []
    thresh = []

    for i, classname in enumerate(classes):
        if classname == "__background__" or classname == "__no_attribute__":
            continue
        rec, prec, ap, scores, npos = calc_detection_voc_prec_rec(pred_boxlists=pred_boxlists, gt_boxlists=gt_boxlists, \
                                                                  classindex=i, iou_thresh=iou_thresh,
                                                                  eval_attributes=eval_attributes,
                                                                  use_07_metric=use_07_metric)
        # Determine per class detection thresholds that maximise f score
        # if npos > 1:
        if npos > 1 and type(scores) != np.int:
            f = np.nan_to_num((prec * rec) / (prec + rec))
            thresh += [scores[np.argmax(f)]]
        else:
            thresh += [0]
        aps += [ap]
        nposs += [float(npos)]
        print('AP for {} = {:.4f} (npos={:,})'.format(classname, ap, npos))
        # if pickle:
        #     with open(os.path.join(output_dir, cls + '_pr.pkl'), 'w') as f:
        #         cPickle.dump({'rec': rec, 'prec': prec, 'ap': ap,
        #             'scores': scores, 'npos':npos}, f)

    # Set thresh to mean for classes with poor results
    thresh = np.array(thresh)
    avg_thresh = np.mean(thresh[thresh != 0])
    thresh[thresh == 0] = avg_thresh
    # if eval_attributes:
    #     filename = 'attribute_thresholds_' + self._image_set + '.txt'
    # else:
    #     filename = 'object_thresholds_' + self._image_set + '.txt'
    # path = os.path.join(output_dir, filename)
    # with open(path, 'wt') as f:
    #     for i, cls in enumerate(classes[1:]):
    #         f.write('{:s} {:.3f}\n'.format(cls, thresh[i]))

    weights = np.array(nposs)
    weights /= weights.sum()
    print('Mean AP = {:.4f}'.format(np.mean(aps)))
    print('Weighted Mean AP = {:.4f}'.format(np.average(aps, weights=weights)))
    print('Mean Detection Threshold = {:.3f}'.format(avg_thresh))
    print('~~~~~~~~')
    print('Results:')
    for ap, npos in zip(aps, nposs):
        print('{:.3f}\t{:.3f}'.format(ap, npos))
    print('{:.3f}'.format(np.mean(aps)))
    print('~~~~~~~~')
    print('')
    print('--------------------------------------------------------------')
    print('Results computed with the **unofficial** PASCAL VOC Python eval code.')
    print('--------------------------------------------------------------')

    # pdb.set_trace()
    return {"ap": aps, "map": np.mean(aps), "weighted map": np.average(aps, weights=weights)}

def calc_detection_voc_prec_rec(pred_boxlists, gt_boxlists, classindex, iou_thresh=0.5, eval_attributes=False,
                                use_07_metric=False):
    """Calculate precision and recall based on evaluation code of PASCAL VOC.
    This function calculates precision and recall of
    predicted bounding boxes obtained from a dataset which has :math:`N`
    images.
    The code is based on the evaluation code used in PASCAL VOC Challenge.
   """
    class_recs = {}
    npos = 0
    image_ids = []
    confidence = []
    BB = []
    for image_index, (gt_boxlist, pred_boxlist) in enumerate(zip(gt_boxlists, pred_boxlists)):
        pred_bbox = pred_boxlist.bbox.numpy()
        gt_bbox = gt_boxlist.bbox.numpy()
        if eval_attributes:
            gt_label = gt_boxlist.get_field("attributes").numpy()
            pred_label = pred_boxlist.get_field("attr_labels").numpy()
            pred_score = pred_boxlist.get_field("attr_scores").numpy()
        else:
            gt_label = gt_boxlist.get_field("labels").numpy()
            pred_label = pred_boxlist.get_field("labels").numpy()
            pred_score = pred_boxlist.get_field("scores").numpy()

        # get the ground truth information for this class
        if eval_attributes:
            gt_mask_l = np.array([classindex in i for i in gt_label])
        else:
            gt_mask_l = gt_label == classindex
        gt_bbox_l = gt_bbox[gt_mask_l]
        gt_difficult_l = np.zeros(gt_bbox_l.shape[0], dtype=bool)
        det = [False] * gt_bbox_l.shape[0]
        npos = npos + sum(~gt_difficult_l)
        class_recs[image_index] = {'bbox': gt_bbox_l,
                                   'difficult': gt_difficult_l,
                                   'det': det}

        # prediction output for each class
        # pdb.set_trace()
        if eval_attributes:
            pred_mask_l = np.logical_and(pred_label == classindex, np.not_equal(pred_score, 0.0)).nonzero()
            pred_bbox_l = pred_bbox[pred_mask_l[0]]
            pred_score_l = pred_score[pred_mask_l]
        else:
            pred_mask_l = pred_label == classindex
            pred_bbox_l = pred_bbox[pred_mask_l]
            pred_score_l = pred_score[pred_mask_l]

        for bbox_tmp, score_tmp in zip(pred_bbox_l, pred_score_l):
            image_ids.append(image_index)
            confidence.append(float(score_tmp))
            BB.append([float(z) for z in bbox_tmp])

    if npos == 0:
        # No ground truth examples
        return 0, 0, 0, 0, npos

    if len(confidence) == 0:
        # No detection examples
        return 0, 0, 0, 0, npos

    confidence = np.array(confidence)
    BB = np.array(BB)

    # sort by confidence
    sorted_ind = np.argsort(-confidence)
    sorted_scores = -np.sort(-confidence)
    BB = BB[sorted_ind, :]
    image_ids = [image_ids[x] for x in sorted_ind]

    # go down dets and mark TPs and FPs
    nd = len(image_ids)
    tp = np.zeros(nd)
    fp = np.zeros(nd)

    for d in range(nd):
        R = class_recs[image_ids[d]]
        bb = BB[d, :].astype(float)
        ovmax = -np.inf
        BBGT = R['bbox'].astype(float)

        if BBGT.size > 0:
            # compute overlaps
            # intersection
            ixmin = np.maximum(BBGT[:, 0], bb[0])
            iymin = np.maximum(BBGT[:, 1], bb[1])
            ixmax = np.minimum(BBGT[:, 2], bb[2])
            iymax = np.minimum(BBGT[:, 3], bb[3])
            iw = np.maximum(ixmax - ixmin + 1., 0.)
            ih = np.maximum(iymax - iymin + 1., 0.)
            inters = iw * ih

            # union
            uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) +
                   (BBGT[:, 2] - BBGT[:, 0] + 1.) *
                   (BBGT[:, 3] - BBGT[:, 1] + 1.) - inters)

            overlaps = inters / uni
            ovmax = np.max(overlaps)
            jmax = np.argmax(overlaps)

        if ovmax > iou_thresh:
            if not R['difficult'][jmax]:
                if not R['det'][jmax]:
                    tp[d] = 1.
                    R['det'][jmax] = 1
                else:
                    fp[d] = 1.
        else:
            fp[d] = 1.

    # compute precision recall
    fp = np.cumsum(fp)
    tp = np.cumsum(tp)
    rec = tp / float(npos)
    # avoid divide by zero in case the first detection matches a difficult
    # ground truth
    prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps)
    ap = voc_ap(rec, prec, use_07_metric)

    return rec, prec, ap, sorted_scores, npos

def voc_ap(rec, prec, use_07_metric=False):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:False).
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.
        for t in np.arange(0., 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.], rec, [1.]))
        mpre = np.concatenate(([0.], prec, [0.]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap

def calc_detection_voc_ap(prec, rec, use_07_metric=False):
    """Calculate average precisions based on evaluation code of PASCAL VOC.
    This function calculates average precisions
    from given precisions and recalls.
    The code is based on the evaluation code used in PASCAL VOC Challenge.
    Args:
        prec (list of numpy.array): A list of arrays.
            :obj:`prec[l]` indicates precision for class :math:`l`.
            If :obj:`prec[l]` is :obj:`None`, this function returns
            :obj:`numpy.nan` for class :math:`l`.
        rec (list of numpy.array): A list of arrays.
            :obj:`rec[l]` indicates recall for class :math:`l`.
            If :obj:`rec[l]` is :obj:`None`, this function returns
            :obj:`numpy.nan` for class :math:`l`.
        use_07_metric (bool): Whether to use PASCAL VOC 2007 evaluation metric
            for calculating average precision. The default value is
            :obj:`False`.
    Returns:
        ~numpy.ndarray:
        This function returns an array of average precisions.
        The :math:`l`-th value corresponds to the average precision
        for class :math:`l`. If :obj:`prec[l]` or :obj:`rec[l]` is
        :obj:`None`, the corresponding value is set to :obj:`numpy.nan`.
    """

    n_fg_class = len(prec)
    ap = np.empty(n_fg_class)
    for l in range(n_fg_class):
        if prec[l] is None or rec[l] is None:
            ap[l] = np.nan
            continue

        if use_07_metric:
            # 11 point metric
            ap[l] = 0
            for t in np.arange(0.0, 1.1, 0.1):
                if np.sum(rec[l] >= t) == 0:
                    p = 0
                else:
                    p = np.max(np.nan_to_num(prec[l])[rec[l] >= t])
                ap[l] += p / 11
        else:
            # correct AP calculation
            # first append sentinel values at the end
            mpre = np.concatenate(([0], np.nan_to_num(prec[l]), [0]))
            mrec = np.concatenate(([0], rec[l], [1]))

            mpre = np.maximum.accumulate(mpre[::-1])[::-1]

            # to calculate area under PR curve, look for points
            # where X axis (recall) changes value
            i = np.where(mrec[1:] != mrec[:-1])[0]

            # and sum (\Delta recall) * prec
            ap[l] = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])

    return ap


import torch.utils.data as data
from pycocotools.coco import COCO
from PIL import Image
import csv
import numpy as np
import os
import pickle
import torch
import xml.etree.ElementTree as ET
import re
from lxml import etree
import cv2
import json
import os

TOTAL_LABELS = 66

class CoCoDataset(data.Dataset):

    def __init__(self, transform, version='train', mask_person=0): # version is train, val, or test
        self.transform = transform
        self.version = version
        self.mask_person = mask_person

        self.img_folder = '/n/fs/visualai-scr/Data/Coco/2014data/{}2014'.format(version)
        self.coco = COCO('/n/fs/visualai-scr/Data/Coco/2014data/annotations/instances_{}2014.json'.format(version))
        gender_data = pickle.load(open('/n/fs/visualai-scr/Data/Coco/2014data/bias_splits/{}.data'.format(version), 'rb'))
        self.gender_info = {int(chunk['img'][10+len(version):22+len(version)]): chunk['annotation'][0] for chunk in gender_data}

        ids = list(self.coco.anns.keys())
        self.image_ids = list(set([self.coco.anns[this_id]['image_id'] for this_id in ids]))

        # dataset will only include gendered people:
        self.image_ids = list(set(self.gender_info.keys()) & set(self.image_ids))
        print("version: {0}, num: {1}".format(version, len(self.image_ids)))

        kept_objs = open('/n/fs/biasamp-scr/biasamp-arrow/mals/data/objs').readlines()
        kept_objs = [obj.rstrip() for obj in kept_objs]

        cats = self.coco.loadCats(self.coco.getCatIds())
        self.labels_to_names = {}
        for cat in cats:
            if cat['name'].replace(' ', '_') in kept_objs:
                self.labels_to_names[cat['id']] = cat['name']

        self.categories = list(self.labels_to_names.keys())
        #self.categories = list(pickle.load(open('/n/fs/biasamp-scr/biasamp-arrow/coco_mask/coco_labels_to_keep.pkl', 'rb')))

    def __getitem__(self, index):
        image_id = self.image_ids[index]
        path = self.coco.loadImgs(int(image_id))[0]["file_name"]

        file_path = os.path.join(self.img_folder, path)
        return self.from_path(file_path)

    def __len__(self):
        return len(self.image_ids)

    def from_path(self, file_path):
        image_id = int(os.path.basename(file_path)[-16:-4])

        image = np.array(Image.open(file_path).convert("RGB"))

        annIds = self.coco.getAnnIds(imgIds=image_id);
        coco_anns = self.coco.loadAnns(annIds) # coco is [x, y, width, height]
        formatted_anns = []
        target = np.zeros(TOTAL_LABELS+1)
        image.setflags(write=1)

        all_people_masks = np.ones((image.shape[0], image.shape[1]))
        all_other_masks = np.ones((image.shape[0], image.shape[1]))

        for ann in coco_anns:
            if ann['category_id'] not in self.categories:
                continue
            target[self.categories.index(ann['category_id'])] = 1

            if ann['category_id'] == 1:
                if self.mask_person == 1:
                    x, y, w, h = ann['bbox']
                    x, y, w, h = int(x), int(y), int(w), int(h)
                    image[y:y+h, x:x+w, :] = 0

                elif self.mask_person == 2:
                    labelMask = np.expand_dims(1 - self.coco.annToMask(ann), 2)
                    image = image * labelMask
                elif self.mask_person == 3:
                    labelMask = self.coco.annToMask(ann)

                    # randomly put half of the mask back
                    random_half = np.ones_like(labelMask).flatten()
                    random_half[:len(random_half)//2] = 0
                    np.random.shuffle(random_half)
                    random_half = random_half.reshape(labelMask.shape)

                    labelMask = labelMask * random_half
                    labelMask = np.expand_dims(1 - labelMask, 2)

                    image = image * labelMask
                elif self.mask_person == 4: # just person
                    labelMask = np.expand_dims(self.coco.annToMask(ann), 2)
                    image = image * labelMask
                elif self.mask_person == 5:
                    labelMask = np.expand_dims(self.coco.annToMask(ann), 2)
                    new_mask = np.zeros_like(labelMask)
                    possible_mid = int(np.sqrt(np.sum(labelMask)))
                    possible_ws = np.arange(10) + (possible_mid - 5)
                    w = possible_ws[np.argmin([np.sum(labelMask) % pos if pos > 1 else float('inf') for pos in possible_ws])]
                    h = np.sum(labelMask) // w
                    h, w = int(np.clip(h, 1, len(labelMask[0]) - 1)), int(np.clip(w, 1, len(labelMask) - 1))
                    y = np.random.randint(len(labelMask[0]) - h)
                    x = np.random.randint(len(labelMask) - w)
                    new_mask[x:x+w, y:y+h] = 1.
                    new_mask = 1 - new_mask
                    image = image * new_mask

                elif self.mask_person == 6:
                    labelMask = np.expand_dims(self.coco.annToMask(ann), 2)
                    new_mask = np.zeros_like(labelMask)
                    possible_mid = int(np.sqrt(np.sum(labelMask)))
                    possible_ws = np.arange(10) + (possible_mid - 5)
                    w = possible_ws[np.argmin([np.sum(labelMask) % pos if pos > 1 else float('inf') for pos in possible_ws])]
                    h = np.sum(labelMask) // w
                    h, w = int(np.clip(h, 1, len(labelMask[0]) - 1)), int(np.clip(w, 1, len(labelMask) - 1))
                    y = np.random.randint(len(labelMask[0]) - h)
                    x = np.random.randint(len(labelMask) - w)
                    new_mask[x:x+w, y:y+h] = 1.
                    image = image * new_mask
                elif self.mask_person == 7: # masking the person, but adds back in any objects
                    all_people_masks = all_people_masks * (1. - self.coco.annToMask(ann))
                elif self.mask_person == 8: # masking the objects, but adds back in the person
                    all_people_masks = all_people_masks * (1. - self.coco.annToMask(ann))
                elif self.mask_person == 9: # like 7 except now only masking half of the person
                    all_people_masks = all_people_masks * (1. - self.coco.annToMask(ann))
            else:
                if self.mask_person == 7:
                    all_other_masks = all_other_masks * (1. - self.coco.annToMask(ann))
                elif self.mask_person == 8:
                    all_other_masks = all_other_masks * (1. - self.coco.annToMask(ann))
                if self.mask_person == 9:
                    all_other_masks = all_other_masks * (1. - self.coco.annToMask(ann))

        if self.mask_person == 7:
            all_other_masks, all_people_masks = all_other_masks.astype(int), all_people_masks.astype(int)
            just_person_mask = np.expand_dims(1. - np.bitwise_and(all_other_masks, 1 - all_people_masks), 2).astype(np.uint8)
            image = image * just_person_mask
        elif self.mask_person == 8:
            all_other_masks, all_people_masks = all_other_masks.astype(int), all_people_masks.astype(int)
            just_objects_mask = np.expand_dims(1. - np.bitwise_and(all_people_masks, 1 - all_other_masks), 2).astype(np.uint8)
            image = image * just_objects_mask
        elif self.mask_person == 9:
            #from PIL import Image as pilimage
            #im = pilimage.fromarray(image)
            #im.save('outfile.jpeg')

            all_other_masks, all_people_masks = all_other_masks.astype(int), all_people_masks.astype(int)

            #just_person_mask = np.expand_dims(1. - np.bitwise_and(all_other_masks, 1 - all_people_masks), 2).astype(np.uint8)
            #image1 = image * just_person_mask
            #im = pilimage.fromarray(image1)
            #im.save('outfile1.jpeg')

            # randomly put half of the mask back
            all_people_masks = 1 - all_people_masks
            random_half = np.ones_like(all_people_masks).flatten()
            random_half[:len(random_half)//2] = 0
            np.random.shuffle(random_half)
            random_half = random_half.reshape(all_people_masks.shape)
            all_people_masks = all_people_masks * random_half
            all_people_masks = 1 - all_people_masks
            just_person_mask = np.expand_dims(1. - np.bitwise_and(all_other_masks, 1 - all_people_masks), 2).astype(np.uint8)
            image = image * just_person_mask

            #im = pilimage.fromarray(image)
            #im.save('outfile2.jpeg')

        image = Image.fromarray(image)
        target[TOTAL_LABELS] = self.gender_info[image_id]

        image = self.transform(image)

        return image, target


#!/usr/bin/env python
# -*- coding:UTF-8 -*-
# File Name     : main-lstm.py
# Purpose       :
# Creation Date : 05-01-2019
# Author        : Shuang Li [sli[at]informatik[dot]uni-hamburg[dot]de]
# Author        : Hongzhuo Liang [liang[at]informatik[dot]uni-hamburg[dot]de]
from __future__ import division, print_function
import argparse
import os
import time
import numpy as np
import torch
import torch.utils.data
import torch.nn as nn
import torch.nn.functional as torch_func
import glob
import rospy
import copy
from geometry_msgs.msg import WrenchStamped
from portaudio_transport.msg import AudioTransport
from std_msgs.msg import Bool, Float32
from audio_pouring.utils.utils import weight2height, config
# from audio_pouring.model.long_preprocess import butter_low_pass_filter
# from audio_pouring.model.long_preprocess import vis_wrench
import warnings

warnings.simplefilter(action="ignore", category=FutureWarning)
try:
    import librosa.util
except ImportError:
    librosa = None
parser = argparse.ArgumentParser(description="audio2height")
parser.add_argument("--cuda", action="store_true")
parser.add_argument("--gpu", type=int, default=0)
parser.add_argument("--bottle", type=int, choices=config["bottle_id_list"], required=False)
parser.add_argument("--cavity-height", type=float, default=40.0)
parser.add_argument("--model-folder-path", type=str, default="./assets/learned_models/robot_experiment/")
parser.add_argument("--multi", action="store_true")
parser.add_argument("--test-robot", action="store_true")
parser.add_argument("--scale", action="store_true", help="add this if a scale that publish ros topic is connected")
parser.add_argument("--test_offline", action="store_true")
parser.add_argument("--minus_wrench_first", action="store_true")
parser.add_argument("--hidden-dim", type=int, default=64)
args = parser.parse_args()

args.cuda = args.cuda if torch.cuda.is_available() else False

if args.cuda:
    torch.cuda.manual_seed(1)

np.random.seed(int(time.time()))
target_frequency = config["target_frequency"]
source_frequency = config["source_frequency"]  # unit Hz
audio_topic = config["audio_topic"]
# ft_topic = config["ur5_ft_topic"]
ft_topic = config["ati_ft_topic"]
audio_length = config["fixed_audio_length"]  # unit second
n_fft = config["n_fft"]
win_size = config["win_size"]
overlap = config["overlap"]
if args.minus_wrench_first:
    ft_mean = np.array(config["minus"]["ft_mean"])
    ft_std = np.array(config["minus"]["ft_std"])
else:
    ft_mean = np.array(config["raw"]["ft_mean"])
    ft_std = np.array(config["raw"]["ft_std"])
win_length = int(win_size * target_frequency)
hop_length = int(win_length * overlap)
thresh_acc = np.array([1, 2, 3, 4, 5])  # unit mm
audio_mean = config["audio_mean"]
audio_std = config["audio_std"]
input_audio_size = config["input_audio_size"]
input_force_size = config["input_force_size_raw"]
force_embedding_size = input_force_size // 2
bottle_upper = torch.tensor([config["bottle_upper"]])
bottle_lower = torch.tensor([config["bottle_lower"]])
max_audio = config["max_audio"]

if args.multi:
    model_name = "0205_multi_lstm2_h64_bs64_bottlerobot_pouringtorobot_pouring_mono_coe0.001_snr-1000.0_a_f_early_fusion_raw_raw_force_bidirectional.model"
else:
    # mixed noise
    model_name = "0207_audio_lstm2_h64_bs64_bottlerobot_pouringtorobot_pouring_mono_coe0.001_snr-1000.0_audio_only_raw_raw_force_bidirectional.model"
    # without noise
    # model_name = "0208_audio_lstm2_h64_bs64_bottlerobot_pouringtorobot_pouring_mono_coe0.001_snr1000.0_audio_only_raw_raw_force_bidirectional.model"
if args.cuda:
    model = torch.load(args.model_folder_path + model_name, map_location="cuda:{}".format(args.gpu))
    model.device_ids = [args.gpu]
else:
    model = torch.load(args.model_folder_path + model_name, map_location="cpu")
model.mini_batch_size = 1
# if args.multi:
#     model.scale.batch_size = 1

device = torch.device("cuda" if args.cuda else "cpu")
if args.cuda:
    if args.gpu != -1:
        torch.cuda.set_device(args.gpu)
        model = model.cuda()
    else:
        device_id = [0, 1]
        torch.cuda.set_device(device_id[0])
        model = nn.DataParallel(model, device_ids=device_id).cuda()
    bottle_upper = bottle_upper.cuda()
    bottle_lower = bottle_lower.cuda()

def test(model_, audio, force=None):
    model_.eval()
    torch.set_grad_enabled(False)
    if args.cuda:
        audio = audio.cuda()
    if args.multi:
        if args.cuda:
            force = force.cuda()
        model.hidden = model.init_hidden(hidden_dim=args.hidden_dim, device=device)
        # force hidden
        model.force_hidden = model.init_hidden(hidden_dim=force_embedding_size, device=device)
        height = model_(audio, force)
    else:
        model_.hidden = model_.init_hidden(hidden_dim=args.hidden_dim, device=device)
        height = model_(audio)
    height = height * (bottle_upper - bottle_lower) + bottle_lower
    return height.cpu().data.numpy()

class Pouring:
    def __init__(self, target_cavity_height, multi=False, scale=False, test_offline=False):
        self.target = target_cavity_height
        self.audio_numpy = np.array([])
        self.force_numpy = np.array([])
        self.last_height = 0
        self.last_2_height = 0
        self.scale_reading = 0.0
        self.scale_reading_last = 0.0
        self.initial_force_magnitude = 0.0
        self.multi = multi
        self.need_print = True
        self.test_offline = test_offline
        self.scale = scale
        self.pub = rospy.Publisher("stop_pour", Bool, queue_size=10)
        self.wrench_first = None
        if self.multi:
            self.network_name = "Multi"
        else:
            self.network_name = "Audio"
        self.pub_network_data = rospy.Publisher("cavity_network_height", Float32, queue_size=10)
        self.pub_real_data = rospy.Publisher("cavity_real_height", Float32, queue_size=10)
        if self.scale:
            self.scale_sub = rospy.Subscriber("/maul_logic/wrench", WrenchStamped, self.scale_callback)
        rospy.sleep(2)  # wait the scale to have reading
        self.pouring_run()
        rospy.spin()

    def scale_callback(self, scale_data):
        self.scale_reading_last = self.scale_reading
        self.scale_reading = scale_data.wrench.force.z

    def pouring_run(self):
        if self.scale:
            while self.scale_reading == self.scale_reading_last:
                pass
        else:
            # some method to judge if the pouring motion is begin
            pass
        if self.multi:
            rospy.Subscriber(ft_topic, WrenchStamped, self.force_callback)
        rospy.Subscriber(audio_topic, AudioTransport, self.audio_callback)
        rospy.sleep(0.5)
        if self.multi:
            while not rospy.is_shutdown():
                copy_audio_numpy = copy.copy(self.audio_numpy)
                audio_spectrum = torch.from_numpy(self.audio_process(copy_audio_numpy).T).type(torch.FloatTensor)
                force_dim = audio_spectrum.shape[0] * input_force_size
                current_force = copy.copy(self.force_numpy)
                while current_force.shape[0] < force_dim / 2:
                    rospy.loginfo("force is not enough, cut {} column audio".
                                  format(len(audio_spectrum) - current_force.shape[0] // 48))
                    # audio_spectrum = audio_spectrum[-(current_force.shape[0] // 48):, :]
                    # force_dim = current_force.shape[0] // 48 * 48
                    exit()
                audio_spectrum = audio_spectrum.unsqueeze(0)
                force_input_numpy = current_force.reshape(-1, 6)
                if args.minus_wrench_first:
                    force_input_numpy -= self.wrench_first
                force_input_numpy -= ft_mean
                force_input_numpy /= ft_std
                final_force = np.array([])
                for i in range(audio_spectrum.shape[1]):
                    final_force = np.hstack([final_force, force_input_numpy[4 * i:4 * i + 8, :].reshape(-1)])
                final_force = final_force.reshape(-1, input_force_size)
                final_force = final_force.astype(np.float32)
                final_force = final_force[-force_dim:]

                # time_record = np.arange(0, len(final_force), 1)
                # vis_wrench(final_force, time_record, raw_force_data=None, lw=1)
                force_input = torch.from_numpy(final_force).type(torch.FloatTensor)
                force_input = force_input.unsqueeze(0)
                all_height = test(model, audio_spectrum, force_input)
                self.process_height(all_height)
        else:
            while not rospy.is_shutdown():
                audio_spectrum = torch.from_numpy(self.audio_process(self.audio_numpy).T).type(torch.FloatTensor)
                audio_spectrum = audio_spectrum.unsqueeze(0)
                all_height = test(model, audio_spectrum)
                self.process_height(all_height)

    def process_height(self, all_height):
        real_height = weight2height(cup_id=args.bottle, cur_weight=self.scale_reading)
        if self.test_offline:  # test with rosbag
            height = all_height[-1][0]
            wait_time = 0  # do not need to wait
            self.pub_network_data.publish(height)
            self.pub_real_data.publish(real_height)
        else:  # real robot experiment
            if self.multi:
                height = all_height[-1][0] - 5
            else:
                height = all_height[-1][0] - 5
            self.pub_network_data.publish(height + 5)
            self.pub_real_data.publish(real_height)
            if self.scale:
                wait_time = 4  # wait for 4 seconds the scale to be stable
            else:
                wait_time = 0  # if no scale is connected, no need to wait
        print("output is {}, go on pouring".format(height))
        if height <= self.target and abs(height - self.last_2_height) < 8:
            rospy.loginfo("Enjoy your drink!")
            start_time = time.time()
            while not rospy.is_shutdown():
                self.pub.publish(True)
                if self.scale:
                    real_height = weight2height(cup_id=args.bottle, cur_weight=self.scale_reading)
                    info = "{} Network {}mm, Real {}mm, Desire {}mm, Scale {}kg"
                    info = info.format(self.network_name, height, real_height, self.target, self.scale_reading)
                else:
                    info = "Network {}mm, Desire {}mm".format(height, self.target)
                if (time.time() - start_time > wait_time) and self.need_print:
                    rospy.loginfo(info)
                    self.need_print = False
        self.last_2_height = self.last_height
        self.last_height = height

    def audio_callback(self, audio_raw):
        audio_numpy = np.array(audio_raw.channels[0].frame_data)
        self.audio_numpy = np.hstack([self.audio_numpy, audio_numpy])
        # get latest 4 second messages
        if self.audio_numpy.shape[0] > source_frequency * audio_length:
            self.audio_numpy = self.audio_numpy[-source_frequency * audio_length:]

    def force_callback(self, force_raw):
        f_x = force_raw.wrench.force.x
        f_y = force_raw.wrench.force.y
        f_z = force_raw.wrench.force.z
        t_x = force_raw.wrench.torque.x
        t_y = force_raw.wrench.torque.y
        t_z = force_raw.wrench.torque.z
        force_data = np.array([f_x, f_y, f_z, t_x, t_y, t_z])
        if self.wrench_first is None:
            self.wrench_first = force_data
        self.force_numpy = np.hstack([self.force_numpy, force_data])
        if len(self.force_numpy) > input_force_size * 251:
            self.force_numpy = self.force_numpy[-input_force_size * 251:]

    @staticmethod
    def audio_process(audio_raw):
        audio_resample = librosa.core.resample(audio_raw, orig_sr=source_frequency, target_sr=target_frequency)
        # audio_resample *= 1. / max_audio * 0.9
        audio_fft = librosa.stft(y=audio_resample, n_fft=n_fft, hop_length=hop_length, win_length=win_length)
        audio_fft = np.abs(audio_fft)
        audio_db = librosa.core.amplitude_to_db(audio_fft)
        audio_db -= audio_mean
        audio_db /= audio_std
        return audio_db

    def velocity_change(self):
        pass

def main():
    if args.test_robot:
        test_files = glob.glob(os.path.join("./dataset/test", "*.npy"))
        test_files.sort()
        for i in range(len(test_files)):
            if len(np.load(test_files[i])) == 5:
                file_name, spectrum, force, height, scale = np.load(test_files[i])
            else:
                file_name, spectrum, height, scale = np.load(test_files[i])
            spectrum -= audio_mean
            spectrum /= audio_std
            spectrum = spectrum.T
            result = test(model, torch.from_numpy(spectrum).unsqueeze(0))
            result = result.squeeze()
            loss_mse = torch_func.mse_loss(torch.from_numpy(height), torch.from_numpy(result))
            loss_f1 = torch_func.l1_loss(torch.from_numpy(height), torch.from_numpy(result), reduction="mean")
            loss_f1 = loss_f1.cpu().data.numpy()
            is_correct = [abs(height.reshape(1, -1) - result.reshape(1, -1)) < thresh for thresh in thresh_acc]
            res_acc = [np.sum(cc) for cc in is_correct]
            acc = float(res_acc[-1]) / float(len(height))
            change_h = result[1:] - result[:-1]
            mono = np.sum(change_h[np.where(change_h > 0)])
            print("mse:{:9.4f}|err:{:9.4f}|mono:{:9.4f}|acc:{:.4f}".format(loss_mse, loss_f1, mono, acc))
    else:
        rospy.init_node("robot_pouring_demo")
        while not rospy.is_shutdown():
            Pouring(target_cavity_height=args.cavity_height, multi=args.multi, scale=args.scale,
                    test_offline=args.test_offline)

if __name__ == "__main__":
    main()


# Copyright 2020,2021 Sony Corporation.
# Copyright 2021 Sony Group Corporation.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import Any, Dict, List, Union, cast

import gym
import numpy as np

import nnabla as nn
import nnabla.solvers as NS
import nnabla_rl.environment_explorers as EE
import nnabla_rl.model_trainers as MT
from nnabla_rl.algorithm import Algorithm, AlgorithmConfig, eval_api
from nnabla_rl.builders import ExplorerBuilder, ModelBuilder, ReplayBufferBuilder, SolverBuilder
from nnabla_rl.environment_explorer import EnvironmentExplorer
from nnabla_rl.environments.environment_info import EnvironmentInfo
from nnabla_rl.model_trainers.model_trainer import ModelTrainer, TrainingBatch
from nnabla_rl.models import QFunction, SACPolicy, SACQFunction, SACVFunction, StochasticPolicy, VFunction
from nnabla_rl.replay_buffer import ReplayBuffer
from nnabla_rl.utils import context
from nnabla_rl.utils.data import add_batch_dimension, marshal_experiences, set_data_to_variable
from nnabla_rl.utils.misc import create_variable, sync_model

@dataclass
class ICML2018SACConfig(AlgorithmConfig):
    '''ICML2018SACConfig
    List of configurations for ICML2018SAC algorithm.

    Args:
        gamma (float): discount factor of rewards. Defaults to 0.99.
        learning_rate (float): learning rate which is set to all solvers. \
            You can customize/override the learning rate for each solver by implementing the \
            (:py:class:`SolverBuilder <nnabla_rl.builders.SolverBuilder>`) by yourself. \
            Defaults to 0.0003.
        batch_size(int): training batch size. Defaults to 256.
        tau (float): target network's parameter update coefficient. Defaults to 0.005.
        environment_steps (int): Number of steps to interact with the environment on each iteration. Defaults to 1.
        gradient_steps (int): Number of parameter updates to perform on each iteration. Defaults to 1.
        reward_scalar (float): Reward scaling factor. Obtained reward will be multiplied by this value. Defaults to 5.0.
        start_timesteps (int): the timestep when training starts.\
            The algorithm will collect experiences from the environment by acting randomly until this timestep.\
            Defaults to 10000.
        replay_buffer_size (int): capacity of the replay buffer. Defaults to 1000000.
        target_update_interval (float): the interval of target v function parameter's update. Defaults to 1.
    '''

    gamma: float = 0.99
    learning_rate: float = 3.0*1e-4
    batch_size: int = 256
    tau: float = 0.005
    environment_steps: int = 1
    gradient_steps: int = 1
    reward_scalar: float = 5.0
    start_timesteps: int = 10000
    replay_buffer_size: int = 1000000
    target_update_interval: int = 1

    def __post_init__(self):
        '''__post_init__

        Check the values are in valid range.

        '''
        self._assert_between(self.tau, 0.0, 1.0, 'tau')
        self._assert_between(self.gamma, 0.0, 1.0, 'gamma')
        self._assert_positive(self.gradient_steps, 'gradient_steps')
        self._assert_positive(self.environment_steps, 'environment_steps')
        self._assert_positive(self.start_timesteps, 'start_timesteps')
        self._assert_positive(self.target_update_interval, 'target_update_interval')

class DefaultVFunctionBuilder(ModelBuilder[VFunction]):
    def build_model(self,  # type: ignore[override]
                    scope_name: str,
                    env_info: EnvironmentInfo,
                    algorithm_config: ICML2018SACConfig,
                    **kwargs) -> VFunction:
        return SACVFunction(scope_name)

class DefaultQFunctionBuilder(ModelBuilder[QFunction]):
    def build_model(self,  # type: ignore[override]
                    scope_name: str,
                    env_info: EnvironmentInfo,
                    algorithm_config: ICML2018SACConfig,
                    **kwargs) -> QFunction:
        return SACQFunction(scope_name)

class DefaultPolicyBuilder(ModelBuilder[StochasticPolicy]):
    def build_model(self,  # type: ignore[override]
                    scope_name: str,
                    env_info: EnvironmentInfo,
                    algorithm_config: ICML2018SACConfig,
                    **kwargs) -> StochasticPolicy:
        return SACPolicy(scope_name, env_info.action_dim)

class DefaultSolverBuilder(SolverBuilder):
    def build_solver(self,  # type: ignore[override]
                     env_info: EnvironmentInfo,
                     algorithm_config: ICML2018SACConfig,
                     **kwargs) -> nn.solver.Solver:
        assert isinstance(algorithm_config, ICML2018SACConfig)
        return NS.Adam(alpha=algorithm_config.learning_rate)

class DefaultReplayBufferBuilder(ReplayBufferBuilder):
    def build_replay_buffer(self,  # type: ignore[override]
                            env_info: EnvironmentInfo,
                            algorithm_config: ICML2018SACConfig,
                            **kwargs) -> ReplayBuffer:
        assert isinstance(algorithm_config, ICML2018SACConfig)
        return ReplayBuffer(capacity=algorithm_config.replay_buffer_size)

class DefaultExplorerBuilder(ExplorerBuilder):
    def build_explorer(self,  # type: ignore[override]
                       env_info: EnvironmentInfo,
                       algorithm_config: ICML2018SACConfig,
                       algorithm: "ICML2018SAC",
                       **kwargs) -> EnvironmentExplorer:
        explorer_config = EE.RawPolicyExplorerConfig(
            warmup_random_steps=algorithm_config.start_timesteps,
            reward_scalar=algorithm_config.reward_scalar,
            initial_step_num=algorithm.iteration_num,
            timelimit_as_terminal=False
        )
        explorer = EE.RawPolicyExplorer(policy_action_selector=algorithm._compute_greedy_action,
                                        env_info=env_info,
                                        config=explorer_config)
        return explorer

class ICML2018SAC(Algorithm):
    '''Soft Actor-Critic (SAC) algorithm.

    This class implements the ICML2018 version of Soft Actor Critic (SAC) algorithm proposed by T. Haarnoja, et al.
    in the paper: "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor"
    For detail see: https://arxiv.org/abs/1801.01290

    This implementation slightly differs from the implementation of Soft Actor-Critic algorithm presented
    also by T. Haarnoja, et al. in the following paper: https://arxiv.org/abs/1812.05905

    You will need to scale the reward received from the environment properly to get the algorithm work.

    Args:
        env_or_env_info \
        (gym.Env or :py:class:`EnvironmentInfo <nnabla_rl.environments.environment_info.EnvironmentInfo>`):
            the environment to train or environment info
        config (:py:class:`ICML2018SACConfig <nnabla_rl.algorithms.icml2018_sac.ICML2018SACConfig>`):
            configuration of the ICML2018SAC algorithm
        v_function_builder (:py:class:`ModelBuilder[VFunction] <nnabla_rl.builders.ModelBuilder>`):
            builder of v function models
        v_solver_builder (:py:class:`SolverBuilder <nnabla_rl.builders.SolverBuilder>`):
            builder of v function solvers
        q_function_builder (:py:class:`ModelBuilder[QFunction] <nnabla_rl.builders.ModelBuilder>`):
            builder of q function models
        q_solver_builder (:py:class:`SolverBuilder <nnabla_rl.builders.SolverBuilder>`):
            builder of q function solvers
        policy_builder (:py:class:`ModelBuilder[StochasticPolicy] <nnabla_rl.builders.ModelBuilder>`):
            builder of actor models
        policy_solver_builder (:py:class:`SolverBuilder <nnabla_rl.builders.SolverBuilder>`):
            builder of policy solvers
        replay_buffer_builder (:py:class:`ReplayBufferBuilder <nnabla_rl.builders.ReplayBufferBuilder>`):
            builder of replay_buffer
        explorer_builder (:py:class:`ExplorerBuilder <nnabla_rl.builders.ExplorerBuilder>`):
            builder of environment explorer
    '''

    # type declarations to type check with mypy
    # NOTE: declared variables are instance variable and NOT class variable, unless it is marked with ClassVar
    # See https://mypy.readthedocs.io/en/stable/class_basics.html for details
    _config: ICML2018SACConfig
    _v: VFunction
    _v_solver: nn.solver.Solver
    _target_v: VFunction
    _q1: QFunction
    _q2: QFunction
    _train_q_functions: List[QFunction]
    _train_q_solvers: Dict[str, nn.solver.Solver]
    _replay_buffer: ReplayBuffer
    _explorer_builder: ExplorerBuilder
    _environment_explorer: EnvironmentExplorer
    _policy_trainer: ModelTrainer
    _q_function_trainer: ModelTrainer
    _v_function_trainer: ModelTrainer

    _eval_state_var: nn.Variable
    _eval_deterministic_action: nn.Variable
    _eval_probabilistic_action: nn.Variable

    _policy_trainer_state: Dict[str, Any]
    _q_function_trainer_state: Dict[str, Any]
    _v_function_trainer_state: Dict[str, Any]

    def __init__(self, env_or_env_info: Union[gym.Env, EnvironmentInfo],
                 config: ICML2018SACConfig = ICML2018SACConfig(),
                 v_function_builder: ModelBuilder[VFunction] = DefaultVFunctionBuilder(),
                 v_solver_builder: SolverBuilder = DefaultSolverBuilder(),
                 q_function_builder: ModelBuilder[QFunction] = DefaultQFunctionBuilder(),
                 q_solver_builder: SolverBuilder = DefaultSolverBuilder(),
                 policy_builder: ModelBuilder[StochasticPolicy] = DefaultPolicyBuilder(),
                 policy_solver_builder: SolverBuilder = DefaultSolverBuilder(),
                 replay_buffer_builder: ReplayBufferBuilder = DefaultReplayBufferBuilder(),
                 explorer_builder: ExplorerBuilder = DefaultExplorerBuilder()):
        super(ICML2018SAC, self).__init__(env_or_env_info, config=config)

        self._explorer_builder = explorer_builder

        with nn.context_scope(context.get_nnabla_context(self._config.gpu_id)):
            self._v = v_function_builder(scope_name="v", env_info=self._env_info, algorithm_config=self._config)
            self._v_solver = v_solver_builder(env_info=self._env_info, algorithm_config=self._config)
            self._target_v = cast(VFunction, self._v.deepcopy('target_' + self._v.scope_name))

            self._q1 = q_function_builder(scope_name="q1", env_info=self._env_info, algorithm_config=self._config)
            self._q2 = q_function_builder(scope_name="q2", env_info=self._env_info, algorithm_config=self._config)

            self._train_q_functions = [self._q1, self._q2]
            self._train_q_solvers = {}
            for q in self._train_q_functions:
                self._train_q_solvers[q.scope_name] = q_solver_builder(
                    env_info=self._env_info, algorithm_config=self._config)

            self._pi = policy_builder(scope_name="pi", env_info=self._env_info, algorithm_config=self._config)
            self._pi_solver = policy_solver_builder(env_info=self._env_info, algorithm_config=self._config)

            self._replay_buffer = replay_buffer_builder(env_info=self._env_info, algorithm_config=self._config)

    @eval_api
    def compute_eval_action(self, state, *, begin_of_episode=False):
        with nn.context_scope(context.get_nnabla_context(self._config.gpu_id)):
            action, _ = self._compute_greedy_action(state, deterministic=True, begin_of_episode=begin_of_episode)
            return action

    def _before_training_start(self, env_or_buffer):
        # set context globally to ensure that the training runs on configured gpu
        context.set_nnabla_context(self._config.gpu_id)
        self._environment_explorer = self._setup_environment_explorer(env_or_buffer)
        self._policy_trainer = self._setup_policy_training(env_or_buffer)
        self._q_function_trainer = self._setup_q_function_training(
            env_or_buffer)
        self._v_function_trainer = self._setup_v_function_training(
            env_or_buffer)

    def _setup_environment_explorer(self, env_or_buffer):
        return None if self._is_buffer(env_or_buffer) else self._explorer_builder(self._env_info, self._config, self)

    def _setup_policy_training(self, env_or_buffer):
        # NOTE: Fix temperature to 1.0. Because This version of SAC adjusts it by scaling the reward
        policy_trainer_config = MT.policy_trainers.SoftPolicyTrainerConfig(fixed_temperature=True)
        temperature = MT.policy_trainers.soft_policy_trainer.AdjustableTemperature(
            scope_name='temperature',
            initial_value=1.0)
        policy_trainer = MT.policy_trainers.SoftPolicyTrainer(
            models=self._pi,
            solvers={self._pi.scope_name: self._pi_solver},
            q_functions=self._train_q_functions,
            temperature=temperature,
            temperature_solver=None,
            env_info=self._env_info,
            config=policy_trainer_config)
        return policy_trainer

    def _setup_q_function_training(self, env_or_buffer):
        q_function_trainer_param = MT.q_value_trainers.VTargetedQTrainerConfig(
            reduction_method='mean',
            q_loss_scalar=0.5)
        q_function_trainer = MT.q_value_trainers.VTargetedQTrainer(
            train_functions=self._train_q_functions,
            solvers=self._train_q_solvers,
            target_functions=self._target_v,
            env_info=self._env_info,
            config=q_function_trainer_param)
        return q_function_trainer

    def _setup_v_function_training(self, env_or_buffer):
        v_function_trainer_config = MT.v_value.SoftVTrainerConfig(
            reduction_method='mean',
            v_loss_scalar=0.5
        )
        v_function_trainer = MT.v_value.SoftVTrainer(
            train_functions=self._v,
            solvers={self._v.scope_name: self._v_solver},
            target_functions=self._train_q_functions,  # Set training q as target
            target_policy=self._pi,
            env_info=self._env_info,
            config=v_function_trainer_config)
        sync_model(self._v, self._target_v, 1.0)

        return v_function_trainer

    def _run_online_training_iteration(self, env):
        for _ in range(self._config.environment_steps):
            self._run_environment_step(env)
        for _ in range(self._config.gradient_steps):
            self._run_gradient_step(self._replay_buffer)

    def _run_offline_training_iteration(self, buffer):
        self._sac_training(buffer)

    def _run_environment_step(self, env):
        experiences = self._environment_explorer.step(env)
        self._replay_buffer.append_all(experiences)

    def _run_gradient_step(self, replay_buffer):
        if self._config.start_timesteps < self.iteration_num:
            self._sac_training(replay_buffer)

    def _sac_training(self, replay_buffer):
        experiences, info = replay_buffer.sample(self._config.batch_size)
        (s, a, r, non_terminal, s_next, *_) = marshal_experiences(experiences)
        batch = TrainingBatch(batch_size=self._config.batch_size,
                              s_current=s,
                              a_current=a,
                              gamma=self._config.gamma,
                              reward=r,
                              non_terminal=non_terminal,
                              s_next=s_next,
                              weight=info['weights'])

        # Train in the order of v -> q -> policy
        self._v_function_trainer_state = self._v_function_trainer.train(batch)
        self._q_function_trainer_state = self._q_function_trainer.train(batch)
        if self.iteration_num % self._config.target_update_interval == 0:
            sync_model(self._v, self._target_v, tau=self._config.tau)
        self._policy_trainer_state = self._policy_trainer.train(batch)

        td_errors = self._q_function_trainer_state['td_errors']
        replay_buffer.update_priorities(td_errors)

    @eval_api
    def _compute_greedy_action(self, s, deterministic=False, *, begin_of_episode=False):
        # evaluation input/action variables
        s = add_batch_dimension(s)
        if not hasattr(self, '_eval_state_var'):
            self._eval_state_var = create_variable(1, self._env_info.state_shape)
            distribution = self._pi.pi(self._eval_state_var)
            self._eval_deterministic_action = distribution.choose_probable()
            self._eval_probabilistic_action = distribution.sample()
        set_data_to_variable(self._eval_state_var, s)
        if deterministic:
            self._eval_deterministic_action.forward()
            return np.squeeze(self._eval_deterministic_action.d, axis=0), {}
        else:
            self._eval_probabilistic_action.forward()
            return np.squeeze(self._eval_probabilistic_action.d, axis=0), {}

    def _models(self):
        models = [self._v, self._target_v, self._q1, self._q2, self._pi]
        return {model.scope_name: model for model in models}

    def _solvers(self):
        solvers = {}
        solvers.update(self._train_q_solvers)
        solvers[self._v.scope_name] = self._v_solver
        solvers[self._pi.scope_name] = self._pi_solver
        return solvers

    @classmethod
    def is_supported_env(cls, env_or_env_info):
        env_info = EnvironmentInfo.from_env(env_or_env_info) if isinstance(env_or_env_info, gym.Env) \
            else env_or_env_info
        return not env_info.is_discrete_action_env()

    @property
    def latest_iteration_state(self):
        latest_iteration_state = super(ICML2018SAC, self).latest_iteration_state
        if hasattr(self, '_policy_trainer_state'):
            latest_iteration_state['scalar'].update({'pi_loss': float(self._policy_trainer_state['pi_loss'])})
        if hasattr(self, '_v_function_trainer_state'):
            latest_iteration_state['scalar'].update({'v_loss': float(self._v_function_trainer_state['v_loss'])})
        if hasattr(self, '_q_function_trainer_state'):
            latest_iteration_state['scalar'].update({'q_loss': float(self._q_function_trainer_state['q_loss'])})
            latest_iteration_state['histogram'].update(
                {'td_errors': self._q_function_trainer_state['td_errors'].flatten()})
        return latest_iteration_state


import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torchsummary import summary
import numpy as np
import time
import gym
import cv2
import os

from agent import Agent
from replay_buffer import ReplayBuffer
from qnetwork import QNetwork

import matplotlib.pyplot as plt
from IPython.display import clear_output

import wandb

env_list = {
    0: "CartPole-v0",
    1: "CartPole-v2",
    2: "LunarLander-v2",
    3: "Breakout-v4",
    4: "BreakoutDeterministic-v4",
    5: "BreakoutNoFrameskip-v4",
    6: "BoxingDeterministic-v4",
    7: "PongDeterministic-v4",
}
env_name = env_list[6]
env = gym.make(env_name)
input_dim = 84
input_frame = 4
print("env_name", env_name)
print(env.unwrapped.get_action_meanings(), env.action_space.n)

update_start_buffer_size = 10000 # Buffer에 random action의 결과를 계속해서 쌓는 것을 언제까지 할지
num_frames = 10000000 # 총 학습 frame 수 (= 학습할 동안 agent가 action을 행하는 횟수)

# e-Greedy exploration 변수.
eps_max = 1.0
eps_min = 0.1
eps_decay = 1/1200000
# 위의 3개는 얼마나 agent가 greedy하게 움직이면서 exploration을 할지 정해주는 매개변수

gamma = 0.99 # Discounted factor. 더 먼 미래에 받을 보상일 수록 얼마나 discount할지를 정해주는 변수

buffer_size = int(8e5) # Replay Buffer 사이즈
batch_size = 32        # Replay Buffer에서 한번에 뽑아서 가져오는 sample의 개수
update_type = 'hard'   # Q-Network을 업데이트할 때, 어떤 방식으로 할지 정하는 것. 일반 DQN 논문들은 일반적으로 hard를 함.
target_update_freq = 250 # Update type이 hard일때, 몇 frame마다 update를 할지 정하는 변수.
soft_update_tau = 0.002  # Update type이 soft일때, target network의 변수를 가져오는 비율.
learning_rate = 0.0001
skipped_frame = 4 # 환경에서 주는 frame을 skip하는 것이 필요한 경우.

# Customized Parameters. GPU device 및 모델 저장 경로등을 설정하는 부분
device_num = 0
rand_seed = None
rand_name = ('').join(map(str, np.random.randint(10, size=(3,))))
folder_name = os.getcwd().split('/')[-1]
model_number = 0
model_save_path = \
f'./model_save/{rand_name}_{env_name}_num_f:{num_frames}_eps_dec:{round(eps_decay,8)}f\
_gamma:{gamma}_tar_up_frq:{target_update_freq}f\
_up_type:{update_type}_soft_tau:{soft_update_tau}f\
_batch:{batch_size}_buffer:{buffer_size}f\
_up_start:{update_start_buffer_size}_lr:{learning_rate}f\
_eps_min:{eps_min}_device:{device_num}_rand:{rand_seed}_{model_number}/'
if not os.path.exists('./model_save/'):
    os.mkdir('./model_save/')
if not os.path.exists(model_save_path):
    os.mkdir(model_save_path)
print("model_save_path:", model_save_path)

# Trained model 의 경로 지정
trained_model_path = ''

# Plot 옵션 선택. 1) wandb 이용 2) inline: 쥬피터 안에서 plot 3) False: plot안하기
plot_options = {0: 'wandb', 1: 'inline', 2: False}
plot_option = plot_options[0]

if plot_option=='wandb':
    call(["wandb", "login", "000c1d3d8ebb4219c3a579d5ae02bc38be380c70"])
    os.environ['WANDB_NOTEBOOK_NAME'] = 'RL_experiment'
    wandb_project_name = 'temporary'
    wandb.init(
            project=wandb_project_name,
            name=f"{rand_name}_{folder_name}_{env_name}",
            config={"env_name": env_name,
                    "input_frame": input_frame,
                    "input_dim": input_dim,
                    "eps_max": eps_max,
                    "eps_min": eps_min,
                    "eps_decay": eps_decay,
                    "num_frames": num_frames,
                    "skipped_frame": skipped_frame,
                    "gamma": gamma,
                    "buffer_size": buffer_size,
                    "update_start_buffer_size": update_start_buffer_size,
                    "batch_size": batch_size,
                    "update_type": update_type,
                    "soft_update_tau": soft_update_tau,
                    "learning_rate": learning_rate,
                    "target_update_freq (unit:frames)": target_update_freq,
                    }
            )

agent = Agent(
    env,
    input_frame,
    input_dim,
    num_frames,
    skipped_frame,
    eps_decay,
    gamma,
    target_update_freq,
    update_type,
    soft_update_tau,
    batch_size,
    buffer_size,
    update_start_buffer_size,
    learning_rate,
    eps_min,
    eps_max,
    device_num,
    rand_seed,
    plot_option,
    model_save_path,
    trained_model_path
)

agent.train()


import multiprocessing
import pathlib
import sys
from typing import Union

import librosa
import numpy as np
import pyaudio

ourpath = pathlib.Path(__file__).parent

class InputConfig(object):
    min_seconds = 2
    sampling_rate = 44100
    stream_chunk_size = 1024
    input_channels = 1

easy_music = next(ourpath.joinpath('easy_music').glob('*.mp3'))
easy_bpm = int(open(str(next(ourpath.joinpath('easy_music').glob('*.txt')))).read())
normal_music = next(ourpath.joinpath('normal_music').glob('*.mp3'))
normal_bpm = int(open(str(next(ourpath.joinpath('normal_music').glob('*.txt')))).read())
hard_music = next(ourpath.joinpath('hard_music').glob('*.mp3'))
hard_bpm = int(open(str(next(ourpath.joinpath('hard_music').glob('*.txt')))).read())
menu_music = next(ourpath.joinpath('menu_music').glob('*.mp3'))
menu_bpm = int(open(str(next(ourpath.joinpath('menu_music').glob('*.txt')))).read())

def read_audio_file(file_path: Union[str, pathlib.Path], config: InputConfig):
    """
    Reads in an audio file and converts to a numpy array.
    >>> f = read_audio_file('./sw.mp3', InputConfig())
    >>> print(f)
    [[-1.2063791e-04 -1.0168436e-04 -3.2342479e-05 ... -7.7938515e-05
      -8.8288434e-05 -5.1688534e-05]
     [ 9.4297948e-06 -2.1362750e-05 -7.0384616e-05 ...  3.9314014e-05
       5.7227458e-05  3.8280330e-05]]

    """
    if not isinstance(file_path, str):
        file_path = str(file_path)
    min_samples = int(config.min_seconds * config.sampling_rate)
    try:
        audio_time_series, sampling_rate = librosa.load(
            file_path, sr=config.sampling_rate, mono=False
        )
        trimmed_audio, trim_indexes = librosa.effects.trim(audio_time_series)

        if trimmed_audio.shape[-1] < min_samples:
            padding = min_samples - len(trimmed_audio)
            offset = padding // 2
            trimmed_audio = np.pad(
                trimmed_audio, (offset, padding - offset), "constant"
            )
        return trimmed_audio
    except Exception as e:
        print(f"Exception while reading {file_path}: {e}", sys.stderr)
        return np.zeros(min_samples, dtype=np.float32)

def play_audio(array: np.ndarray, config: InputConfig, loop=False):
    """
    Plays a numpy array as audio.

    from: https://www.scivision.dev/playing-sounds-from-numpy-arrays-in-python/

    >>> i = InputConfig()
    >>> f = read_audio_file('./sw.mp3', i)
    >>> play_audio(f, i)

    # song: https://www.youtube.com/watch?v=d-nxW9qBtxQ

    """
    if len(array.shape) == 2:
        channels = int(array.shape[0])
    else:
        channels = 1

    array = np.transpose(array)
    p = pyaudio.PyAudio()
    stream = p.open(rate=config.sampling_rate, format=pyaudio.paFloat32, channels=channels, output=True)
    arrby = array.tobytes()
    stream.write(arrby)
    while loop:
        stream.write(arrby)
    stream.close()
    p.terminate()

def get_audio_thread(music: Union[str, pathlib.Path], loop=False):
    audio_cfg = InputConfig()
    np_file = read_audio_file(music, audio_cfg)
    audio_thread = multiprocessing.Process(target=play_audio, args=(np_file, audio_cfg, loop))
    return audio_thread