#!/usr/bin/env python
# coding=utf-8
'''

yxhu@NPU-ASLP in Sogou inc.
modified by yxhu in Tencent AiLab 2020

'''

import numpy as np
import scipy
import torch
import random
import torch.nn as nn
from torch.utils.data import Dataset
import torch.utils.data as tud
import os
import sys
sys.path.append(
    os.path.dirname(__file__))
import soundfile as sf
import multiprocessing as mp
from Augmentation import  Mixer, load_scp, clip_data, load_wav

class AutoDataset(Dataset):

    def __init__(
            self,
            scps,
            rir_scps,
            segement_length=8,
            sample_rate=16000,
            load_memory=False,
            repeate=50,
            mono=False,
        ):
        '''
            scp_file_name: the list include:[input_wave_path, output_wave_path, duration]
            spk_emb_scp: a speaker embedding ark's scp
            segement_length: to clip data in a fix length segment, default: 4s
            sample_rate: the sample rate of wav, default: 16000
            processer: a processer class to handle wave data
            gender2spk: a list include gender2spk, default: None
        '''
        mgr = mp.Manager()
        self.mono = mono
        self.index = mgr.list()#[d for b in buckets for d in b]
        self.load_memory = load_memory
        self.sample_rate = sample_rate
        self.wav_list = load_scp(scps[0], load_memory, mono=mono, sample_rate=self.sample_rate)
        print('load target scp success')
        self.segement_length = segement_length * sample_rate
        _dochunk(self.wav_list, self.index, self.segement_length,sample_rate=self.sample_rate)
        print('chunk data success', len(self.index), len(self.wav_list))
        print(scps)
        np.random.shuffle(self.index)
        self.index *= repeate
        self.mixer = Mixer(mix_scps=[scps[1]], rir_scps=rir_scps, snr_range=[-5,20], mix_mode='first', load_memory=load_memory, mono=mono, sample_rate=sample_rate)
        self.num_states = len(self.index) % 3000
        self.randstates = [ np.random.RandomState(idx*10) for idx in range(self.num_states)]

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

    def __getitem__(self, index):
        wav_idx,start_time = self.index[index]
        data_info = self.wav_list[wav_idx]
        spkid, path = data_info['spkid'], data_info['path']
        if self.load_memory:
            path = clip_data(path, start_time, self.segement_length)
        Y, X, D, R, fac,snr,scale = self.mixer(
                                    (spkid, path),
                                    start=start_time,
                                    randstat=self.randstates[index%self.num_states],
                                    segement_length=self.segement_length,
                                    emb_length=18,
                                )
        '''
        X = load_wav(path,self.mono)[0][start_time:start_time+self.segement_length]
        Y = X
        emb =X
        if not self.mono :
            X = np.concatenate(X,-1).T
            Y = Y.T
        else:
            X = np.stack(X,0)
        '''
        return Y, X[0]

def worker(target_list, result_list, start, end, segement_length, sample_rate):
    for idx in range(len(target_list[start:end])):
        item = target_list[idx]
        duration = item['duration']
        length = duration
        if length < segement_length:
            sample_index = -1
            if length * 2 < segement_length and length*4 > segement_length:
                sample_index = -2
            elif length * 2 > segement_length:
                sample_index = -1
            else:
                continue
            result_list.append([idx, sample_index])
        else:
            sample_index = 0
            if isinstance(item['path'], str):
                data = load_wav(item['path'], sample_rate=sample_rate,mono=True)[0]
            else:
                data = item['path']
            data = np.abs(data)
            threshold=1e-3 # this is just for vocal

            while sample_index + segement_length < length:
                # Attention
                # there is an easy vad to filter
                # blank
                non_zero = np.sum(data[sample_index:sample_index+segement_length]>threshold)

                if non_zero*10 > segement_length:
                    result_list.append(
                        [idx, sample_index])
                    sample_index += (segement_length//3)
                else:
                    #print(segement_length, duration, item['spkid'], non_zero, sample_index)
                    sample_index += segement_length//2

            if sample_index < length:
                non_zero = np.sum(data[-length+segement_length:]>threshold)
                if non_zero*4 > segement_length:
                    result_list.append([
                        idx,
                        int(length - segement_length),
                    ])

def _dochunk(wav_list, index, segement_length=16000*4, sample_rate=44100,num_threads=12):
        # mutliproccesing
        pc_list = []
        stride = len(wav_list) // num_threads
        if stride < 100:
            p = mp.Process(
                            target=worker,
                            args=(
                                    wav_list,
                                    index,
                                    0,
                                    len(wav_list),
                                    segement_length,
                                    sample_rate
                                )
                        )
            p.start()
            pc_list.append(p)
        else:
            for idx in range(num_threads):
                if idx == num_threads-1:
                    end = len(wav_list)
                else:
                    end = (idx+1)*stride
                p = mp.Process(
                                target=worker,
                                args=(
                                    wav_list,
                                    index,
                                    idx*stride,
                                    end,
                                    segement_length,
                                    sample_rate
                                )
                            )
                p.start()
                pc_list.append(p)
        for p in pc_list:
            p.join()

class Sampler(tud.sampler.Sampler):
    '''

    '''
    def __init__(self, data_source, batch_size):
        it_end = len(data_source) - batch_size + 1
        self.batches = [range(i, i+batch_size)
                        for i in range(0, it_end, batch_size)]
        self.data_source = data_source

    def __iter__(self):
        random.shuffle(self.batches)
        return (i for b in self.batches for i in b)

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

def zero_pad_concat(inputs):
    max_t = max(inp.shape[0] for inp in inputs)
    shape = np.array([len(inputs), max_t, inputs[0].shape[1]])
    inputs_mat = np.zeros(shape, np.float32)
    for idx, inp in enumerate(inputs):
        inputs_mat[idx, :inp.shape[0],:] = inp
    return inputs_mat

def collate_fn(data):
    #inputs, labels, embeddings, spkid= zip(*data)
    inputs, labels = zip(*data)
    inputs = np.array(inputs, dtype=np.float32)
    labels = np.array(labels, dtype=np.float32)
    return torch.from_numpy(inputs),\
           torch.from_numpy(labels)

def make_dataloader(vocals_scp, background_scp, rir_scps, batch_size, sample_rate=16000,num_workers=12, training=True, repeate=10, load_memory=False, segement_length=6, mono=True, sampler=None):
    dataset = AutoDataset(
                            [vocals_scp, background_scp],
                            rir_scps=rir_scps,
                            load_memory=load_memory,
                            repeate=repeate,
                            sample_rate=sample_rate,
                            segement_length=segement_length,
                            mono=mono
                         )
    loader = tud.DataLoader(
                            dataset,
                            batch_size=batch_size,
                            num_workers=num_workers,
                            collate_fn=collate_fn,
                            shuffle=(training and sampler is None),
                            drop_last=False,
                            sampler=sampler
                        )
    return loader

if __name__ == '__main__':
    torch.manual_seed(20)
    test_mode = 0
    if test_mode == 0:
        laoder = make_dataloader(
                            '../data/tr_clean.lst',
                            '../data/tr_noise.lst',
                            '../data/tr_rir.lst',
                            sample_rate=16000,
                            batch_size=32,
                            num_workers=8,
                            training=True,
                            load_memory=False,
                            repeate=30,
                            mono=True
            )
    else:
        laoder= make_dataloader(
                            './vocals',
                            './background',
            32, num_workers=8, training=False)
    import time
    stime = time.time()
    print_freq=10
    for epoch in range(3):
        if epoch == 0 or epoch == 4:
            torch.manual_seed(20)
        else:
            torch.manual_seed(20+epoch)
        '''
        idx=str(99)
        print(dataset)
        epoch=str(epoch)
        inputs = dataset[0][0]
        sf.write(epoch+'_'+idx+'_noisy.wav', inputs, 16000)
        '''
        print('num_batch', len(laoder))
        for idx, data in enumerate(laoder):
            if (idx+1)%print_freq == 0:
                inputs, labels = data
                etime = time.time()
                print(epoch, idx, inputs.size(), (etime-stime)/print_freq)
                idx=str(idx)
                epoch=str(epoch)
        #        sf.write(epoch+'_'+idx+'_noisy.wav', inputs.numpy().T, 44100)
        #        sf.write(epoch+'_'+idx+'_clean.wav', labels.numpy().T, 44100)
        #        sf.write(epoch+'_'+idx+'_emb.wav',  inputs[2].numpy().T, 44100)
                stime = etime
        #        break
        break


# Copyright (c) 1996-2015 PSERC. All rights reserved.
# Use of this source code is governed by a BSD-style
# license that can be found in the LICENSE file.

"""Tests if two matrices are identical to some tolerance.
"""

from numpy import array, max, abs, nonzero, argmax, zeros

from pypower.t.t_ok import t_ok
from pypower.t.t_globals import TestGlobals

def t_is(got, expected, prec=5, msg=''):
    """Tests if two matrices are identical to some tolerance.

    Increments the global test count and if the maximum difference
    between corresponding elements of C{got} and C{expected} is less
    than 10**(-C{prec}) then it increments the passed tests count,
    otherwise increments the failed tests count. Prints 'ok' or 'not ok'
    followed by the MSG, unless the global variable t_quiet is true.
    Intended to be called between calls to C{t_begin} and C{t_end}.

    @author: Ray Zimmerman (PSERC Cornell)
    """
    if isinstance(got, int) or isinstance(got, float):
        got = array([got], float)
    elif isinstance(got, list) or isinstance(got, tuple):
        got = array(got, float)

    if isinstance(expected, int) or isinstance(expected, float):
        expected = array([expected], float)
    elif isinstance(expected, list) or isinstance(expected, tuple):
        expected = array(expected, float)

    if (got.shape == expected.shape) or (expected.shape == (0,)):
        got_minus_expected = got - expected
        max_diff = max(max(abs(got_minus_expected)))
        condition = ( max_diff < 10**(-prec) )
    else:
        condition = False
        max_diff = 0

    t_ok(condition, msg)
    if (not condition and not TestGlobals.t_quiet):
        s = ''
        if max_diff != 0:
            idx = nonzero(not abs(got_minus_expected) < 10**(-prec))
            if len(idx) == 1:  # 1D array
                idx = (idx[0], zeros( len(got_minus_expected) ))
            i, j = idx

            k = i + (j-1) * expected.shape[0]

            got = got.flatten()
            expected = expected.flatten()
            got_minus_expected = got_minus_expected.flatten()

            kk = argmax( abs(got_minus_expected[ k.astype(int) ]) )

            s += '  row     col          got             expected          got - exp\n'
            s += '-------  ------  ----------------  ----------------  ----------------'
            for u in range(len(i)):
                s += '\n%6d  %6d  %16g  %16g  %16g' % \
                    (i[u], j[u], got[k[u]], expected[k[u]], got_minus_expected[k[u]])
                if u == kk:
                    s += '  *'
            s += '\nmax diff @ (%d,%d) = %g > allowed tol of %g\n\n' % \
                (i[kk], j[kk], max_diff, 10**(-prec))
        else:
            s += '    dimension mismatch:\n'

            if len(got.shape) == 1:  # 1D array
                s += '             got: %d\n' % got.shape
            else:
                s += '             got: %d x %d\n' % got.shape

            if len(expected.shape) == 1:  # 1D array
                s += '        expected: %d\n' % expected.shape
            else:
                s += '        expected: %d x %d\n' % expected.shape

        print(s)

if __name__ == '__main__':
    a = array([[1,2,3], [4,5,6]])
    b = array([[2,3,4], [5,6,7]])
    TestGlobals.t_quiet = False
    t_is(a, b)


import json
import sys
import numpy as np
import cv2
import json
import queue
import paho.mqtt.client as mqtt
import threading
import requests
import os
import binascii
import random

class GlobalState:
    jobQueue = queue.Queue()

class Image:
    def __init__(self, width, height, ndims=3):
        self.image = np.zeros((height,width,ndims), np.uint8)
        self.prevPosition = None

    def save(self, path):
        cv2.imwrite(path, self.image)

    def add(self, coords, color=(255,255,255), size=10):
        position = (int(coords[0]), int(coords[1]))
        if self.prevPosition is not None:
            cv2.line(self.image, self.prevPosition, position, color, size*2)
        cv2.circle(self.image, position, size, color, -1)
        self.prevPosition = position

    def get(self):
        return self.image

    def resetPrevious(self):
        self.prevPosition = None

class Artistry:
    def __init__(self):
        self.fgbg = cv2.createBackgroundSubtractorMOG2()
        self.kernel = np.ones((100,100),np.uint8)

    def mask_color(self, image, ranges):
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        clmask = cv2.inRange(hsv, ranges[0], ranges[1])
        masked_hsv = cv2.bitwise_and(
            hsv, hsv, mask=clmask)
        masked_color = cv2.cvtColor(masked_hsv, cv2.COLOR_HSV2BGR)
        return masked_color

    def mask_background(self, image):
        fgmask = self.fgbg.apply(image)
        masked_background = cv2.bitwise_and(image, image, mask=fgmask)
        return masked_background

    def thresh_image(self, image):
        cimg = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        thresh = cv2.morphologyEx(cimg, cv2.MORPH_CLOSE, self.kernel)
        return thresh

    def find_thresh_positions(self, image):
        results = []
        img, contours, hierarchy = cv2.findContours(
            image,
            cv2.RETR_TREE,
            cv2.CHAIN_APPROX_SIMPLE
        )
        # There is likely a better way to do this.
        if len(contours) > 0:
            for i in contours:
                k = i.sum(axis=1)
                x = 0
                y = 0
                for j in k:
                    x += j[0]
                    y += j[1]
                x = x/len(k)
                y = y/len(k)
                results.append((len(i), (x, y)))

        return results

class ArtUI:
    def __init__(self, config):
        self.config = config
        self.img = Image(config['width'], config['height'])
        self.lower_range = np.array(config['ranges']['lower'])
        self.upper_range = np.array(config['ranges']['upper'])
        self.current_color = (255,255,0)
        self.brushDown = False
        self.assets = {}
        self.assets['bob_ross_hair'] = cv2.imread('assets/bob_ross_hair.png')
        self.select_episode()

    def start(self):
        art = Artistry()

        cap = cv2.VideoCapture(self.config['video_input'])
        cap.set(3, self.config['width'])
        cap.set(4, self.config['height'])

        while True:
            # Check if we have anything new in our queue.
            ret, frame = cap.read()
            try:
                job = GlobalState.jobQueue.get_nowait()
                self.perform_job(job, frame)
            except queue.Empty:
                pass

            ret, frame = cap.read()
            # Mask out the background
            masked_background = art.mask_background(frame)
            # Pick out a color.
            masked_color = art.mask_color(
                masked_background, (self.lower_range, self.upper_range))
            # Build a mesh for the object.
            thresh = art.thresh_image(masked_color)
            positions = art.find_thresh_positions(thresh)
            # Display the best position
            try:
                _, coords = sorted(positions)[0]
                if self.brushDown:
                    self.img.add(coords, color=self.current_color)
            except IndexError:
                pass
            # Our UI
            overlay = cv2.add(frame, self.img.get())

            if config.get('bob_ross_hair', False):
                overlay = cv2.addWeighted(
                    self.assets['bob_ross_hair'], 1,
                    overlay, 1, 0)

            if config.get('developer', False):
                output = np.hstack(
                    (
                        cv2.flip(masked_color, 1),
                        cv2.flip(overlay, 1),
                        cv2.flip(self.img.get(), 1)
                    )
                )
            else:
                output = np.hstack(
                    (
                        cv2.flip(self.screenshot, 1),
                        cv2.flip(overlay, 1)
                    )
                )

            if config.get('fullscreen', False):
                cv2.namedWindow('window', cv2.WND_PROP_FULLSCREEN)
                cv2.setWindowProperty(
                    'window',cv2.WND_PROP_FULLSCREEN,cv2.WINDOW_FULLSCREEN)
            cv2.imshow('window', output)

            # User input
            key = cv2.waitKey(1)
            if key == ord('q'): break
            elif key == ord('r'):
                self._reset_image()
            elif key == ord('s'):
                self._backup_image(frame)
                self._reset_image()
            elif key == ord('d'):
                self.brushDown = not(self.brushDown)

        cap.release()
        cv2.destroyAllWindows()

    def select_episode(self):
        self.episode = random.choice(self.config['episodes'])
        self.screenshot = cv2.imread(self.episode['file'])

    def _set_color(self, job):
        r = eval('0x'+job['data'][0:2])
        g = eval('0x'+job['data'][2:4])
        b = eval('0x'+job['data'][4:6])
        self.current_color = (b, g, r)
        self.img.resetPrevious()

    def _reset_image(self):
        self.select_episode()
        self.img = Image(self.config['width'], self.config['height'])

    def perform_job(self, job, frame):
        global GlobalState
        if job['command'] == 'brushUp':
            self.img.resetPrevious()
            self.brushDown = False
        elif  job['command'] == 'brushDown':
            self.brushDown = True
        elif job['command'] == 'colour':
            self._set_color(job)
        elif job['command'] == 'save':
            # Push to the website.
            self._backup_image(frame)
        elif job['command'] == 'reset':
            self._reset_image()

    def _backup_image(self, frame):
        # save it to /tmp/
        image_id = binascii.hexlify(os.urandom(16)).decode('ascii')+'.jpg'
        image_path = '/tmp/{}'.format(image_id)
        self.img.image = cv2.add(self.img.image, frame)
        self.img.save(image_path)
        # Now we can go post it to the form.
        print(self.episode)
        r = requests.post(
            self.config['form']['host'],
            files={
                'file': open(image_path, 'rb')
            },
            data={
                'episode': self.episode['name']
            }
        )
        self._reset_image()

class MqttClient:
    def __init__(self, config):
        self.client = mqtt.Client()
        self.client.on_connect = MqttClient.on_connect
        self.client.on_message = MqttClient.on_message
        self.config = config

    def start(self):
        self.client.connect(
            self.config['host'],
            self.config['port'],
            60
        )
        self.client.loop_forever()

    @staticmethod
    def on_connect(client, userdaat, flags, rc):
        client.subscribe("HackShef/bobro")

    @staticmethod
    def on_message(client, userdata, msg):
        global GlobalState
        raw = json.loads(msg.payload)
        GlobalState.jobQueue.put(raw)

if __name__ == "__main__":
    if len(sys.argv) < 1:
        exit()
    with open(sys.argv[1]) as config_file:
        config = json.load(config_file)
        art_c = ArtUI(config)
        threading.Thread(
            target=art_c.start
        ).start()
        mqtt_c = MqttClient(config['mqtt'])
        threading.Thread(
            target=mqtt_c.start
        ).start()


from keras import backend as K
from keras.models import load_model
from keras.optimizers import Adam
from imageio import imread
import numpy as np
from matplotlib import pyplot as plt

from models.keras_ssd512 import ssd_512
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization
from data_generator.object_detection_2d_data_generator import DataGenerator
from eval_utils.average_precision_evaluator import Evaluator

import os
img_height = 512
img_width = 512
n_classes = 20
model_mode = 'inference'

K.clear_session() # Clear previous models from memory.

use_model = False
if not use_model:
    model = ssd_512(image_size=(img_height, img_width, 3),
                    n_classes=20,
                    mode='inference',
                    l2_regularization=0.0005,
                    scales=[0.07, 0.15, 0.3, 0.45, 0.6, 0.75, 0.9, 1.05], # The scales for MS COCO are [0.04, 0.1, 0.26, 0.42, 0.58, 0.74, 0.9, 1.06]
                    aspect_ratios_per_layer=[[1.0, 2.0, 0.5],
                                             [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                             [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                             [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                             [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                             [1.0, 2.0, 0.5],
                                             [1.0, 2.0, 0.5]],
                   two_boxes_for_ar1=True,
                   steps=[8, 16, 32, 64, 128, 256, 512],
                   offsets=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
                   clip_boxes=False,
                   variances=[0.1, 0.1, 0.2, 0.2],
                   normalize_coords=True,
                   subtract_mean=[123, 117, 104],
                   swap_channels=[2, 1, 0],
                   confidence_thresh=0.5,
                   iou_threshold=0.45,
                   top_k=200,
                   nms_max_output_size=400)

    weights_path = 'VGG_VOC0712Plus_SSD_512x512_iter_240000.h5'
    # weights_path = 'final_weights.h5'
    model.load_weights(weights_path, by_name=True)

    # 3: Compile the model so that Keras won't complain the next time you load it.
    adam = Adam(lr=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-08, decay=0.0)
    ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0)
    model.compile(optimizer=adam, loss=ssd_loss.compute_loss)
else:
    model_path = 'final_model.h5'
    ssd_loss = SSDLoss(neg_pos_ratio=3, alpha=1.0)
    model = load_model(model_path, custom_objects={'AnchorBoxes': AnchorBoxes,
                                                   'L2Normalization': L2Normalization,
                                                   'DecodeDetections': DecodeDetections,
                                                   'compute_loss': ssd_loss.compute_loss})

dataset = DataGenerator()
# TODO: Set the paths to the dataset here.
Pascal_VOC_dataset_images_dir = './datasets/VOCdevkit/VOC2007/JPEGImages/'
Pascal_VOC_dataset_annotations_dir = './datasets/VOCdevkit/VOC2007/Annotations/'
Pascal_VOC_dataset_image_set_filename = './datasets/VOCdevkit/VOC2007/ImageSets/Main/test.txt'

# The XML parser needs to now what object class names to look for and in which order to map them to integers.
classes = ['background',
           'aeroplane', 'bicycle', 'bird', 'boat',
           'bottle', 'bus', 'car', 'cat',
           'chair', 'cow', 'diningtable', 'dog',
           'horse', 'motorbike', 'person', 'pottedplant',
           'sheep', 'sofa', 'train', 'tvmonitor']

dataset.parse_xml(images_dirs=[Pascal_VOC_dataset_images_dir],
                  image_set_filenames=[Pascal_VOC_dataset_image_set_filename],
                  annotations_dirs=[Pascal_VOC_dataset_annotations_dir],
                  classes=classes,
                  include_classes='all',
                  exclude_truncated=False,
                  exclude_difficult=False,
                  ret=False)

evaluator = Evaluator(model=model,
                      n_classes=n_classes,
                      data_generator=dataset,
                      model_mode=model_mode)

results = evaluator(img_height=img_height,
                    img_width=img_width,
                    batch_size=8,
                    data_generator_mode='resize',
                    round_confidences=False,
                    matching_iou_threshold=0.5,
                    border_pixels='include',
                    sorting_algorithm='quicksort',
                    average_precision_mode='sample',
                    num_recall_points=11,
                    ignore_neutral_boxes=True,
                    return_precisions=True,
                    return_recalls=True,
                    return_average_precisions=True,
                    verbose=True)

mean_average_precision, average_precisions, precisions, recalls = results

print('average_precisions =',average_precisions)

for i in range(1, len(average_precisions)):
    print("{:<14}{:<6}{}".format(classes[i], 'AP', round(average_precisions[i], 3)))
print()
print("{:<14}{:<6}{}".format('','mAP', round(mean_average_precision, 3)))

m = max((n_classes + 1) // 2, 2)
n = 2

fig, cells = plt.subplots(m, n, figsize=(n*8,m*8))
for i in range(m):
    for j in range(n):
        if n*i+j+1 > n_classes: break
        cells[i, j].plot(recalls[n*i+j+1], precisions[n*i+j+1], color='blue', linewidth=1.0)
        cells[i, j].set_xlabel('recall', fontsize=14)
        cells[i, j].set_ylabel('precision', fontsize=14)
        cells[i, j].grid(True)
        cells[i, j].set_xticks(np.linspace(0,1,11))
        cells[i, j].set_yticks(np.linspace(0,1,11))
        cells[i, j].set_title("{}, AP: {:.3f}".format(classes[n*i+j+1], average_precisions[n*i+j+1]), fontsize=16)

evaluator.get_num_gt_per_class(ignore_neutral_boxes=True,
                               verbose=False,
                               ret=False)

evaluator.match_predictions(ignore_neutral_boxes=True,
                            matching_iou_threshold=0.5,
                            border_pixels='include',
                            sorting_algorithm='quicksort',
                            verbose=True,
                            ret=False)

precisions, recalls = evaluator.compute_precision_recall(verbose=True, ret=True)

average_precisions = evaluator.compute_average_precisions(mode='integrate',
                                                          num_recall_points=11,
                                                          verbose=True,
                                                          ret=True)

mean_average_precision = evaluator.compute_mean_average_precision(ret=True)

for i in range(1, len(average_precisions)):
    print("{:<14}{:<6}{}".format(classes[i], 'AP', round(average_precisions[i], 3)))
print()
print("{:<14}{:<6}{}".format('','mAP', round(mean_average_precision, 3)))


import numpy as np
import cv2
import pickle
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

##############Project Settings###################
width = 640
height = 4480
threshold = 0.7
#################################################

cap = cv2.VideoCapture(0)
cap.set(3, width)
cap.set(4, height)

pickle_in = open("model_trained.p", "rb")
model = pickle.load(pickle_in)

def preProcessing(img):
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = cv2.equalizeHist(img)
    # Normalization
    img = img / 255
    return img

while True:
    success, imgOriginal = cap.read()
    img = np.asarray(imgOriginal)
    img = cv2.resize(img, (32, 32))
    img = preProcessing(img)
    cv2.imshow("Processed Image", img)
    img = img.reshape(1, 32, 32, 1)

    #Predict
    classIndex = int(model.predict_classes(img))
    print(classIndex)
    predictions = model.predict(img)
    probValue = np.amax(predictions)
    print(classIndex, probValue)
    if probValue > threshold:
        cv2.putText(imgOriginal, str(classIndex)+"  " + str(probValue), (50, 50), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 1)

    cv2.imshow("Original Image", imgOriginal)
    if cv2.waitKey(1) & 0xff == ord('q'):
        break


"""
Coppied from Rowan Zellers, modified by Ji
"""
"""
Adapted from Danfei Xu. In particular, slow code was removed
"""
import os
import numpy as np
import logging
from six.moves import cPickle as pickle
import json
import csv
from tqdm import tqdm

from core.config import cfg
from functools import reduce
# from lib.fpn.box_intersections_cpu.bbox import bbox_overlaps
from utils.boxes import bbox_overlaps, bbox_pair_overlaps, boxes_union
from datasets.voc_eval_rel import voc_eval, prepare_mAP_dets

np.set_printoptions(precision=3)

logger = logging.getLogger(__name__)

topk = 100

def eval_rel_results(all_results, output_dir, do_val):

    if do_val:
        if cfg.TEST.DATASETS[0].find('vg') >= 0:
            prd_k_set = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20)
        elif cfg.TEST.DATASETS[0].find('vrd') >= 0:
            prd_k_set = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 20, 70)
        else:
            prd_k_set = (1, 2, 3, 4, 5, 6, 7, 8, 9)

        if cfg.TEST.DATASETS[0].find('vg') >= 0:
            eval_sets = (False,)
        else:
            eval_sets = (False, True)

        for phrdet in eval_sets:
            eval_metric = 'phrdet' if phrdet else 'reldet'
            print('{}:'.format(eval_metric))

            for prd_k in prd_k_set:
                print('prd_k = {}:'.format(prd_k))

                recalls = {20: 0, 50: 0, 100: 0}
                all_gt_cnt = 0
                topk_dets = []
                for im_i, res in enumerate(tqdm(all_results)):
                    # in oi_all_rel some images have no dets
                    if res['prd_scores'] is None:
                        det_boxes_s_top = np.zeros((0, 4), dtype=np.float32)
                        det_boxes_o_top = np.zeros((0, 4), dtype=np.float32)
                        det_labels_s_top = np.zeros(0, dtype=np.int32)
                        det_labels_p_top = np.zeros(0, dtype=np.int32)
                        det_labels_o_top = np.zeros(0, dtype=np.int32)
                        det_scores_top = np.zeros(0, dtype=np.float32)
                    else:
                        det_boxes_sbj = res['sbj_boxes']  # (#num_rel, 4)
                        det_boxes_obj = res['obj_boxes']  # (#num_rel, 4)
                        det_labels_sbj = res['sbj_labels']  # (#num_rel,)
                        det_labels_obj = res['obj_labels']  # (#num_rel,)
                        det_scores_sbj = res['sbj_scores']  # (#num_rel,)
                        det_scores_obj = res['obj_scores']  # (#num_rel,)
                        det_scores_prd = res['prd_scores'][:, 1:]

                        det_labels_prd = np.argsort(-det_scores_prd, axis=1)
                        det_scores_prd = -np.sort(-det_scores_prd, axis=1)

                        det_scores_so = det_scores_sbj * det_scores_obj
                        det_scores_spo = det_scores_so[:, None] * det_scores_prd[:, :prd_k]
                        # det_scores_spo = det_scores_prd[:, :prd_k]
                        det_scores_inds = argsort_desc(det_scores_spo)[:topk]
                        det_scores_top = det_scores_spo[det_scores_inds[:, 0], det_scores_inds[:, 1]]
                        det_boxes_so_top = np.hstack(
                            (det_boxes_sbj[det_scores_inds[:, 0]], det_boxes_obj[det_scores_inds[:, 0]]))
                        det_labels_p_top = det_labels_prd[det_scores_inds[:, 0], det_scores_inds[:, 1]]
                        det_labels_spo_top = np.vstack(
                            (det_labels_sbj[det_scores_inds[:, 0]], det_labels_p_top, det_labels_obj[det_scores_inds[:, 0]])).transpose()

                        det_boxes_s_top = det_boxes_so_top[:, :4]
                        det_boxes_o_top = det_boxes_so_top[:, 4:]
                        det_labels_s_top = det_labels_spo_top[:, 0]
                        det_labels_p_top = det_labels_spo_top[:, 1]
                        det_labels_o_top = det_labels_spo_top[:, 2]

                    topk_dets.append(dict(image=res['image'],
                                          det_boxes_s_top=det_boxes_s_top,
                                          det_boxes_o_top=det_boxes_o_top,
                                          det_labels_s_top=det_labels_s_top,
                                          det_labels_p_top=det_labels_p_top,
                                          det_labels_o_top=det_labels_o_top,
                                          det_scores_top=det_scores_top))

                    gt_boxes_sbj = res['gt_sbj_boxes']  # (#num_gt, 4)
                    gt_boxes_obj = res['gt_obj_boxes']  # (#num_gt, 4)
                    gt_labels_sbj = res['gt_sbj_labels']  # (#num_gt,)
                    gt_labels_obj = res['gt_obj_labels']  # (#num_gt,)
                    gt_labels_prd = res['gt_prd_labels']  # (#num_gt,)
                    gt_boxes_so = np.hstack((gt_boxes_sbj, gt_boxes_obj))
                    gt_labels_spo = np.vstack((gt_labels_sbj, gt_labels_prd, gt_labels_obj)).transpose()
                    # Compute recall. It's most efficient to match once and then do recall after
                    # det_boxes_so_top is (#num_rel, 8)
                    # det_labels_spo_top is (#num_rel, 3)
                    if phrdet:
                        det_boxes_r_top = boxes_union(det_boxes_s_top, det_boxes_o_top)
                        gt_boxes_r = boxes_union(gt_boxes_sbj, gt_boxes_obj)
                        pred_to_gt = _compute_pred_matches(
                            gt_labels_spo, det_labels_spo_top,
                            gt_boxes_r, det_boxes_r_top,
                            phrdet=phrdet)
                    else:
                        pred_to_gt = _compute_pred_matches(
                            gt_labels_spo, det_labels_spo_top,
                            gt_boxes_so, det_boxes_so_top,
                            phrdet=phrdet)
                    all_gt_cnt += gt_labels_spo.shape[0]
                    for k in recalls:
                        if len(pred_to_gt):
                             match = reduce(np.union1d, pred_to_gt[:k])
                        else:
                            match = []
                        recalls[k] += len(match)

                    topk_dets[-1].update(dict(gt_boxes_sbj=gt_boxes_sbj,
                                              gt_boxes_obj=gt_boxes_obj,
                                              gt_labels_sbj=gt_labels_sbj,
                                              gt_labels_obj=gt_labels_obj,
                                              gt_labels_prd=gt_labels_prd))
                for k in recalls:
                    recalls[k] = float(recalls[k]) / (float(all_gt_cnt) + 1e-12)
                print_stats(recalls)
    else:
        prd_k = 2
        topk_dets = []
        for im_i, res in enumerate(tqdm(all_results)):
            # in oi_all_rel some images have no dets
            if res['prd_scores'] is None:
                det_boxes_s_top = np.zeros((0, 4), dtype=np.float32)
                det_boxes_o_top = np.zeros((0, 4), dtype=np.float32)
                det_labels_s_top = np.zeros(0, dtype=np.int32)
                det_labels_p_top = np.zeros(0, dtype=np.int32)
                det_labels_o_top = np.zeros(0, dtype=np.int32)
                det_scores_top = np.zeros(0, dtype=np.float32)
            else:
                det_boxes_sbj = res['sbj_boxes']  # (#num_rel, 4)
                det_boxes_obj = res['obj_boxes']  # (#num_rel, 4)
                det_labels_sbj = res['sbj_labels']  # (#num_rel,)
                det_labels_obj = res['obj_labels']  # (#num_rel,)
                det_scores_sbj = res['sbj_scores']  # (#num_rel,)
                det_scores_obj = res['obj_scores']  # (#num_rel,)
                det_scores_prd = res['prd_scores'][:, 1:]

                det_labels_prd = np.argsort(-det_scores_prd, axis=1)
                det_scores_prd = -np.sort(-det_scores_prd, axis=1)

                det_scores_so = det_scores_sbj * det_scores_obj
                det_scores_spo = det_scores_so[:, None] * det_scores_prd[:, :prd_k]
                # det_scores_spo = det_scores_prd[:, :prd_k]
                det_scores_inds = argsort_desc(det_scores_spo)[:topk]
                det_scores_top = det_scores_spo[det_scores_inds[:, 0], det_scores_inds[:, 1]]
                det_boxes_so_top = np.hstack(
                    (det_boxes_sbj[det_scores_inds[:, 0]], det_boxes_obj[det_scores_inds[:, 0]]))
                det_labels_p_top = det_labels_prd[det_scores_inds[:, 0], det_scores_inds[:, 1]]
                det_labels_spo_top = np.vstack(
                    (det_labels_sbj[det_scores_inds[:, 0]], det_labels_p_top, det_labels_obj[det_scores_inds[:, 0]])).transpose()

                det_boxes_s_top = det_boxes_so_top[:, :4]
                det_boxes_o_top = det_boxes_so_top[:, 4:]
                det_labels_s_top = det_labels_spo_top[:, 0]
                det_labels_p_top = det_labels_spo_top[:, 1]
                det_labels_o_top = det_labels_spo_top[:, 2]

            topk_dets.append(dict(image=res['image'],
                                  det_boxes_s_top=det_boxes_s_top,
                                  det_boxes_o_top=det_boxes_o_top,
                                  det_labels_s_top=det_labels_s_top,
                                  det_labels_p_top=det_labels_p_top,
                                  det_labels_o_top=det_labels_o_top,
                                  det_scores_top=det_scores_top))
        print('Saving topk dets...')
        topk_dets_f = os.path.join(output_dir, 'rel_detections_topk.pkl')
        with open(topk_dets_f, 'wb') as f:
            pickle.dump(topk_dets, f, pickle.HIGHEST_PROTOCOL)
        logger.info('topk_dets size: {}'.format(len(topk_dets)))
        print('Done.')

def print_stats(recalls):
    # print('====================== ' + 'sgdet' + ' ============================')
    for k, v in recalls.items():
        print('R@%i: %f' % (k, v))

def _compute_pred_matches(gt_triplets, pred_triplets,
                 gt_boxes, pred_boxes, iou_thresh=0.5, phrdet=False):
    """
    Given a set of predicted triplets, return the list of matching GT's for each of the
    given predictions
    :param gt_triplets:
    :param pred_triplets:
    :param gt_boxes:
    :param pred_boxes:
    :param iou_thresh:
    :return:
    """
    # This performs a matrix multiplication-esque thing between the two arrays
    # Instead of summing, we want the equality, so we reduce in that way
    # The rows correspond to GT triplets, columns to pred triplets
    keeps = intersect_2d(gt_triplets, pred_triplets)
    gt_has_match = keeps.any(1)
    pred_to_gt = [[] for x in range(pred_boxes.shape[0])]
    for gt_ind, gt_box, keep_inds in zip(np.where(gt_has_match)[0],
                                         gt_boxes[gt_has_match],
                                         keeps[gt_has_match],
                                         ):
        boxes = pred_boxes[keep_inds]
        if phrdet:
            gt_box = gt_box.astype(dtype=np.float32, copy=False)
            boxes = boxes.astype(dtype=np.float32, copy=False)
            rel_iou = bbox_overlaps(gt_box[None, :], boxes)[0]

            inds = rel_iou >= iou_thresh
        else:
            gt_box = gt_box.astype(dtype=np.float32, copy=False)
            boxes = boxes.astype(dtype=np.float32, copy=False)
            sub_iou = bbox_overlaps(gt_box[None,:4], boxes[:, :4])[0]
            obj_iou = bbox_overlaps(gt_box[None,4:], boxes[:, 4:])[0]

            inds = (sub_iou >= iou_thresh) & (obj_iou >= iou_thresh)

        for i in np.where(keep_inds)[0][inds]:
            pred_to_gt[i].append(int(gt_ind))
    return pred_to_gt

def intersect_2d(x1, x2):
    """
    Given two arrays [m1, n], [m2,n], returns a [m1, m2] array where each entry is True if those
    rows match.
    :param x1: [m1, n] numpy array
    :param x2: [m2, n] numpy array
    :return: [m1, m2] bool array of the intersections
    """
    if x1.shape[1] != x2.shape[1]:
        raise ValueError("Input arrays must have same #columns")

    # This performs a matrix multiplication-esque thing between the two arrays
    # Instead of summing, we want the equality, so we reduce in that way
    res = (x1[..., None] == x2.T[None, ...]).all(1)
    return res

def argsort_desc(scores):
    """
    Returns the indices that sort scores descending in a smart way
    :param scores: Numpy array of arbitrary size
    :return: an array of size [numel(scores), dim(scores)] where each row is the index you'd
             need to get the score.
    """
    return np.column_stack(np.unravel_index(np.argsort(-scores.ravel()), scores.shape))