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

import os
import random
import logging
import sys
import numbers
import math
import sklearn
import datetime
import numpy as np
import cv2
from PIL import Image
from io import BytesIO

import mxnet as mx
from mxnet import ndarray as nd
from mxnet import io
from mxnet import recordio

import glob
import os.path as osp
from collections import defaultdict
import ipdb

logger = logging.getLogger()

class FaceImageIter(io.DataIter):

    def __init__(self, batch_size, data_shape,
                 path_imgrec = None,
                 shuffle=False, aug_list=None, mean = None,
                 rand_mirror=False, cutoff=0, color_jittering = 0,
                 data_name='data', label_name='softmax_label', **kwargs):
        super(FaceImageIter, self).__init__()

        path_to_data  = path_imgrec
        self.train_dir = path_to_data
        self.train_list = osp.join(path_to_data, 'list_train.txt')

        self._check_before_run()
        train_set, num_train_pids, num_train_imgs = self._process_dir(self.train_dir, self.train_list)

        if True:
            print("=> {} loaded".format(path_to_data))
            print("Dataset statistics:")
            print("  ------------------------------")
            print("  subset   | # ids | # images")
            print("  ------------------------------")
            print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_imgs))
            print("  ------------------------------")

        self.num_instances = 4
        self.train_set = train_set
        self.num_train_pids = num_train_pids
        self.index_dic = defaultdict(list)
        for index, (_, pid) in enumerate(self.train_set):
            self.index_dic[pid].append(index)
        self.pids = list(self.index_dic.keys())
        self.num_identities = len(self.pids)

        self.mean = mean
        self.nd_mean = None
        if self.mean:
          self.mean = np.array(self.mean, dtype=np.float32).reshape(1,1,3)
          self.nd_mean = mx.nd.array(self.mean).reshape((1,1,3))

        self.check_data_shape(data_shape)
        self.provide_data = [(data_name, (batch_size,) + data_shape)]
        self.batch_size = batch_size
        self.data_shape = data_shape
        self.shuffle = shuffle
        self.image_size = '%d,%d'%(data_shape[1],data_shape[2])
        self.rand_mirror = rand_mirror
        self.cutoff = cutoff
        self.color_jittering = color_jittering
        self.CJA = mx.image.ColorJitterAug(0.125, 0.125, 0.125)
        self.provide_label = [(label_name, (batch_size,))]
        #print(self.provide_label[0][1])
        self.cur = 0
        self.nbatch = 0
        self.is_init = False

    def _check_before_run(self):
        """Check if all files are available before going deeper"""
        if not osp.exists(self.train_dir):
            raise RuntimeError("'{}' is not available".format(self.train_dir))
        '''
        if not osp.exists(self.gallery_dir):
            raise RuntimeError("'{}' is not available".format(self.gallery_dir))
        if not osp.exists(self.probe_dir):
            raise RuntimeError("'{}' is not available".format(self.probe_dir))
        '''

    def _process_dir(self, dir_path, list_path):
        with open(list_path, 'r') as txt:
            lines = txt.readlines()
        dataset = []
        pid_container = set()
        for img_idx, img_info in enumerate(lines):
            img_path, pid = img_info.split(' ')
            pid = int(pid) # no need to relabel
            img_path = osp.join(dir_path, img_path)
            dataset.append((img_path, pid))
            pid_container.add(pid)
        num_imgs = len(dataset)
        num_pids = len(pid_container)
        # check if pid starts from 0 and increments with 1
        for idx, pid in enumerate(pid_container):
            assert idx == pid, "See code comment for explanation "+list_path+' %d'%pid
        return dataset, num_pids, num_imgs

    def reset_random(self):
        """Resets the iterator to the beginning of the data."""
        print('call reset()')
        self.cur = 0
        self.seq = self.imgidx
        if self.shuffle:
          random.shuffle(self.seq)
        if self.seq is None and self.imgrec is not None:
            self.imgrec.reset()

    def reset(self):
        """Resets the iterator to the beginning of the data."""
        print('call reset()')
        self.cur = 0
        indices = np.random.permutation(self.num_identities)
        self.seq = []
        for i in indices:
            pid = self.pids[i]
            t = self.index_dic[pid]
            replace = False if len(t) >= self.num_instances else True
            t = np.random.choice(t, size=self.num_instances, replace=replace)
            self.seq.extend(t)

    def num_samples(self):
      return len(self.train_set)

    def next_sample(self):
        """Helper function for reading in next sample."""
        while True:
            if self.cur >= len(self.seq):
                raise StopIteration
            idx = self.seq[self.cur]
            self.cur += 1
            if self.train_set is not None:
                img_path, label = self.train_set[idx]
                img = self.read_image(img_path)
                img = self.imdecode(img)
                return label, img, None, None

    def brightness_aug(self, src, x):
      alpha = 1.0 + random.uniform(-x, x)
      src *= alpha
      return src

    def contrast_aug(self, src, x):
      alpha = 1.0 + random.uniform(-x, x)
      coef = nd.array([[[0.299, 0.587, 0.114]]])
      gray = src * coef
      gray = (3.0 * (1.0 - alpha) / gray.size) * nd.sum(gray)
      src *= alpha
      src += gray
      return src

    def saturation_aug(self, src, x):
      alpha = 1.0 + random.uniform(-x, x)
      coef = nd.array([[[0.299, 0.587, 0.114]]])
      gray = src * coef
      gray = nd.sum(gray, axis=2, keepdims=True)
      gray *= (1.0 - alpha)
      src *= alpha
      src += gray
      return src

    def color_aug(self, img, x):
      #augs = [self.brightness_aug, self.contrast_aug, self.saturation_aug]
      #random.shuffle(augs)
      #for aug in augs:
      #  #print(img.shape)
      #  img = aug(img, x)
      #  #print(img.shape)
      #return img
      return self.CJA(img)

    def mirror_aug(self, img):
      _rd = random.randint(0,1)
      if _rd==1:
        for c in xrange(img.shape[2]):
          img[:,:,c] = np.fliplr(img[:,:,c])
      return img

    def compress_aug(self, img):
      buf = BytesIO()
      img = Image.fromarray(img.asnumpy(), 'RGB')
      q = random.randint(2, 20)
      img.save(buf, format='JPEG', quality=q)
      buf = buf.getvalue()
      img = Image.open(BytesIO(buf))
      return nd.array(np.asarray(img, 'float32'))

    def next(self):
        if not self.is_init:
          self.reset()
          self.is_init = True
        """Returns the next batch of data."""
        #print('in next', self.cur, self.labelcur)
        self.nbatch+=1
        batch_size = self.batch_size
        c, h, w = self.data_shape
        batch_data = nd.empty((batch_size, c, h, w))
        if self.provide_label is not None:
          batch_label = nd.empty(self.provide_label[0][1])
        i = 0
        try:
            while i < batch_size:
                label, _data, bbox, landmark = self.next_sample()
                if _data.shape[0]!=self.data_shape[1]:
                  _data = mx.image.resize_short(_data, self.data_shape[1])
                if self.rand_mirror:
                  _rd = random.randint(0,1)
                  if _rd==1:
                    _data = mx.ndarray.flip(data=_data, axis=1)
                if self.color_jittering>0:
                  if self.color_jittering>1:
                    _rd = random.randint(0,1)
                    if _rd==1:
                      _data = self.compress_aug(_data)
                  #print('do color aug')
                  _data = _data.astype('float32', copy=False)
                  #print(_data.__class__)
                  _data = self.color_aug(_data, 0.125)
                if self.cutoff>0:
                  _rd = random.randint(0,1)
                  if _rd==1:
                    #print('do cutoff aug', self.cutoff)
                    centerh = random.randint(0, _data.shape[0]-1)
                    centerw = random.randint(0, _data.shape[1]-1)
                    half = self.cutoff//2
                    starth = max(0, centerh-half)
                    endh = min(_data.shape[0], centerh+half)
                    startw = max(0, centerw-half)
                    endw = min(_data.shape[1], centerw+half)
                    #print(starth, endh, startw, endw, _data.shape)
                    _data[starth:endh, startw:endw, :] = 128
                if self.nd_mean is not None:
                  _data = _data.astype('float32', copy=False)
                  _data -= self.nd_mean
                  _data *= 0.0078125
                data = [_data]
                try:
                    self.check_valid_image(data)
                except RuntimeError as e:
                    logging.debug('Invalid image, skipping:  %s', str(e))
                    continue
                #print('aa',data[0].shape)
                #data = self.augmentation_transform(data)
                #print('bb',data[0].shape)
                for datum in data:
                    assert i < batch_size, 'Batch size must be multiples of augmenter output length'
                    #print(datum.shape)
                    batch_data[i][:] = self.postprocess_data(datum)
                    batch_label[i][:] = label
                    i += 1
        except StopIteration:
            if i<batch_size:
                raise StopIteration

        return io.DataBatch([batch_data], [batch_label], batch_size - i)

    def check_data_shape(self, data_shape):
        """Checks if the input data shape is valid"""
        if not len(data_shape) == 3:
            raise ValueError('data_shape should have length 3, with dimensions CxHxW')
        if not data_shape[0] == 3:
            raise ValueError('This iterator expects inputs to have 3 channels.')

    def check_valid_image(self, data):
        """Checks if the input data is valid"""
        if len(data[0].shape) == 0:
            raise RuntimeError('Data shape is wrong')

    def imdecode(self, s):
        """Decodes a string or byte string to an NDArray.
        See mx.img.imdecode for more details."""
        img = mx.image.imdecode(s) #mx.ndarray
        return img

    def read_image(self, fname):
        """Reads an input image `fname` and returns the decoded raw bytes.

        Example usage:
        ----------
        >>> dataIter.read_image('Face.jpg') # returns decoded raw bytes.
        """
        with open(fname, 'rb') as fin:
            img = fin.read()
        return img

    def augmentation_transform(self, data):
        """Transforms input data with specified augmentation."""
        for aug in self.auglist:
            data = [ret for src in data for ret in aug(src)]
        return data

    def postprocess_data(self, datum):
        """Final postprocessing step before image is loaded into the batch."""
        return nd.transpose(datum, axes=(2, 0, 1))

class FaceImageIterList(io.DataIter):
  def __init__(self, iter_list):
    assert len(iter_list)>0
    self.provide_data = iter_list[0].provide_data
    self.provide_label = iter_list[0].provide_label
    self.iter_list = iter_list
    self.cur_iter = None

  def reset(self):
    self.cur_iter.reset()

  def next(self):
    self.cur_iter = random.choice(self.iter_list)
    while True:
      try:
        ret = self.cur_iter.next()
      except StopIteration:
        self.cur_iter.reset()
        continue
      return ret

if __name__=='__main__':
  from config import config, default, generate_config

  generate_config('shuffse', 'faceid', 'arcface')
  path_to_data = config.dataset_path
  data_shape = (config.image_shape[2], config.image_shape[0], config.image_shape[1])

  train_dataiter = FaceImageIter(
          batch_size           = 32,
          data_shape           = data_shape,
          path_to_data         = path_to_data,
          shuffle              = True,
          rand_mirror          = config.data_rand_mirror,
          mean                 = None,
          cutoff               = config.data_cutoff,
          color_jittering      = config.data_color,
          images_filter        = config.data_images_filter,
  )

  batchx = train_dataiter.next()
  ipdb.set_trace()


import numpy as np
from jina import Document, DocumentArray, Executor, requests
from jina.logging.logger import JinaLogger

from executors.pn_encoder.pn import get_bottleneck_model

class PNEncoder(Executor):
    def __init__(self, ckpt_path: str, **kwargs):
        super().__init__(**kwargs)
        self.embedding_model = get_bottleneck_model(ckpt_path=ckpt_path)
        self.logger = JinaLogger('pn-encoder')

    @requests(on=['/index', '/search'])
    def encode(self, docs: DocumentArray, **kwargs):
        embeds = self.embedding_model.predict(np.stack(docs.get_attributes('blob')))

        for layer_embeds in embeds:
            for d, b in zip(docs, layer_embeds):
                d.chunks.append(Document(embedding=b))

        # set content to uri to reduce document size
        for d in docs:
            d.uri = d.tags['glb_path']


import cv2
import os
import json
import argparse
import tensorflow as tf
from tensorflow.keras.applications import vgg16, mobilenet, mobilenet_v2
import numpy as np
from glob import glob
from networks import SSD_VGG16, SSD_MOBILENET, SSD_MOBILENETV2
from utils import inference_utils, textboxes_utils, command_line_utils

parser = argparse.ArgumentParser(
    description='run inference on an input image.')
parser.add_argument('test_file', type=str, help='path to the test set file.')
parser.add_argument('images_dir', type=str, help='path to images dir.')
parser.add_argument('labels_dir', type=str, help='path to labels dir.')
parser.add_argument('config', type=str, help='path to config file.')
parser.add_argument('weights', type=str, help='path to config file.')
parser.add_argument('--label_maps', type=str, help='path to label maps file.')
parser.add_argument('--num_predictions', type=int,
                    help='the number of detections to be output as final detections', default=10)
parser.add_argument('--output_dir', type=str,
                    help='ouput', default="output")
args = parser.parse_args()

assert os.path.exists(args.config), "config file does not exist"
assert args.num_predictions > 0, "num_predictions must be larger than zero"
# assert args.confidence_threshold > 0, "confidence_threshold must be larger than zero."
# assert args.confidence_threshold <= 1, "confidence_threshold must be smaller than or equal to 1."

with open(args.config, "r") as config_file:
    config = json.load(config_file)

if not os.path.exists(args.output_dir):
    os.makedirs(args.output_dir)

input_size = config["model"]["input_size"]
model_config = config["model"]

if model_config["name"] == "ssd_mobilenetv2":
    model, process_input_fn, label_maps = inference_utils.inference_ssd_mobilenetv2(
        config, args)
elif model_config["name"] == "ssd_vgg16":
    model, process_input_fn, label_maps = inference_utils.ssd_vgg16(config, args)
else:
    print(
        f"model with name ${model_config['name']} has not been implemented yet")
    exit()

model.load_weights(args.weights)

with open(args.test_file, "r") as test_set_file:
    tests = test_set_file.readlines()
    for idx, sample in enumerate(tests):
        print(f"{idx+1}/{len(tests)}")
        image_file, label_file = sample.split(" ")
        filename = image_file[:image_file.index(".")]
        image = cv2.imread(os.path.join(args.images_dir, image_file))
        image = np.array(image, dtype=np.float)
        image = np.uint8(image)
        image_height, image_width, _ = image.shape
        height_scale, width_scale = input_size/image_height, input_size/image_width

        image = cv2.resize(image, (input_size, input_size))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = process_input_fn(image)

        image = np.expand_dims(image, axis=0)
        y_pred = model.predict(image)[0]

        with open(os.path.join(args.output_dir, f"{filename}.txt"), "w") as outfile:
            for i, pred in enumerate(y_pred):
                classname = label_maps[int(pred[0]) - 1].lower()
                confidence_score = pred[1]
                pred[[2, 4]] /= width_scale
                pred[[3, 5]] /= height_scale
                outfile.write(f"{classname} {confidence_score} {int(pred[2])} {int(pred[3])} {int(pred[4])} {int(pred[5])}\n")


import torch
from vision.ssd.vgg_ssd import create_vgg_ssd, create_vgg_ssd_predictor
from vision.ssd.mobilenetv1_ssd import create_mobilenetv1_ssd, create_mobilenetv1_ssd_predictor
from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite, create_mobilenetv1_ssd_lite_predictor
from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite, create_squeezenet_ssd_lite_predictor
from vision.datasets.voc_dataset import VOCDataset
from vision.datasets.open_images import OpenImagesDataset
from vision.datasets.coco_dataset import CocoDetection
from vision.utils import box_utils, measurements
from vision.utils.misc import str2bool, Timer
import argparse
import pathlib
import numpy as np
import logging
import sys
import os
from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite, create_mobilenetv2_ssd_lite_predictor

parser = argparse.ArgumentParser(description="SSD Evaluation on VOC Dataset.")
parser.add_argument('--net', default="vgg16-ssd",
                    help="The network architecture, it should be of mb1-ssd, mb1-ssd-lite, mb2-ssd-lite or vgg16-ssd.")
parser.add_argument("--trained_model", type=str)

parser.add_argument("--dataset_type", default="voc", type=str,
                    help='Specify dataset type. Currently support voc and open_images.')
parser.add_argument("--dataset", type=str, help="The root directory of the VOC dataset or Open Images dataset.")
parser.add_argument("--label_file", type=str, help="The label file path.")
parser.add_argument("--use_cuda", type=str2bool, default=True)
parser.add_argument("--use_2007_metric", type=str2bool, default=True)
parser.add_argument("--nms_method", type=str, default="hard")
parser.add_argument("--iou_threshold", type=float, default=0.5, help="The threshold of Intersection over Union.")
parser.add_argument("--eval_dir", default="eval_results", type=str, help="The directory to store evaluation results.")
parser.add_argument('--mb2_width_mult', default=1.0, type=float,
                    help='Width Multiplifier for MobilenetV2')
args = parser.parse_args()
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() and args.use_cuda else "cpu")

def group_annotation_by_class(dataset):
    true_case_stat = {}
    all_gt_boxes = {}
    all_difficult_cases = {}
    set_class = []
    for i in range(len(dataset)):
        image_id, annotation = dataset.get_annotation(i)
        gt_boxes, classes, is_difficult = annotation
        set_class.extend(classes.tolist())
        gt_boxes = torch.from_numpy(gt_boxes)
        for i, difficult in enumerate(is_difficult):
            class_index = int(classes[i])
            gt_box = gt_boxes[i]
            if not difficult:
                true_case_stat[class_index] = true_case_stat.get(class_index, 0) + 1

            if class_index not in all_gt_boxes:
                all_gt_boxes[class_index] = {}
            if image_id not in all_gt_boxes[class_index]:
                all_gt_boxes[class_index][image_id] = []
            all_gt_boxes[class_index][image_id].append(gt_box)
            if class_index not in all_difficult_cases:
                all_difficult_cases[class_index]={}
            if image_id not in all_difficult_cases[class_index]:
                all_difficult_cases[class_index][image_id] = []
            all_difficult_cases[class_index][image_id].append(difficult)

    for class_index in all_gt_boxes:
        for image_id in all_gt_boxes[class_index]:
            all_gt_boxes[class_index][image_id] = torch.stack(all_gt_boxes[class_index][image_id])
    for class_index in all_difficult_cases:
        for image_id in all_difficult_cases[class_index]:
            all_gt_boxes[class_index][image_id] = torch.tensor(all_gt_boxes[class_index][image_id])
    return true_case_stat, all_gt_boxes, all_difficult_cases

def coco_ap_per_class(num_true_cases, gt_boxes, difficult_cases,
                      prediction_file, use_2007_metric):
    aps, precs, recalls = [],[],[]
    _measurement_func = measurements.compute_voc2007_average_precision if use_2007_metric else measurements.compute_average_precision
    with open(prediction_file) as f:
        image_ids = []
        boxes = []
        scores = []
        for line in f:
            t = line.rstrip().split(" ")
            image_ids.append(t[0])
            scores.append(float(t[1]))
            box = torch.tensor([float(v) for v in t[2:]]).unsqueeze(0)
            box -= 1.0  # convert to python format where indexes start from 0
            boxes.append(box)
        scores = np.array(scores)
        sorted_indexes = np.argsort(-scores)
        boxes = [boxes[i] for i in sorted_indexes]
        image_ids = [image_ids[i] for i in sorted_indexes]

    max_ious = []
    max_args = []
    for i, image_id in enumerate(image_ids):
        box = boxes[i]
        image_id = int(image_id)
        # if image_id not in gt_boxes.keys():
        #     false_positive[i] = 1
        #     continue
        if image_id in gt_boxes.keys():
            gt_box = gt_boxes[image_id]
            ious = box_utils.iou_of(box, gt_box)
            max_iou = torch.max(ious).item()
            max_arg = torch.argmax(ious).item()
        else:
            max_iou = 0.0
            max_arg = 0
        max_ious.append(max_iou)
        max_args.append(max_arg)
    for _iou_thresh  in [0.5+0.05 * i for i in range(0,10)]:
        true_positive = np.zeros(len(image_ids))
        false_positive = np.zeros(len(image_ids))
        matched = set()
        for i,(max_iou,max_arg,image_id) in enumerate(zip(max_ious,max_args,image_ids)):
            image_id = int(image_id)
            if max_iou > _iou_thresh:
                if difficult_cases[image_id][max_arg] == 0:
                    if (image_id, max_arg) not in matched:
                        true_positive[i] = 1
                        matched.add((image_id, max_arg))
                    else:
                        false_positive[i] = 1
            else:
                false_positive[i] = 1

        true_positive = true_positive.cumsum()
        false_positive = false_positive.cumsum()
        precision = true_positive / (true_positive + false_positive + 1e-8 )
        recall = true_positive / num_true_cases

        _vtp = true_positive[-1] if len(true_positive)!=0 else 0
        _vfp = false_positive[-1] if len(false_positive)!=0 else 0
        ap = _measurement_func(precision,recall)
        _prec = _vtp / (_vtp+ _vfp + 1e-8)
        _recall =  _vtp/num_true_cases
        aps.append(ap)
        precs.append(_prec)
        recalls.append(_recall)
    return sum(aps)/len(aps), sum(precs)/len(precs), sum(recalls) / len(recalls)
def compute_average_precision_per_class(num_true_cases, gt_boxes, difficult_cases,
                                        prediction_file, iou_threshold, use_2007_metric):
    with open(prediction_file) as f:
        image_ids = []
        boxes = []
        scores = []
        for line in f:
            t = line.rstrip().split(" ")
            image_ids.append(t[0])
            scores.append(float(t[1]))
            box = torch.tensor([float(v) for v in t[2:]]).unsqueeze(0)
            box -= 1.0  # convert to python format where indexes start from 0
            boxes.append(box)
        scores = np.array(scores)
        sorted_indexes = np.argsort(-scores)
        boxes = [boxes[i] for i in sorted_indexes]
        image_ids = [image_ids[i] for i in sorted_indexes]
        true_positive = np.zeros(len(image_ids))
        false_positive = np.zeros(len(image_ids))
        matched = set()
        for i, image_id in enumerate(image_ids):
            box = boxes[i]

            if args.dataset_type=="voc":
                pass
            else:
                image_id = int(image_id)

            if image_id not in gt_boxes.keys():
                false_positive[i] = 1
                continue

            gt_box = gt_boxes[image_id]
            ious = box_utils.iou_of(box, gt_box)
            max_iou = torch.max(ious).item()
            max_arg = torch.argmax(ious).item()
            if max_iou > iou_threshold:
                if difficult_cases[image_id][max_arg] == 0:
                    if (image_id, max_arg) not in matched:
                        true_positive[i] = 1
                        matched.add((image_id, max_arg))
                    else:
                        false_positive[i] = 1
            else:
                false_positive[i] = 1

    true_positive = true_positive.cumsum()
    false_positive = false_positive.cumsum()
    precision = true_positive / (true_positive + false_positive + 1e-8 )
    recall = true_positive / num_true_cases
    # import pdb;pdb.set_trace()
    _vtp = true_positive[-1] if len(true_positive)!=0 else 0
    _vfp = false_positive[-1] if len(false_positive)!=0 else 0
    if use_2007_metric:
        return measurements.compute_voc2007_average_precision(precision, recall), _vtp / (_vtp+ _vfp + 1e-8), _vtp/num_true_cases
    else:
        return measurements.compute_average_precision(precision, recall)

if __name__ == '__main__':
    # eval_path = pathlib.Path(args.eval_dir)
    # eval_path.mkdir(exist_ok=True)
    if not os.path.exists(args.eval_dir):
        os.mkdir(args.eval_dir)
    timer = Timer()
    class_names = [name.strip() for name in open(args.label_file).readlines()]

    if args.dataset_type == "voc":
        dataset = VOCDataset(args.dataset, is_test=True)
    elif args.dataset_type == 'open_images':
        dataset = OpenImagesDataset(args.dataset, dataset_type="test")
    elif args.dataset_type == "coco":
        # dataset = CocoDetection("/home/wenyen4desh/datasets/coco/test2017","/home/wenyen4desh/datasets/annotations/image_info_test2017.json")
        dataset = CocoDetection("../../dataset/datasets/coco/val2017","../../dataset/datasets/coco/annotations/instances_val2017.json")
        # dataset = CocoDetection("/home/wenyen4desh/datasets/coco/train2017","/home/wenyen4desh/datasets/coco/annotations/instances_train2017.json")

    true_case_stat, all_gb_boxes, all_difficult_cases = group_annotation_by_class(dataset)
    if args.net == 'vgg16-ssd':
        net = create_vgg_ssd(len(class_names), is_test=True)
    elif args.net == 'mb1-ssd':
        net = create_mobilenetv1_ssd(len(class_names), is_test=True)
    elif args.net == 'mb1-ssd-lite':
        net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
    elif args.net == 'sq-ssd-lite':
        net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
    elif args.net == 'mb2-ssd-lite':
        net = create_mobilenetv2_ssd_lite(len(class_names), width_mult=args.mb2_width_mult, is_test=True)
    else:
        logging.fatal("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
        parser.print_help(sys.stderr)
        sys.exit(1)

    #train_transform = MatchPrior(config.priors, config.center_variance,
    #                              config.size_variance, 0.5)

    #test_transform = TestTransform(config.image_size, config.image_mean, config.image_std)

    timer.start("Load Model")
    net.load(args.trained_model)
    net = net.to(DEVICE)
    print('It took {} seconds to load the model.'.format(timer.end("Load Model")))

    if args.net == 'vgg16-ssd':
        predictor = create_vgg_ssd_predictor(net, nms_method=args.nms_method, device=DEVICE)
    elif args.net == 'mb1-ssd':
        predictor = create_mobilenetv1_ssd_predictor(net, nms_method=args.nms_method, device=DEVICE)
    elif args.net == 'mb1-ssd-lite':
        predictor = create_mobilenetv1_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
    elif args.net == 'sq-ssd-lite':
        predictor = create_squeezenet_ssd_lite_predictor(net,nms_method=args.nms_method, device=DEVICE)
    elif args.net == 'mb2-ssd-lite':
        predictor = create_mobilenetv2_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
    else:
        logging.fatal("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
        parser.print_help(sys.stderr)
        sys.exit(1)

    results = []
    # Predict Bounding Box
    for i in range(len(dataset)):
        print("process image", i)
        timer.start("Load Image")
        image = dataset.get_image(i)
        print("Load Image: {:4f} seconds.".format(timer.end("Load Image")))
        timer.start("Predict")
        boxes, labels, probs = predictor.predict(image)
        print("Prediction: {:4f} seconds.".format(timer.end("Predict")))
        indexes = torch.ones(labels.size(0), 1, dtype=torch.float32) * i
        results.append(torch.cat([
            indexes.reshape(-1, 1),
            labels.reshape(-1, 1).float(),
            probs.reshape(-1, 1),
            boxes + 1.0  # matlab's indexes start from 1
        ], dim=1))
    results = torch.cat(results)

    # Write the result to file
    for class_index, class_name in enumerate(class_names):
        if class_index == 0: continue  # ignore background
        file_name = "det_test_{}.txt".format(class_name)
        prediction_path = os.path.join(args.eval_dir,file_name)
        with open(prediction_path, "w") as f:
            sub = results[results[:, 1] == class_index, :]
            for i in range(sub.size(0)):
                prob_box = sub[i, 2:].numpy()
                image_id,_ = dataset.get_annotation(int(sub[i, 0]))
                f.write(str(image_id) + " " + " ".join([str(v) for v in prob_box])+"\n")
                # image_id = dataset.ids[int(sub[i, 0])]
                # print(str(image_id) + " " + " ".join([str(v) for v in prob_box]), file=f)

    aps = []
    prcs = []
    recalls = []
    print("\n\nAverage Precision Per-class:")
    for class_index, class_name in enumerate(class_names):
        if class_index == 0:
            continue
        file_name = "det_test_{}.txt".format(class_name)
        prediction_path = os.path.join(args.eval_dir,file_name)
        # Pascal@0.5 evaluation method
        ap, precision, recall = compute_average_precision_per_class(
            true_case_stat[class_index],
            all_gb_boxes[class_index],
            all_difficult_cases[class_index],
            prediction_path,
            args.iou_threshold,
            args.use_2007_metric
        )

        # # COCO eval

        # ap, precision, recall = coco_ap_per_class(
        #     true_case_stat[class_index],
        #     all_gb_boxes[class_index],
        #     all_difficult_cases[class_index],
        #     prediction_path,
        #     args.use_2007_metric
        # )

        aps.append(ap)
        prcs.append(precision)
        recalls.append(recall)
        print("{}: {}".format(class_name,ap))

    print("\nAverage Precision Across All Classes:{}".format(sum(aps)/len(aps)))
    print("\nAverage Precision :{}".format(sum(prcs)/len(prcs)))
    print("\nAverage Recall :{}".format(sum(recalls)/len(recalls)))


# -*- coding: utf-8 -*-
import tensorflow as tf
import numpy as np
from data_utils import *
from textsum_model import Neuralmodel
from gensim.models import KeyedVectors
from rouge import Rouge
import os
import math
import pickle
from tqdm import tqdm

#configuration
FLAGS=tf.app.flags.FLAGS

tf.app.flags.DEFINE_string("log_path","../log/","path of summary log.")
tf.app.flags.DEFINE_string("tra_data_path","../src/neuralsum/dailymail/tra/","path of training data.")
tf.app.flags.DEFINE_string("tst_data_path","../src/neuralsum/dailymail/tst/","path of test data.")
tf.app.flags.DEFINE_string("val_data_path","../src/neuralsum/dailymail/val/","path of validation data.")
tf.app.flags.DEFINE_string("vocab_path","../cache/vocab","path of vocab frequency list")
tf.app.flags.DEFINE_integer("vocab_size",199900,"maximum vocab size.")

tf.app.flags.DEFINE_float("learning_rate",0.0001,"learning rate")

tf.app.flags.DEFINE_integer("is_frozen_step", 400, "how many steps before fine-tuning the embedding.")
tf.app.flags.DEFINE_integer("decay_step", 5000, "how many steps before decay learning rate.")
tf.app.flags.DEFINE_float("decay_rate", 0.1, "Rate of decay for learning rate.")
tf.app.flags.DEFINE_string("ckpt_dir","../ckpt/","checkpoint location for the model")
tf.app.flags.DEFINE_integer("batch_size", 20, "Batch size for training/evaluating.")
tf.app.flags.DEFINE_integer("embed_size", 150,"embedding size")
tf.app.flags.DEFINE_integer("input_y2_max_length", 40,"the max length of a sentence in abstracts")
tf.app.flags.DEFINE_integer("max_num_sequence", 30,"the max number of sequence in documents")
tf.app.flags.DEFINE_integer("max_num_abstract", 4,"the max number of abstract in documents")
tf.app.flags.DEFINE_integer("sequence_length", 100,"the max length of a sentence in documents")
tf.app.flags.DEFINE_integer("hidden_size", 300,"the hidden size of the encoder and decoder")
tf.app.flags.DEFINE_boolean("use_highway_flag", True,"using highway network or not.")
tf.app.flags.DEFINE_integer("highway_layers", 1,"How many layers in highway network.")
tf.app.flags.DEFINE_integer("document_length", 1000,"the max vocabulary of documents")
tf.app.flags.DEFINE_integer("beam_width", 4,"the beam search max width")
tf.app.flags.DEFINE_integer("attention_size", 150,"the attention size of the decoder")
tf.app.flags.DEFINE_boolean("extract_sentence_flag", True,"using sentence extractor")
tf.app.flags.DEFINE_boolean("is_training", True,"is traning.true:tranining,false:testing/inference")
tf.app.flags.DEFINE_integer("num_epochs",10,"number of epochs to run.")
tf.app.flags.DEFINE_integer("validate_every", 1, "Validate every validate_every epochs.")
tf.app.flags.DEFINE_boolean("use_embedding",True,"whether to use embedding or not.")
tf.app.flags.DEFINE_string("word2vec_model_path","../w2v/benchmark_sg1_e150_b.vector","word2vec's vocabulary and vectors")
filter_sizes = [1,2,3,4,5,6,7]
feature_map = [20,20,30,40,50,70,70]
cur_learning_steps = [500,2500]

def main(_):
    config = tf.ConfigProto()
    config.gpu_options.allow_growth=True
    with tf.Session(config=config) as sess:
        # instantiate model
        Model = Neuralmodel(FLAGS.extract_sentence_flag, FLAGS.is_training, FLAGS.vocab_size, FLAGS.batch_size, FLAGS.embed_size, FLAGS.learning_rate, cur_learning_steps, FLAGS.decay_step, FLAGS.decay_rate, FLAGS.max_num_sequence, FLAGS.sequence_length,
                            filter_sizes, feature_map, FLAGS.use_highway_flag, FLAGS.highway_layers, FLAGS.hidden_size, FLAGS.document_length, FLAGS.max_num_abstract, FLAGS.beam_width, FLAGS.attention_size, FLAGS.input_y2_max_length)
        # initialize saver
        saver = tf.train.Saver()
        if os.path.exists(FLAGS.ckpt_dir+"checkpoint"):
            print("Restoring Variables from Checkpoint.")
            saver.restore(sess,tf.train.latest_checkpoint(FLAGS.ckpt_dir))
            summary_writer = tf.summary.FileWriter(logdir=FLAGS.log_path, graph=sess.graph)
        else:
            print('Initializing Variables')
            sess.run(tf.global_variables_initializer())
            summary_writer = tf.summary.FileWriter(logdir=FLAGS.log_path, graph=sess.graph)
            if FLAGS.use_embedding: #load pre-trained word embedding
                assign_pretrained_word_embedding(sess, FLAGS.vocab_path, FLAGS.vocab_size, Model,FLAGS.word2vec_model_path)
        curr_epoch=sess.run(Model.epoch_step)

        batch_size=FLAGS.batch_size
        iteration=0
        for epoch in range(curr_epoch,FLAGS.num_epochs):
            loss, counter =  0.0, 0
            train_gen = Batch(FLAGS.tra_data_path,FLAGS.vocab_path,FLAGS.batch_size,FLAGS)
            for batch in tqdm(train_gen):
                iteration=iteration+1
                if epoch==0 and counter==0:
                    print("train_batch", batch['abstracts_len'])
                feed_dict={}
                if FLAGS.extract_sentence_flag:
                    feed_dict[Model.dropout_keep_prob] = 0.5
                    feed_dict[Model.input_x] = batch['article_words']
                    feed_dict[Model.input_y1] = batch['label_sentences']
                    feed_dict[Model.input_y1_length] = batch['article_len']
                    feed_dict[Model.tst] = FLAGS.is_training
                    feed_dict[Model.cur_learning] = True if cur_learning_steps[1] > iteration and epoch == 0 else False
                else:
                    feed_dict[Model.dropout_keep_prob] = 0.5
                    feed_dict[Model.input_x] = batch['article_words']
                    feed_dict[Model.input_y2_length] = batch['abstracts_len']
                    feed_dict[Model.input_y2] = batch['abstracts_inputs']
                    feed_dict[Model.input_decoder_x] = batch['abstracts_targets']
                    feed_dict[Model.value_decoder_x] = batch['article_value']
                    feed_dict[Model.tst] = FLAGS.is_training
                train_op = Model.train_op_frozen if FLAGS.is_frozen_step > iteration and epoch == 0 else Model.train_op
                curr_loss,lr,_,_,summary,logits=sess.run([Model.loss_val,Model.learning_rate,train_op,Model.global_increment,Model.merge,Model.logits],feed_dict)
                summary_writer.add_summary(summary, global_step=iteration)
                loss,counter=loss+curr_loss,counter+1
                if counter %50==0:
                    print("Epoch %d\tBatch %d\tTrain Loss:%.3f\tLearning rate:%.5f" %(epoch,counter,loss/float(counter),lr))
                if iteration % 1000 == 0:
                    eval_loss = do_eval(sess, Model)
                    print("Epoch %d Validation Loss:%.3f\t " % (epoch, eval_loss))
                    # TODO eval_loss, acc_score = do_eval(sess, Model)
                    # TODO print("Epoch %d Validation Loss:%.3f\t Acc:%.3f" % (epoch, eval_loss, acc_score))
                    # save model to checkpoint
                    save_path = FLAGS.ckpt_dir + "model.ckpt"
                    saver.save(sess, save_path, global_step=epoch)
            #epoch increment
            print("going to increment epoch counter....")
            sess.run(Model.epoch_increment)
            print(epoch,FLAGS.validate_every,(epoch % FLAGS.validate_every==0))
            if epoch % FLAGS.validate_every==0:
                #save model to checkpoint
                save_path=FLAGS.ckpt_dir+"model.ckpt"
                saver.save(sess,save_path,global_step=epoch)
        summary_writer.close()

def do_eval(sess, Model):
    eval_loss, eval_counter= 0.0, 0
    # eval_loss, eval_counter, acc_score= 0.0, 0, 0.0
    batch_size = 20
    valid_gen = Batch(FLAGS.tst_data_path,FLAGS.vocab_path,batch_size,FLAGS)
    for batch in valid_gen:
        feed_dict={}
        if FLAGS.extract_sentence_flag:
            feed_dict[Model.dropout_keep_prob] = 1.0
            feed_dict[Model.input_x] = batch['article_words']
            feed_dict[Model.input_y1] = batch['label_sentences']
            feed_dict[Model.input_y1_length] = batch['article_len']
            feed_dict[Model.tst] = not FLAGS.is_training
            feed_dict[Model.cur_learning] = False
        else:
            feed_dict[Model.dropout_keep_prob] = 1.0
            feed_dict[Model.input_x] = batch['article_words']
            feed_dict[Model.input_y2] = batch['abstracts_inputs']
            feed_dict[Model.input_y2_length] = batch['abstracts_len']
            feed_dict[Model.input_decoder_x] = batch['abstracts_targets']
            feed_dict[Model.value_decoder_x] = batch['article_value']
            feed_dict[Model.tst] = not FLAGS.is_training
        curr_eval_loss,logits=sess.run([Model.loss_val,Model.logits],feed_dict)
        # curr_acc_score = compute_label(logits, batch)
        # acc_score += curr_acc_score
        eval_loss += curr_eval_loss
        eval_counter += 1

    return eval_loss/float(eval_counter) # acc_score/float(eval_counter)

def compute_label(logits, batch): # TODO
    imp_pos = np.argsort(logits)
    lab_num = [ len(res['label']) for res in batch['original']]
    lab_pos = [ res['label'] for res in batch['original']]
    abs_num = [ res['abstract'] for res in batch['original']]
    sen_pos = [ pos[:num] for pos, num in zip(imp_pos, lab_num)]

    # compute
    acc_list = []
    for sen, lab, abst in zip(sen_pos, lab_pos, abs_num):
        sen = set(sen)
        lab = set(lab)
        if len(lab) == 0 or len(abst) == 0:
            continue
        score = float(len(sen&lab)) / len(abst)
        acc = 1.0 if score > 1.0 else score
        acc_list.append(acc)
    acc_score = np.mean(acc_list)

    return acc_score

def assign_pretrained_word_embedding(sess,vocab_path,vocab_size,Model,word2vec_model_path):
    print("using pre-trained word emebedding.started.word2vec_model_path:",word2vec_model_path)
    vocab = Vocab(vocab_path, vocab_size)
    word2vec_model = KeyedVectors.load_word2vec_format(word2vec_model_path, binary=True)
    bound = np.sqrt(6.0) / np.sqrt(vocab_size)  # bound for random variables.
    count_exist = 0;
    count_not_exist = 0
    word_embedding_2dlist = [[]] * vocab_size  # create an empty word_embedding list.
    word_embedding_2dlist[0] = np.zeros(FLAGS.embed_size, dtype=np.float32)  # assign empty for first word:'PAD'
    for i in range(1, vocab_size):  # loop each word
        word = vocab.id2word(i)
        embedding = None
        try:
            embedding = word2vec_model[word]  # try to get vector:it is an array.
        except Exception:
            embedding = None
        if embedding is not None:  # the 'word' exist a embedding
            word_embedding_2dlist[i] = embedding;
            count_exist = count_exist + 1  # assign array to this word.
        else:  # no embedding for this word
            word_embedding_2dlist[i] = np.random.uniform(-bound, bound, FLAGS.embed_size)
            count_not_exist = count_not_exist + 1  # init a random value for the word.
    word_embedding_final = np.array(word_embedding_2dlist)  # covert to 2d array.
    word_embedding = tf.constant(word_embedding_final, dtype=tf.float32)  # convert to tensor
    t_assign_embedding = tf.assign(Model.Embedding,word_embedding)  # assign this value to our embedding variables of our model.
    sess.run(t_assign_embedding)

    word_embedding_2dlist_ = [[]] * 2  # create an empty word_embedding list for GO END.
    word_embedding_2dlist_[0] = np.random.uniform(-bound, bound, FLAGS.hidden_size) # GO
    word_embedding_2dlist_[1] = np.random.uniform(-bound, bound, FLAGS.hidden_size) # END
    word_embedding_final_ = np.array(word_embedding_2dlist_)  # covert to 2d array.
    word_embedding_ = tf.constant(word_embedding_final_, dtype=tf.float32)  # convert to tensor
    t_assign_embedding_ = tf.assign(Model.Embedding_,word_embedding_)  # assign this value to our embedding variables of our model.
    sess.run(t_assign_embedding_)
    print("word. exists embedding:", count_exist, " ;word not exist embedding:", count_not_exist)
    print("using pre-trained word emebedding.ended...")

if __name__ == "__main__":
    tf.app.run()


import os, torch
import numpy as np
from torch import nn
from torch.nn import functional as F

class SuperResolutionGenerator:
    def __init__(self):
        pass
    def SetTorchFilePath(self,path):
        self.pretrained_net = torch.load(path)
        self.pretrained_net.eval()
        torch.set_grad_enabled(False)
        torch.cuda.empty_cache()

    def SetMeanMax(self,mean,max):
        self.meamVal = mean
        self.maxVal = max

    def Generate(self,img):
        torch.cuda.empty_cache()
        torch.set_grad_enabled(False)
        lrImg = np.array(img, dtype=np.float32)
        lrImg = (lrImg - self.meamVal) / (self.maxVal / 2)
        lrImg = np.expand_dims(lrImg, axis=0)
        lrImg = np.expand_dims(lrImg, axis=0)
        lrImg = torch.from_numpy(lrImg).float()
        lrImg = lrImg.to(torch.device('cuda:1'))
        pre2 = self.pretrained_net(lrImg)
        saveImg = pre2.cpu().data.numpy()[0, 0, :, :, :]
        saveImg *= (self.maxVal / 2)
        saveImg += self.meamVal
        saveImg = np.uint8(np.maximum(np.minimum(saveImg, 255),0))
        return saveImg


from tensorflow.keras.applications.imagenet_utils import decode_predictions
from application.initializer import LoggerInstance
from PIL import Image
import numpy as np
import tensorflow as tf
import ssl

ssl._create_default_https_context = ssl._create_unverified_context
logger = LoggerInstance().get_logger(__name__)
classifier_model = None

class InferenceTask:

    @staticmethod
    async def load_model(classifier_model_name):
        if classifier_model_name == "MOBILENET_V2":
            model = tf.keras.applications.MobileNetV2(weights="imagenet")
        elif classifier_model_name == "INCEPTION_V3":
            model = tf.keras.applications.InceptionV3(weights="imagenet")
        else:
            model = tf.keras.applications.MobileNetV2(weights="imagenet")

        return model

    async def predict(self, classifier_model_name, image: Image.Image, shape):
        global classifier_model
        if classifier_model is None:
            classifier_model = await self.load_model(classifier_model_name)
        image = np.asarray(image.resize(shape))
        image = np.array(image)[..., :3] / 255.0
        image = np.expand_dims(image, 0)
        result = classifier_model.predict(image, 2)
        result = decode_predictions(
                result, top=5
        )[0]
        response = []
        for i, res in enumerate(result):
            resp = dict()
            resp["class"] = res[1]
            resp["confidence"] = f"{res[2] * 100:0.2f} %"
            response.append(resp)
        return response


import usb.core
import usb.util
import numpy
import logging
from .classes import AbstractPush2Section
from .exceptions import Push2USBDeviceConfigurationError, Push2USBDeviceNotFound
from .constants import ABLETON_VENDOR_ID, PUSH2_PRODUCT_ID, USB_TRANSFER_TIMEOUT, DISPLAY_FRAME_HEADER, \
    DISPLAY_BUFFER_SIZE, DISPLAY_FRAME_XOR_PATTERN, DISPLAY_N_LINES, DISPLAY_LINE_PIXELS, DISPLAY_LINE_FILLER_BYTES, \
    FRAME_FORMAT_BGR565, FRAME_FORMAT_RGB565, FRAME_FORMAT_RGB

NP_DISPLAY_FRAME_XOR_PATTERN = numpy.array(DISPLAY_FRAME_XOR_PATTERN, dtype=numpy.uint16)  # Numpy array version of the constant

# Utils for image format processing

def get_bit(value, n):
    """Get bit value at binary position from value"""
    return ((value >> n & 1) != 0)

def set_bit(value, n):
    """Set bit to True at given binary position in value"""
    return value | (1 << n)

def build_rgb565_to_bgr565_map():
    """Creates a dictionary with 2^16 items to allow easy conversion from rgb565 format to bgr565 format"""
    rgb565_to_bgr565_map = dict()
    rgb565_to_bgr565_bit_shifting_table = (
        (0, 11),  # r 1
        (1, 12),  # r 2
        (2, 13),  # r 3
        (3, 14),  # r 4
        (4, 15),  # r 5
        (11, 0),  # b 1
        (12, 1),  # b 2
        (13, 2),  # b 3
        (14, 3),  # b 4
        (15, 4),  # b 5
        (5, 5),   # g 1
        (6, 6),   # g 2
        (7, 7),   # g 3
        (8, 8),   # g 4
        (9, 9),   # g 5
        (10, 10),  # g 6
    )
    for rgb565_value in range(0, pow(2, 16)):
        bgr565_value = 0
        for i, j in rgb565_to_bgr565_bit_shifting_table:
            if get_bit(rgb565_value, i):
                bgr565_value = set_bit(bgr565_value, j)
        rgb565_to_bgr565_map[rgb565_value] = int(bgr565_value)
    return rgb565_to_bgr565_map

# Vectorized function to convert from rgb565 to bgr565
rgb565_to_bgr565_map = build_rgb565_to_bgr565_map()
rgb565_to_bgr565_vecotrized = numpy.vectorize(
    lambda x: rgb565_to_bgr565_map[x], otypes=[numpy.uint16])

# Non-vectorized function for converting from rgb to bgr565
def rgb_to_bgr565(rgb_frame):
    # RGB is defined here as an 2d array with (r, g, b) tuples with values between [0.0, 1.0]
    # NOTE: this is really slow, should only be used offline
    out_array = numpy.zeros(shape=(rgb_frame.shape[1], rgb_frame.shape[0]), dtype=numpy.uint16)
    for i in range(0, rgb_frame.shape[0]):
        for j in range(0, rgb_frame.shape[1]):
            r = rgb_frame[i][j][0]
            g = rgb_frame[i][j][1]
            b = rgb_frame[i][j][2]
            r = int(round(r * (pow(2, 5) - 1)))
            g = int(round(g * (pow(2, 6) - 1)))
            b = int(round(b * (pow(2, 5) - 1)))
            out_array[j][i] = int(
                '{b:05b}{g:06b}{r:05b}'.format(r=r, g=g, b=b), 2)
    return out_array

class Push2Display(AbstractPush2Section):
    """Class to interface with Ableton's Push2 display.
    See https://github.com/Ableton/push-interface/blob/master/doc/AbletonPush2MIDIDisplayInterface.asc#display-interface
    """
    usb_endpoint = None
    last_prepared_frame = None

    def configure_usb_device(self):
        """Connect to Push2 USB device and get the Endpoint object used to send data
        to Push2's display.
        See https://github.com/Ableton/push-interface/blob/master/doc/AbletonPush2MIDIDisplayInterface.asc#31-usb-display-interface-access
        """
        usb_device = None
        try:
            usb_device = usb.core.find(
                idVendor=ABLETON_VENDOR_ID, idProduct=PUSH2_PRODUCT_ID)
        except usb.core.NoBackendError:
            logging.error('No backend is available for pyusb. Please make sure \'libusb\' is installed in your system.')

        if usb_device is None:
            raise Push2USBDeviceNotFound

        device_configuration = usb_device.get_active_configuration()
        if device_configuration is None:
            usb_device.set_configuration()

        interface = device_configuration[(0, 0)]
        out_endpoint = usb.util.find_descriptor(
            interface,
            custom_match=lambda e:
            usb.util.endpoint_direction(e.bEndpointAddress) ==
            usb.util.ENDPOINT_OUT)

        if out_endpoint is None:
            raise Push2USBDeviceConfigurationError

        self.usb_endpoint = out_endpoint

    def prepare_frame(self, frame, input_format=FRAME_FORMAT_BGR565):
        """Prepare the given image frame to be shown in the Push2's display.
        Depending on the input_format argument, "frame" must be a numpy array with the following characteristics:

        * for FRAME_FORMAT_BGR565: numpy array of shape 910x160 and of uint16. Each uint16 element specifies rgb
          color with the following bit position meaning: [b4 b3 b2 b1 b0 g5 g4 g3 g2 g1 g0 r4 r3 r2 r1 r0].

        * for FRAME_FORMAT_RGB565: numpy array of shape 910x160 and of uint16. Each uint16 element specifies rgb
          color with the following bit position meaning: [r4 r3 r2 r1 r0 g5 g4 g3 g2 g1 g0 b4 b3 b2 b1 b0].

        * for FRAME_FORMAT_RGB: numpy array of shape 910x160x3 with the third dimension representing rgb colors
          with separate float values for rgb channels (float values in range [0.0, 1.0]).

        Preferred format is brg565 as it requires no conversion before sending to Push2. Using brg565 it should be
        possible to achieve frame rates as high as 36fps. With rgb565 the conversion slows down the process but should
        still allow frame rates of 14fps. Sending data in rgb will result in very long frame preparation times that can
        take seconds. Therefore rgb format should only be used for displaying static images that are prepared offline.
        "prepare_frame" method expects numpy array elements to be big endian.
        In addition to format changing (if needed), "prepare_frame" prepares the frame to be sent to push by adding
        filler bytes and performing bitwise XOR as decribed in the Push2 specification.
        See https://github.com/Ableton/push-interface/blob/master/doc/AbletonPush2MIDIDisplayInterface.asc#326-allocating-libusb-transfers
        """

        assert input_format in [FRAME_FORMAT_BGR565, FRAME_FORMAT_RGB565, FRAME_FORMAT_RGB], 'Invalid frame format'

        if input_format == FRAME_FORMAT_RGB:
            # If format is rgb, do conversion before the rest as frame must be reshaped
            # from (w, h, 3) to (w, h)
            frame = rgb_to_bgr565(frame)

        assert type(frame) == numpy.ndarray
        assert frame.dtype == numpy.dtype('uint16')
        assert frame.shape[0] == DISPLAY_LINE_PIXELS, 'Wrong number of pixels in line ({0})'.format(
            frame.shape[0])
        assert frame.shape[1] == DISPLAY_N_LINES, 'Wrong number of lines in frame ({0})'.format(
            frame.shape[1])

        width = DISPLAY_LINE_PIXELS + DISPLAY_LINE_FILLER_BYTES // 2
        height = DISPLAY_N_LINES
        prepared_frame = numpy.zeros(shape=(width, height), dtype=numpy.uint16)
        prepared_frame[0:frame.shape[0], 0:frame.shape[1]] = frame
        prepared_frame = prepared_frame.transpose().flatten()
        if input_format == FRAME_FORMAT_RGB565:
            prepared_frame = rgb565_to_bgr565_vecotrized(prepared_frame)
        elif input_format == FRAME_FORMAT_BGR565:
            pass  # Nothing to do as this is already the requested format
        elif input_format == FRAME_FORMAT_RGB:
            pass  # Nothing as conversion was done before
        prepared_frame = prepared_frame.byteswap()  # Change to little endian
        prepared_frame = numpy.bitwise_xor(prepared_frame, NP_DISPLAY_FRAME_XOR_PATTERN)

        self.last_prepared_frame = prepared_frame
        return prepared_frame.byteswap().tobytes()

    def send_to_display(self, prepared_frame):
        """Sends a prepared frame to Push2 display.
        First sends frame header and then sends prepared_frame in buffers of BUFFER_SIZE.
        'prepared_frame' must be a flattened array of (DISPLAY_LINE_PIXELS + (DISPLAY_LINE_FILLER_BYTES // 2)) * DISPLAY_N_LINES 16bit BGR 565 values
        as returned by the 'Push2Display.prepare_frame' method.
        See https://github.com/Ableton/push-interface/blob/master/doc/AbletonPush2MIDIDisplayInterface.asc#326-allocating-libusb-transfers
        """

        if self.usb_endpoint is not None:
            self.usb_endpoint.write(
                DISPLAY_FRAME_HEADER, USB_TRANSFER_TIMEOUT)

            self.usb_endpoint.write(prepared_frame, USB_TRANSFER_TIMEOUT)
            for i in range(0, len(prepared_frame), DISPLAY_BUFFER_SIZE):
                buffer_data = prepared_frame[i: i + DISPLAY_BUFFER_SIZE]
                self.usb_endpoint.write(buffer_data, USB_TRANSFER_TIMEOUT)

    def display_frame(self, frame, input_format=FRAME_FORMAT_BGR565):
        self.send_to_display(self.prepare_frame(frame, input_format=input_format))

    def display_prepared_frame(self, prepared_frame):
        self.send_to_display(prepared_frame)

    def display_last_frame(self):
        self.send_to_display(self.last_prepared_frame)


import numpy as np
import tensorflow as tf
import os, argparse
from jamo import hangul_to_jamo
from models.tacotron import Tacotron
from util.text import text_to_sequence, sequence_to_text
from util.plot_alignment import plot_alignment
from util.hparams import *

sentences = [
  '정말로 사랑한담 기다려주세요'
]

checkpoint_dir = './checkpoint/1'

class Synthesizer:
    def load(self, step):
        enc_input = tf.placeholder(tf.int32, [1, None])
        sequence_length = tf.placeholder(tf.int32, [1])
        dec_input = tf.placeholder(tf.float32, [1, None, mel_dim])

        self.model = Tacotron()
        self.model.initialize(enc_input, sequence_length, dec_input)

        self.enc_input = self.model.enc_input[0]
        self.sequence_length = self.model.sequence_length[0]
        self.mel_output = self.model.mel_output[0]
        self.alignment = self.model.alignment[0]

        self.session = tf.Session()
        self.session.run(tf.global_variables_initializer())
        saver = tf.train.Saver()
        saver.restore(self.session, os.path.join(checkpoint_dir, 'model.ckpt-{}'.format(step)))

    def synthesize(self, args, text, idx):
        seq = text_to_sequence(text)
        dec_input = np.zeros((1, 200, mel_dim), dtype='float32')
        pred = []
        for i in range(1, 201):
            mel_out, alignment = self.session.run([self.mel_output, self.alignment],
                                       feed_dict={self.model.enc_input: [np.asarray(seq, dtype=np.int32)],
                                                  self.model.sequence_length: np.asarray([len(seq)], dtype=np.int32),
                                                  self.model.dec_input: dec_input})
            if i < 200:
                dec_input[:, i, :] = mel_out[5*i - 1, :]
            pred.extend(mel_out[5*(i-1):5*i, :])

        np.save(os.path.join(args.save_dir, 'mel-{}'.format(idx)), pred, allow_pickle=False)

        input_seq = sequence_to_text(seq)
        alignment_dir = os.path.join(args.save_dir, 'align-{}.png'.format(idx))
        plot_alignment(alignment, alignment_dir, input_seq)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--step', required=True)
    parser.add_argument('--save_dir', default='./output')
    args = parser.parse_args()
    os.makedirs(args.save_dir, exist_ok=True)
    synth = Synthesizer()
    synth.load(args.step)
    for i, text in enumerate(sentences):
        jamo = ''.join(list(hangul_to_jamo(text)))
        synth.synthesize(args, jamo, i)


from ipywidgets import interact, interactive, fixed, interact_manual
import ipywidgets as widgets
from IPython.display import display
from gtts import gTTS
import os

import numpy as np
import matplotlib.pyplot as plt
#%matplotlib inline

import pandas as pd
import ctcsound

pan = 0
index = 10
cs = ctcsound.Csound()

csd = '''
<CsoundSynthesizer>
<CsOptions>
-odac -d

</CsOptions>
<CsInstruments>

sr = 44100
ksmps = 32
nchnls = 2
0dbfs = 1

instr 1

;aMod1 poscil 200, 700, 1
aMod1 poscil p4, p5, 1   ; p4 = amp1, p5 = f1, p6 = amp2, p7 = f2
;aMod2 poscil 1800, 290, 1
aMod2 poscil p6, p7, 1
kenv linen p9 , 0.3 , p3, p9
aSig poscil kenv, 440+aMod1+aMod2, 1
outs aSig*(1-p8), aSig*p8
endin

</CsInstruments>
<CsScore>
f 0 14400
f 1 0 1024 10 1

</CsScore>

</CsoundSynthesizer>
'''

cs.compileCsdText(csd)
cs.start()
pt = ctcsound.CsoundPerformanceThread(cs.csound())
pt.play()

def f(percentage):
	global index
	index = percentage

def on_button_clicked(b):
    in_min  = 0
    in_max = 100
    out_min=690
    out_max = 710

    global index
    tts = gTTS(text=(str(index)+'percent'), lang='en')
    tts.save("num.mp3")
    os.system("afplay num.mp3")
    if (index>50):
        pan = 1
    if (index<50):
        pan = 0
    if (index==50):
        pan = 0.5
    freq = (index - in_min) * (out_max - out_min) / (in_max - in_min) + out_min
    #print(freq)

    pt.scoreEvent(False, 'i', (1, 0, 4, 200, 700, 200, freq, pan, 0.5))


import logging as log
import os
import sys
import time

import cv2
import numpy as np
from openvino.inference_engine import IENetwork, IECore

from text_spotting import ModelHandler
from .tracker import StaticIOUTracker

SOS_INDEX = 0
EOS_INDEX = 1

class TextSpottingModel:

    def __init__(self, device='CPU', prob_threshold=0.3, max_seq_len=7,
                 iou_threshold=0.4, model_type='FP32', rgb2bgr=True, performance_counts=False, verbose=True):

        assert (model_type == 'FP32') or (model_type == 'FP16')

        mask_rcnn_model_xml = ModelHandler.get_models()['{}/text-spotting-0002-detector.xml'.format(model_type)]
        mask_rcnn_model_bin = ModelHandler.get_models()['{}/text-spotting-0002-detector.bin'.format(model_type)]

        text_enc_model_xml = ModelHandler.get_models()['{}/text-spotting-0002-recognizer-encoder.xml'.format(model_type)]
        text_enc_model_bin = ModelHandler.get_models()['{}/text-spotting-0002-recognizer-encoder.bin'.format(model_type)]

        text_dec_model_xml = ModelHandler.get_models()['{}/text-spotting-0002-recognizer-decoder.xml'.format(model_type)]
        text_dec_model_bin = ModelHandler.get_models()['{}/text-spotting-0002-recognizer-decoder.bin'.format(model_type)]

        # Plugin initialization for specified device and load extensions library if specified.
        log.info('Creating Inference Engine...')
        ie = IECore()
        # Read IR
        log.info('Loading network files:\n\t{}\n\t{}'.format(mask_rcnn_model_xml, mask_rcnn_model_bin))
        mask_rcnn_net = IENetwork(model=mask_rcnn_model_xml, weights=mask_rcnn_model_bin)

        log.info('Loading network files:\n\t{}\n\t{}'.format(text_enc_model_xml, text_enc_model_bin))
        text_enc_net = IENetwork(model=text_enc_model_xml, weights=text_enc_model_bin)

        log.info('Loading network files:\n\t{}\n\t{}'.format(text_dec_model_xml, text_dec_model_bin))
        text_dec_net = IENetwork(model=text_dec_model_xml, weights=text_dec_model_bin)

        supported_layers = ie.query_network(mask_rcnn_net, 'CPU')
        not_supported_layers = [layer for layer in mask_rcnn_net.layers.keys() if layer not in supported_layers]
        if len(not_supported_layers) != 0:
            log.error('Following layers are not supported by the plugin for specified device {}:\n {}'.
                      format(device, ', '.join(not_supported_layers)))
            log.error("Please try to specify cpu extensions library path in sample's command line parameters using -l "
                      "or --cpu_extension command line argument")
            sys.exit(1)

        required_input_keys = {'im_data', 'im_info'}
        assert required_input_keys == set(mask_rcnn_net.inputs.keys()), \
            'Demo supports only topologies with the following input keys: {}'.format(', '.join(required_input_keys))
        required_output_keys = {'boxes', 'scores', 'classes', 'raw_masks', 'text_features'}
        assert required_output_keys.issubset(mask_rcnn_net.outputs.keys()), \
            'Demo supports only topologies with the following output keys: {}'.format(', '.join(required_output_keys))

        n, c, h, w = mask_rcnn_net.inputs['im_data'].shape
        assert n == 1, 'Only batch 1 is supported by the demo application'
        self.shape = [n, c, h, w]
        self.verbose = verbose
        log.info('Loading IR to the plugin...')
        self.mask_rcnn_exec_net = ie.load_network(network=mask_rcnn_net, device_name=device, num_requests=2)
        self.text_enc_exec_net = ie.load_network(network=text_enc_net, device_name=device)
        self.text_dec_exec_net = ie.load_network(network=text_dec_net, device_name=device)
        self.hidden_shape = text_dec_net.inputs['prev_hidden'].shape
        self.prob_threshold = prob_threshold
        self.tracker = None
        self.visualizer = None
        self.iou_threshold = iou_threshold
        self.max_seq_len = max_seq_len
        self.rgb2bgr = rgb2bgr
        self.perf_counts = performance_counts
        self.device_names = None
        log.info('Model ready...')

    def predict(self, frame):
        """
        returns: texts, boxes, scores, frame
        boxes are [left, top, right, bottom]
        """
        [n, c, h, w] = self.shape
        # Resize the image to keep the same aspect ratio and to fit it to a window of a target size.
        scale_x = scale_y = min(h / frame.shape[0], w / frame.shape[1])
        input_image = cv2.resize(frame, None, fx=scale_x, fy=scale_y)

        if self.rgb2bgr:
            input_image = input_image[:, :, ::-1]  # reverse channels order

        input_image_size = input_image.shape[:2]
        input_image = np.pad(input_image, ((0, h - input_image_size[0]),
                                           (0, w - input_image_size[1]),
                                           (0, 0)),
                             mode='constant', constant_values=0)
        # Change data layout from HWC to CHW.
        input_image = input_image.transpose((2, 0, 1))
        input_image = input_image.reshape((n, c, h, w)).astype(np.float32)
        input_image_info = np.asarray([[input_image_size[0], input_image_size[1], 1]], dtype=np.float32)
        del n, c, h, w
        # Run the net.
        inf_start = time.time()
        log.info('running main network')
        outputs = self.mask_rcnn_exec_net.infer({'im_data': input_image, 'im_info': input_image_info})
        log.info('main network finished')
        if len(outputs['boxes']) == 0:
            return [], [], [], []
        # Parse detection results of the current request
        boxes = outputs['boxes']
        scores = outputs['scores']
        classes = outputs['classes'].astype(np.uint32)
        raw_masks = outputs['raw_masks']
        text_features = outputs['text_features']

        # Filter out detections with low confidence.
        detections_filter = scores > self.prob_threshold
        scores = scores[detections_filter]
        classes = classes[detections_filter]
        boxes = boxes[detections_filter]
        raw_masks = raw_masks[detections_filter]
        text_features = text_features[detections_filter]

        # initialize text_features_secondary
        text_features_secondary = [None for t in text_features]
        matches_secondary = [None for t in text_features]

        boxes[:, 0::2] /= scale_x
        boxes[:, 1::2] /= scale_y

        masks = []
        for box, cls, raw_mask in zip(boxes, classes, raw_masks):
            raw_cls_mask = raw_mask[cls, ...]
            mask = self.segm_postprocess(box, raw_cls_mask, frame.shape[0], frame.shape[1])
            masks.append(mask)

        texts = []
        alphabet = '  0123456789abcdefghijklmnopqrstuvwxyz'
        for feature in text_features:
            feature = self.text_enc_exec_net.infer({'input': feature})['output']
            feature = np.reshape(feature, (feature.shape[0], feature.shape[1], -1))
            feature = np.transpose(feature, (0, 2, 1))

            hidden = np.zeros(self.hidden_shape)
            prev_symbol_index = np.ones((1,)) * SOS_INDEX

            text = ''
            for i in range(self.max_seq_len):
                decoder_output = self.text_dec_exec_net.infer({
                    'prev_symbol': prev_symbol_index,
                    'prev_hidden': hidden,
                    'encoder_outputs': feature})
                symbols_distr = decoder_output['output']
                prev_symbol_index = int(np.argmax(symbols_distr, axis=1))
                if prev_symbol_index == EOS_INDEX:
                    break
                text += alphabet[prev_symbol_index]
                hidden = decoder_output['hidden']

            texts.append(text)

        inf_end = time.time()
        inf_time = inf_end - inf_start

        render_start = time.time()

        if len(boxes) and self.verbose:
            log.info('Detected boxes:')
            log.info('  Text | Confidence |     Left |     Top |     Right |     Bottom ')
            for box, text, score, mask in zip(boxes, texts, scores, masks):
                log.info('{} | {:>10f} | {:>8.2f} | {:>8.2f} | {:>8.2f} | {:>8.2f} '.format(text, score, *box))

        render_time = 0
        # Visualize masks.
        # frame = visualizer(frame, boxes, classes, scores, masks, texts, masks_tracks_ids)

        # Draw performance stats.
        inf_time_message = 'Inference and post-processing time: {:.3f} ms'.format(inf_time * 1000)
        render_time_message = 'OpenCV rendering time: {:.3f} ms'.format(render_time * 1000)

        # Print performance counters.
        if self.perf_counts:
            perf_counts = self.mask_rcnn_exec_net.requests[0].get_perf_counts()
            log.info('Performance counters:')
            log.info('{:<70} {:<15} {:<15} {:<15} {:<10}'.format('name', 'layer_type', 'exet_type', 'status',
                                                              'real_time, us'))
            for layer, stats in perf_counts.items():
                log.debug('{:<70} {:<15} {:<15} {:<15} {:<10}'.format(layer, stats['layer_type'], stats['exec_type'],
                                                                      stats['status'], stats['real_time']))

        render_end = time.time()
        render_time = render_end - render_start
        log.info(f"Render time = {render_time}")

        return texts, boxes, scores, frame

    @staticmethod
    def segm_postprocess(box, raw_cls_mask, im_h, im_w):
        # Add zero border to prevent upsampling artifacts on segment borders.
        raw_cls_mask = np.pad(raw_cls_mask, ((1, 1), (1, 1)), 'constant', constant_values=0)
        extended_box = TextSpottingModel.expand_box(box, raw_cls_mask.shape[0] / (raw_cls_mask.shape[0] - 2.0)).astype(int)
        w, h = np.maximum(extended_box[2:] - extended_box[:2] + 1, 1)
        x0, y0 = np.clip(extended_box[:2], a_min=0, a_max=[im_w, im_h])
        x1, y1 = np.clip(extended_box[2:] + 1, a_min=0, a_max=[im_w, im_h])

        raw_cls_mask = cv2.resize(raw_cls_mask, (w, h)) > 0.5
        mask = raw_cls_mask.astype(np.uint8)
        # Put an object mask in an image mask.
        im_mask = np.zeros((im_h, im_w), dtype=np.uint8)
        im_mask[y0:y1, x0:x1] = mask[(y0 - extended_box[1]):(y1 - extended_box[1]),
                                (x0 - extended_box[0]):(x1 - extended_box[0])]
        return im_mask

    @staticmethod
    def expand_box(box, scale):
        w_half = (box[2] - box[0]) * .5
        h_half = (box[3] - box[1]) * .5
        x_c = (box[2] + box[0]) * .5
        y_c = (box[3] + box[1]) * .5
        w_half *= scale
        h_half *= scale
        box_exp = np.zeros(box.shape)
        box_exp[0] = x_c - w_half
        box_exp[2] = x_c + w_half
        box_exp[1] = y_c - h_half
        box_exp[3] = y_c + h_half
        return box_exp


import os.path as op
import pandas as pd
import numpy as np
from glob import glob

df_generic = pd.read_csv('../stims/stimuli-expressive_selection-train+test.tsv', sep='\t', index_col=0)

tsvs = sorted(glob('data/ratings_complete/*expressive_ratings.tsv'))
dfs = []
for tsv in tsvs:
    sub = op.basename(tsv).split("_")[0]
    print(f"Processing {sub} ...")

    df = pd.read_csv(tsv, sep='\t', index_col=0).loc[:, ['data_split', 'rep', 'rating', 'rating_type', 'trial_type']]
    df = df.query("rating_type != 'emotion'")
    df_orig = df.copy()
    df['rating'] = df['rating'].astype(float)
    df = df.groupby(['trial_type', 'rating_type']).mean()['rating'].reset_index()
    df = df.pivot(index='trial_type', columns='rating_type', values='rating')
    df['sub'] = sub
    df['data_split'] = df_generic.loc[df.index, 'data_split']
    df = df.loc[:, ['sub', 'data_split'] + [col for col in df.columns[:-1] if col != 'sub']]
    dfs.append(df)

df = pd.concat(dfs, axis=0)#.set_index('trial_type')
df.to_csv('data/featurespace-circumplex.tsv', sep='\t')
print(df.shape)