import os
import cv2
import json
import ctypes
import numpy as np
import pandas as pd
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit  # This is needed for initializing CUDA driver
from backend.utils import timeit, draw_boxed_text

conf_th = 0.3
INPUT_HW = (300, 300)
OUTPUT_LAYOUT=7

with open(os.path.join('models/ssd_mobilenet/labels.json')) as json_data:
    CLASS_NAMES = json.load(json_data)

def _preprocess_trt(img, shape=(300, 300)):
    """Preprocess an image before TRT SSD inferencing."""
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, shape)
    img = img.transpose((2, 0, 1)).astype(np.float32)
    img = (2.0/255.0) * img - 1.0
    return img

class Detector():
    """Class ssd version gpu"""

    def _load_plugins(self):
        ctypes.CDLL("models/ssd_mobilenet/libflattenconcat.so")
        trt.init_libnvinfer_plugins(self.trt_logger, '')

    def _load_engine(self):
        TRTbin = 'models/ssd_mobilenet/TRT_ssd_mobilenet_v2_coco.bin'
        with open(TRTbin, 'rb') as f, trt.Runtime(self.trt_logger) as runtime:
            return runtime.deserialize_cuda_engine(f.read())

    def _create_context(self):
        for binding in self.engine:
            size = trt.volume(self.engine.get_binding_shape(binding)) * \
                   self.engine.max_batch_size
            host_mem = cuda.pagelocked_empty(size, np.float32)
            cuda_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(cuda_mem))
            if self.engine.binding_is_input(binding):
                self.host_inputs.append(host_mem)
                self.cuda_inputs.append(cuda_mem)
            else:
                self.host_outputs.append(host_mem)
                self.cuda_outputs.append(cuda_mem)
        return self.engine.create_execution_context()

    @timeit
    def __init__(self):
        self.colors = np.random.uniform(0, 255, size=(100, 3))
        self.input_shape = INPUT_HW
        self.trt_logger = trt.Logger(trt.Logger.INFO)
        self._load_plugins()
        self.engine = self._load_engine()

        self.host_inputs = []
        self.cuda_inputs = []
        self.host_outputs = []
        self.cuda_outputs = []
        self.bindings = []
        self.stream = cuda.Stream()
        self.context = self._create_context()

    def __del__(self):
        """Free CUDA memories."""
        del self.stream
        del self.cuda_outputs
        del self.cuda_inputs

    @timeit
    def prediction(self, img):
        img_resized = _preprocess_trt(img, self.input_shape)
        np.copyto(self.host_inputs[0], img_resized.ravel())

        cuda.memcpy_htod_async(
            self.cuda_inputs[0], self.host_inputs[0], self.stream)
        self.context.execute_async(
            batch_size=1,
            bindings=self.bindings,
            stream_handle=self.stream.handle)
        cuda.memcpy_dtoh_async(
            self.host_outputs[1], self.cuda_outputs[1], self.stream)
        cuda.memcpy_dtoh_async(
            self.host_outputs[0], self.cuda_outputs[0], self.stream)
        self.stream.synchronize()

        output = self.host_outputs[0]
        return np.reshape(output, (-1, OUTPUT_LAYOUT))

    @timeit
    def filter_prediction(self, output, image, conf_th=0.3, conf_class=[]):
        height, width = image.shape[:-1]
        df = pd.DataFrame(
                output,
                columns=[
                    '_', 'class_id', 'confidence', 'x1', 'y1', 'x2', 'y2'])
        df = df.assign(
                x1=lambda x: (x['x1'] * width).astype(int).clip(0),
                y1=lambda x: (x['y1'] * height).astype(int).clip(0),
                x2=lambda x: (x['x2'] * width).astype(int),
                y2=lambda x: (x['y2'] * height).astype(int),
                class_name=lambda x: (
                    x['class_id'].astype(int).astype(str).replace(CLASS_NAMES)
                    ),
                )
        df['label'] = (df['class_name'] + ': ' +
                df['confidence'].astype(str).str.slice(stop=4))
        df = df[df['confidence'] > conf_th]
        if len(conf_class) > 0:
            df = df[df['class_id'].isin(conf_class)]
        return df

    def draw_boxes(self, image, df):
        for idx, box in df.iterrows():
            x_min, y_min, x_max, y_max = box['x1'], box['y1'], box['x2'], box['y2']
            color = self.colors[int(box['class_id'])]
            cv2.rectangle(image, (x_min, y_min), (x_max, y_max), color, 2)
            txt_loc = (max(x_min+2, 0), max(y_min+2, 0))
            txt = box['label']
            image = draw_boxed_text(image, txt, txt_loc, color)
        return image

if __name__ == "__main__":
    image = cv2.imread("./imgs/image.jpeg")

    detector = Detector()
    output = detector.prediction(image)
    df = detector.filter_prediction(output, image, conf_th=0.3)
    image = detector.draw_boxes(image, df)
    print(df)
    cv2.imwrite("./imgs/outputcv.jpg", image)


# -*- coding: utf-8 -*-
"""
@object: Twitter
@task: Main function of recursive NN (4 classes)
@author: majing
@structure: bottom-up recursive neural networks
@variable: Nepoch, lr, obj, fold
@time: Jan 24, 2018
"""

import os
import sys

import BU_RvNN
import BU_Transformer
import torch.optim as optim
import torch
import numpy as np
import torch.nn as nn
from numpy.testing import assert_array_almost_equal

import time
import datetime
import random
from evaluate import *

obj = "Twitter16" # choose dataset, you can choose either "Twitter15" or "Twitter16"
fold = "2" # fold index, choose from 0-4
tag = ""
vocabulary_size = 5000
hidden_dim = 64
Nclass = 4
Nepoch = 150
lr = 0.005

unit="BU_RvNN-"+obj+str(fold)+'-vol.'+str(vocabulary_size)+tag

treePath = '../resource/data.BU_RvNN.vol_'+str(vocabulary_size)+tag+'.txt'

trainPath = "../nfold/RNNtrainSet_"+obj+str(fold)+"_tree.txt"
testPath = "../nfold/RNNtestSet_"+obj+str(fold)+"_tree.txt"
labelPath = "../resource/"+obj+"_label_All.txt"

################################### tools #####################################
def str2matrix(Str, MaxL): # str = index:wordfreq index:wordfreq
    wordFreq, wordIndex = [], []
    l = 0
    for pair in Str.split(' '):
        wordFreq.append(float(pair.split(':')[1]))
        wordIndex.append(int(pair.split(':')[0]))
        l += 1
    ladd = [ 0 for i in range( MaxL-l ) ]
    wordFreq += ladd
    wordIndex += ladd
    return wordFreq, wordIndex

def loadLabel(label, l1, l2, l3, l4):
    labelset_nonR, labelset_f, labelset_t, labelset_u = ['news', 'non-rumor'], ['false'], ['true'], ['unverified']
    if label in labelset_nonR:
       y_train = [1,0,0,0]
       l1 += 1
    if label in labelset_f:
       y_train = [0,1,0,0]
       l2 += 1
    if label in labelset_t:
       y_train = [0,0,1,0]
       l3 += 1
    if label in labelset_u:
       y_train = [0,0,0,1]
       l4 += 1
    return y_train, l1,l2,l3,l4

def constructTree(tree):
    ## 1. ini tree node
    index2node = {}
    for i in tree:
        node = BU_RvNN.Node_tweet(idx=i)
        index2node[i] = node
    ## 2. construct tree
    for j in tree:
        indexC = j
        indexP = tree[j]['parent']
        nodeC = index2node[indexC]
        wordFreq, wordIndex = str2matrix( tree[j]['vec'], tree[j]['maxL'] )
        nodeC.index = wordIndex
        nodeC.word = wordFreq
        ## not root node ##
        if not indexP == 'None':
           nodeP = index2node[int(indexP)]
           nodeC.parent = nodeP
           nodeP.children.append(nodeC)
        ## root node ##
        else:
           root = nodeC
    ## 3. convert tree to DNN input
    degree = tree[j]['max_degree']
    x_word, x_index, tree = BU_RvNN.gen_nn_inputs(root, max_degree=degree, only_leaves_have_vals=False)
    return x_word, x_index, tree

################################# loas data ###################################
def loadData():
    print("loading tree label",)
    labelDic = {}
    for line in open(labelPath):
        line = line.rstrip()
        label, eid = line.split('\t')[0], line.split('\t')[2]
        labelDic[eid] = label.lower()
    print(len(labelDic))

    print("reading tree") ## X
    treeDic = {}
    for line in open(treePath):
        line = line.rstrip()
        eid, indexP, indexC = line.split('\t')[0], line.split('\t')[1], int(line.split('\t')[2])
        max_degree, maxL, Vec = int(line.split('\t')[3]), int(line.split('\t')[4]), line.split('\t')[5]

        if not treeDic.__contains__(eid):
           treeDic[eid] = {}
        treeDic[eid][indexC] = {'parent':indexP, 'max_degree':max_degree, 'maxL':maxL, 'vec':Vec}
    print('tree no:', len(treeDic))

    print("loading train set", )
    tree_train, word_train, index_train, y_train, c = [], [], [], [], 0
    l1,l2,l3,l4 = 0,0,0,0
    for eid in open(trainPath):
        eid = eid.rstrip()
        if not labelDic.__contains__(eid): continue
        if not treeDic.__contains__(eid): continue
        if len(treeDic[eid]) < 2: continue
        ## 1. load label
        label = labelDic[eid]
        y, l1,l2,l3,l4 = loadLabel(label, l1, l2, l3, l4)
        y_train.append(y)
        ## 2. construct tree
        x_word, x_index, tree = constructTree(treeDic[eid])
        tree_train.append(tree)
        word_train.append(x_word)
        index_train.append(x_index)
        c += 1
    print(l1,l2,l3,l4)

    print("loading test set", )
    tree_test,  word_test, index_test, y_test, c = [], [], [], [], 0
    l1,l2,l3,l4 = 0,0,0,0
    for eid in open(testPath):
        #if c > 4: break
        eid = eid.rstrip()
        if not labelDic.__contains__(eid): continue
        if not treeDic.__contains__(eid): continue
        if len(treeDic[eid]) < 2: continue
        ## 1. load label
        label = labelDic[eid]
        y, l1,l2,l3,l4 = loadLabel(label, l1, l2, l3, l4)
        y_test.append(y)
        ## 2. construct tree
        x_word, x_index, tree = constructTree(treeDic[eid])
        tree_test.append(tree)
        word_test.append(x_word)
        index_test.append(x_index)
        c += 1
    print(l1,l2,l3,l4)
    print("train no:", len(tree_train), len(word_train), len(index_train),len(y_train) )
    print("test no:", len(tree_test), len(word_test), len(index_test), len(y_test))
    print("dim1 for 0:", len(tree_train[0]), len(word_train[0]), len(index_train[0]))
    print("case 0:", tree_train[0][0], word_train[0][0], index_train[0][0])
    return tree_train, word_train, index_train, y_train, tree_test, word_test, index_test, y_test

##################################### MAIN ####################################
def CompLoss(pred, ylabel):
    return (torch.tensor(ylabel, dtype=torch.float) - pred).pow(2).sum()

## 1. load tree & word & index & label
tree_train, word_train, index_train, y_train, tree_test, word_test, index_test, y_test = loadData()
# max_degree = max( list(
#                         map(lambda tree: max(list(map(lambda layer:len(layer), tree))),
#                         tree_train)
#                     )
#                 )
# print("tree_train:", max_degree)
## 2. ini RNN model
t0 = time.time()
# model = BU_RvNN.RvNN(vocabulary_size, hidden_dim, Nclass) #GRU 改用maxpooling之后，twitter16上的最好效果达到81.6%
model = BU_Transformer.AttentionGRU(vocabulary_size, hidden_dim, Nclass)  #AttentionGRU,最好效果７２左右
# model = BU_Transformer.MultiAttentionGRU(vocabulary_size, hidden_dim, Nclass)  #MultiHeadAttentionGRU,最好效果在59左右
# model = BU_Transformer.AttentionGRU(vocabulary_size, hidden_dim, Nclass)  #MultiHeadAttentionFCN
# model = BU_Transformer.TreeLSTM(vocabulary_size, hidden_dim, Nclass)  #TreeLSTM
# model = BU_Transformer.TransformerEncoder(vocabulary_size, hidden_dim, Nclass)
# model = BU_Transformer.TransformerEncoderPoolV2(vocabulary_size, hidden_dim, Nclass)
model = BU_Transformer.StarTransformer(vocabulary_size, hidden_dim, Nclass)
for p in model.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform(p)
best_model = torch.load("../resource/GRU_0.816.pkl")
model.E_bu = best_model.E_bu
del best_model
t1 = time.time()
print('Recursive model established,', (t1-t0)/60)
## 3. looping SGD

paras = [{'param':parameter, 'lr':0.1, 'weight_decay':0.001} if not 'E_bu' in name else {'param':parameter, 'lr':0.05} for name, parameter in model.named_parameters()]
optimizer = optim.Adagrad(model.parameters())
losses_5, losses = [], []
num_examples_seen = 0
batch_size = 5
best_test_acc = 0.0

indexs = [i for i in range(len(y_train))]
for epoch in range(Nepoch):
    ## one SGD
    model.train()
    random.shuffle(indexs)
    for i in range(0, len(indexs), batch_size):
        batch_indexs = indexs[i:min(i+batch_size, len(indexs))]
        loss = torch.tensor(0.0)
        for i in batch_indexs:
            pred = model.forward(word_train[i], index_train[i], tree_train[i])
            loss += CompLoss(pred, y_train[i])
        loss = loss/(1.0*len(batch_indexs))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        losses.append(float(loss.data))
        num_examples_seen += 1
        print("epoch=%d: idx=%d,loss=%f"%( epoch, i, np.mean(losses)))
    sys.stdout.flush()

    model.eval()
    ## cal loss & evaluate
    if epoch % 1 == 0:
       losses_5.append((num_examples_seen, np.mean(losses)))
       time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
       print("%s: Loss after num_examples_seen=%d epoch=%d: %f" % (time, num_examples_seen, epoch, np.mean(losses)))
       sys.stdout.flush()
       prediction = []
       for j in range(len(y_test)):
           prediction.append(model.forward(word_test[j], index_test[j], tree_test[j]).data.tolist() )
       res = evaluation_4class(prediction, y_test)
       if best_test_acc < res[1]:
           best_test_acc = res[1]
           print("best_performance:")
           torch.save(model, "../resource/AttnGRU_T_%.3f.pkl"%best_test_acc)
       print('results:', res)
       sys.stdout.flush()
       ## Adjust the learning rate if loss increases
       if len(losses_5) > 1 and losses_5[-1][1] > losses_5[-2][1]:
          lr = lr * 0.5
          print("Setting learning rate to %f" % lr)
          sys.stdout.flush()
    sys.stdout.flush()
    losses = []


import random
from collections import deque
import numpy as np
import cv2
import time
import tensorflow as tf
import keras.backend.tensorflow_backend as backend
from keras.models import load_model
from ZMyCode import CarEnv, MEMORY_FRACTION

MODEL_PATH = 'models/Xception____53.00max_-223.40avg_-387.00min__1567913413.model'

if __name__ == '__main__':

    # Memory fraction
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=MEMORY_FRACTION)
    backend.set_session(tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)))

    # Load the model
    model = load_model(MODEL_PATH)

    # Create environment
    env = CarEnv()

    # For agent speed measurements - keeps last 60 frametimes
    fps_counter = deque(maxlen=60)

    # Initialize predictions - first prediction takes longer as of initialization that has to be done
    # It's better to do a first prediction then before we start iterating over episode steps
    model.predict(np.ones((1, env.im_height, env.im_width, 3)))

    # Loop over episodes
    while True:

        print('Restarting episode')

        # Reset environment and get initial state
        current_state = env.reset()
        env.collision_hist = []

        done = False

        # Loop over steps
        while True:

            # For FPS counter
            step_start = time.time()

            # Show current frame
            cv2.imshow(f'Agent - preview', current_state)
            cv2.waitKey(1)

            # Predict an action based on current observation space
            qs = model.predict(np.array(current_state).reshape(-1, *current_state.shape)/255)[0]
            action = np.argmax(qs)

            # Step environment (additional flag informs environment to not break an episode by time limit)
            new_state, reward, done, _ = env.step(action)

            # Set current step for next loop iteration
            current_state = new_state

            # If done - agent crashed, break an episode
            if done:
                break

            # Measure step time, append to a deque, then print mean FPS for last 60 frames, q values and taken action
            frame_time = time.time() - step_start
            fps_counter.append(frame_time)
            print(f'Agent: {len(fps_counter)/sum(fps_counter):>4.1f} FPS | Action: [{qs[0]:>5.2f}, {qs[1]:>5.2f}, {qs[2]:>5.2f}] {action}')

        # Destroy an actor at end of episode
        for actor in env.actor_list:
            actor.destroy()


# this file helps with encoding under HLA-I alleles

import numpy as np
import pandas as pd
import random

from random import shuffle
from collections import defaultdict

# the version of padding function used here follows
# bradley's aligning approach as in
# https://github.com/phbradley/conga/blob/master/conga/tcrdist/tcr_distances.py

# contains hard-coded elements:
# max length 27 for CDR3 sequence padding
# and the following two shapes under one-hot-encoding
# (40, 21), the shape of encoded HLA-I alleles
# (27, 21), the shape of encoded CDR3
# (12, 21), the shape of encoded cdr1
# (10, 22), the shape of encoded cdr2
# (6, 21), the shape of encoded cdr25
def encode_set(negative_ori, positive_ori, allele_dict, \
               HLA_enc, V_enc, CDR3len_enc, CDR3_enc, cdr1_enc, cdr2_enc, \
               cdr25_enc, enc_method = "one_hot",\
               rseed = 2134):
    # set random seed
    random.seed(rseed)
    # define padding function for later use
    def pad_cdr3(cdr3, max_length, pad_letter):
        if len(cdr3) == max_length:
            return cdr3
        else:
            gap_start = min( 6, 3 + (len(cdr3)-5)//2 )
            return cdr3[:gap_start] + [pad_letter] * (max_length - len(cdr3)) + \
                cdr3[gap_start:]
    # load info for cdr1, cdr2, cdr25 encoding
    V_info = pd.read_csv("../data/combo_xcr.tsv", sep='\t')
    V_sub_info = V_info.loc[(V_info.organism == 'human') \
                          & (V_info.chain == 'B') \
                          & (V_info.region == 'V')]
    cdr1_dict = defaultdict(str)
    cdr2_dict = defaultdict(str)
    cdr25_dict = defaultdict(str)
    for allele, cdrs in zip(V_sub_info.id.tolist(), V_sub_info.cdrs.tolist()):
        cdr_seqs = cdrs.split(";")
        if len(cdr_seqs) != 4:
            print("length of cdr sequences list is not 4")
            break
        cdr1_dict[allele] = cdr_seqs[0]
        cdr2_dict[allele] = cdr_seqs[1]
        cdr25_dict[allele]= cdr_seqs[2]
    # add the 'not_found' key
    cdr1_dict['not_found']  = ''.join(['.' for i in range(12)])
    cdr2_dict['not_found']  = ''.join(['.' for i in range(10)])
    cdr25_dict['not_found'] = ''.join(['.' for i in range(6)])
    # this file translates the v_allele name to the format in
    # "../data/combo_xcr.tsv"
    df_v_allele_trans = pd.read_csv(\
            "../data/step56_v_allele_translate_table_fillin.csv", \
            header = 0)
    v_allele_trans_dict = defaultdict(str)
    for v_allele, translate in \
        zip(df_v_allele_trans.v_allele_new.tolist(),\
            df_v_allele_trans.v_allele_translate.tolist()):
        v_allele_trans_dict[v_allele] = translate
    # encode negative part
    # HLA
    ne_HLA_aa_list = [allele_dict[name] for _, name in negative_ori]
    # V gene
    ne_V_list = [[tcr.split(",")[0]] for tcr, _ in negative_ori]
    ne_V_array = np.array(ne_V_list)
    ne_V_encoded = V_enc.transform(ne_V_array).toarray()
    # CDR3
    ne_CDR3_list = [pad_cdr3(list(tcr.split(",")[1]), 27, '.') \
                    for tcr, _ in negative_ori]
    # cdr1, cdr2, cdr25
    ne_V_allele_trans = [v_allele_trans_dict[tcr.split(",")[0]] \
                         for tcr, _ in negative_ori]
    ne_cdr1_list = [list(cdr1_dict[trans]) for trans in ne_V_allele_trans]
    ne_cdr2_list = [list(cdr2_dict[trans]) for trans in ne_V_allele_trans]
    ne_cdr25_list = [list(cdr25_dict[trans]) for trans in ne_V_allele_trans]
    # length of cdr3
    ne_CDR3_len_list = [[len(tcr.split(",")[1])] for tcr, _ in negative_ori]
    ne_CDR3_len = np.array(ne_CDR3_len_list)
    ne_CDR3len_encoded = CDR3len_enc.transform(ne_CDR3_len).toarray()
    # encode positive part
    # HLA
    po_HLA_aa_list = [allele_dict[name] for _, name in positive_ori]
    # V gene
    po_V_list = [[tcr.split(",")[0]] for tcr, _ in positive_ori]
    po_V_array = np.array(po_V_list)
    po_V_encoded = V_enc.transform(po_V_array).toarray()
    # CDR3
    po_CDR3_list = [pad_cdr3(list(tcr.split(",")[1]), 27, '.') \
                    for tcr, _ in positive_ori]
    # cdr1, cdr2, cdr25
    po_V_allele_trans = [v_allele_trans_dict[tcr.split(",")[0]] \
                         for tcr, _ in positive_ori]
    po_cdr1_list = [list(cdr1_dict[trans]) for trans in po_V_allele_trans]
    po_cdr2_list = [list(cdr2_dict[trans]) for trans in po_V_allele_trans]
    po_cdr25_list = [list(cdr25_dict[trans]) for trans in po_V_allele_trans]
    # length of cdr3
    po_CDR3_len_list = [[len(tcr.split(",")[1])] for tcr, _ in positive_ori]
    po_CDR3_len = np.array(po_CDR3_len_list)
    po_CDR3len_encoded = CDR3len_enc.transform(po_CDR3_len).toarray()
    # encoding for amino acis in HLA and CDR3
    if enc_method == "one_hot":
        # HLA in negative
        ne_HLA_aa_array = np.array(ne_HLA_aa_list)
        ne_HLA_encoded = HLA_enc.transform(ne_HLA_aa_array).toarray()
        ne_HLA_encoded = ne_HLA_encoded.reshape(len(negative_ori), 40, 21)
        # CDR3 in negative
        ne_CDR3_array = np.array(ne_CDR3_list)
        ne_CDR3_encoded = CDR3_enc.transform(ne_CDR3_array).toarray()
        ne_CDR3_encoded = ne_CDR3_encoded.reshape(len(negative_ori), 27, 21)
        # cdr1, cdr2, cdr25 in negative
        ne_cdr1_array = np.array(ne_cdr1_list)
        ne_cdr1_encoded = cdr1_enc.transform(ne_cdr1_array).toarray()
        ne_cdr1_encoded = ne_cdr1_encoded.reshape(len(negative_ori), 12, 21)
        ne_cdr2_array = np.array(ne_cdr2_list)
        ne_cdr2_encoded = cdr2_enc.transform(ne_cdr2_array).toarray()
        ne_cdr2_encoded = ne_cdr2_encoded.reshape(len(negative_ori), 10, 22)
        ne_cdr25_array = np.array(ne_cdr25_list)
        ne_cdr25_encoded = cdr25_enc.transform(ne_cdr25_array).toarray()
        ne_cdr25_encoded = ne_cdr25_encoded.reshape(len(negative_ori), 6, 21)
        # HLA in positive
        po_HLA_aa_array = np.array(po_HLA_aa_list)
        po_HLA_encoded = HLA_enc.transform(po_HLA_aa_array).toarray()
        po_HLA_encoded = po_HLA_encoded.reshape(len(positive_ori), 40, 21)
        # CDR3 in positive
        po_CDR3_array = np.array(po_CDR3_list)
        po_CDR3_encoded = CDR3_enc.transform(po_CDR3_array).toarray()
        po_CDR3_encoded = po_CDR3_encoded.reshape(len(positive_ori), 27, 21)
        # cdr1, cdr2, cdr25 in positive
        po_cdr1_array = np.array(po_cdr1_list)
        po_cdr1_encoded = cdr1_enc.transform(po_cdr1_array).toarray()
        po_cdr1_encoded = po_cdr1_encoded.reshape(len(positive_ori), 12, 21)
        po_cdr2_array = np.array(po_cdr2_list)
        po_cdr2_encoded = cdr2_enc.transform(po_cdr2_array).toarray()
        po_cdr2_encoded = po_cdr2_encoded.reshape(len(positive_ori), 10, 22)
        po_cdr25_array = np.array(po_cdr25_list)
        po_cdr25_encoded = cdr25_enc.transform(po_cdr25_array).toarray()
        po_cdr25_encoded = po_cdr25_encoded.reshape(len(positive_ori), 6, 21)
        # blosum62 and atchley
    elif enc_method in ["blosum62", "atchley", "pca"]:
        # negative
        ne_HLA_encoded = np.array(HLA_enc(ne_HLA_aa_list))
        ne_CDR3_encoded = np.array(CDR3_enc(ne_CDR3_list))
        ne_cdr1_encoded = np.array(cdr1_enc(ne_cdr1_list))
        ne_cdr2_encoded = np.array(cdr2_enc(ne_cdr2_list))
        ne_cdr25_encoded = np.array(cdr25_enc(ne_cdr25_list))
        # positive
        po_HLA_encoded = np.array(HLA_enc(po_HLA_aa_list))
        po_CDR3_encoded = np.array(CDR3_enc(po_CDR3_list))
        po_cdr1_encoded = np.array(cdr1_enc(po_cdr1_list))
        po_cdr2_encoded = np.array(cdr1_enc(po_cdr2_list))
        po_cdr25_encoded = np.array(cdr1_enc(po_cdr25_list))
    else:
        print("Error: enc_method is a value not coded for yet.")
        return
    # concatenate together
    HLA_encoded = np.concatenate((ne_HLA_encoded, po_HLA_encoded), axis = 0)
    V_encoded = np.concatenate((ne_V_encoded, po_V_encoded), axis = 0)
    CDR3_encoded = np.concatenate((ne_CDR3_encoded, po_CDR3_encoded), axis = 0)
    CDR3len_encoded = np.concatenate((ne_CDR3len_encoded, po_CDR3len_encoded), axis = 0)
    cdr1_encoded = np.concatenate((ne_cdr1_encoded, po_cdr1_encoded), axis = 0)
    cdr2_encoded = np.concatenate((ne_cdr2_encoded, po_cdr2_encoded), axis = 0)
    cdr25_encoded = np.concatenate((ne_cdr25_encoded, po_cdr25_encoded), axis = 0)
    # y label
    ne_y = [0 for _ in range(len(negative_ori))]
    po_y = [1 for _ in range(len(positive_ori))]
    y = np.array(ne_y + po_y).reshape(len(negative_ori) + \
                                              len(positive_ori),1)
    # shuffle
    index_list = list(range(len(negative_ori) + len(positive_ori)))
    shuffle(index_list)
    HLA_encoded = HLA_encoded[index_list]
    V_encoded = V_encoded[index_list]
    CDR3_encoded = CDR3_encoded[index_list]
    CDR3len_encoded = CDR3len_encoded[index_list]
    cdr1_encoded = cdr1_encoded[index_list]
    cdr2_encoded = cdr2_encoded[index_list]
    cdr25_encoded = cdr25_encoded[index_list]
    y = y[index_list]
    # return encoded data
    return HLA_encoded, V_encoded, CDR3_encoded, CDR3len_encoded, \
           cdr1_encoded, cdr2_encoded, cdr25_encoded, y


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

"""Rotor parameters."""

from makani.config import mconfig
from makani.control import system_types as m
import numpy as np

@mconfig.Config(deps={
    'flight_plan': 'common.flight_plan',
    'propellers': 'prop.propellers',
    'wing_serial': 'common.wing_serial',
})
def MakeParams(params):
  # Motor rotor moment-of-inertia [kg-m^2].
  yasa_rotor_moment_of_inertia = 0.33

  bottom_row = [m.kMotorSbo, m.kMotorSbi, m.kMotorPbi, m.kMotorPbo]

  # Assign propeller versions.
  propeller_versions = [None for _ in range(m.kNumMotors)]
  if params['wing_serial'] == m.kWingSerial01:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] in [m.kWingSerial04Crosswind]:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] == m.kWingSerial04Hover:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] in [m.kWingSerial05Crosswind]:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] == m.kWingSerial05Hover:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] in [m.kWingSerial06Crosswind]:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] == m.kWingSerial06Hover:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] in [m.kWingSerial07Crosswind]:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSti] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  elif params['wing_serial'] == m.kWingSerial07Hover:
    propeller_versions[m.kMotorSbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorSbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbi] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPbo] = m.kPropVersionRev4NegativeX
    propeller_versions[m.kMotorPto] = m.kPropVersionRev4PositiveX
    propeller_versions[m.kMotorPti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSti] = m.kPropVersionRev1Trimmed
    propeller_versions[m.kMotorSto] = m.kPropVersionRev4PositiveX
  else:
    assert False, 'Unknown wing serial.'

  rotors = [None for _ in range(m.kNumMotors)]
  for r in range(m.kNumMotors):
    rotors[r] = {
        # Normal vector to the propeller plane.
        'axis': [np.cos(np.deg2rad(3.0)), 0.0, np.sin(np.deg2rad(3.0))],

        # Direction cosine matrix from body to rotor frame.
        'dcm_b2r': {'d': [[np.cos(np.deg2rad(-3.0)), 0.0,
                           np.sin(np.deg2rad(-3.0))],
                          [0.0, 1.0, 0.0],
                          [-np.sin(np.deg2rad(-3.0)), 0.0,
                           np.cos(np.deg2rad(-3.0))]]},

        # Local pressure coefficient [#] at the rotor position.  The
        # pressure coefficient, C_P, is related to local airspeed
        # through the equation:
        #
        #    C_P = 1 - (v / v_freestream)^2
        #
        # There is a significant difference in airspeeds between the top
        # and bottom propellers caused by the lift of the wing. These
        # pressure coefficients are derived from CFD with the slatted
        # kite at 4 deg alpha (https://goo.gl/yfkJJS)
        'local_pressure_coeff': 0.1448 if r in bottom_row else -0.1501,

        # The rotor direction, diameter [m] and moment of inertia [kg
        # m^2] are set from the corresponding propeller's information.
        'version': propeller_versions[r],
        'dir': params['propellers'][propeller_versions[r]]['dir'],
        'D': params['propellers'][propeller_versions[r]]['D'],
        'I': (yasa_rotor_moment_of_inertia +
              params['propellers'][propeller_versions[r]]['I']),
    }
    # We check that the rotor axis is normalized. because it is used
    # to determine the force-moment conversion matrix in
    # rotor_control.py.
    assert abs(np.linalg.norm(rotors[r]['axis']) - 1.0) < 1e-9

  # Rotor positions [m].
  #
  # Updated on 2015-01-22 based on the COM positions given by the Mass
  # and Balance spreadsheet.
  rotors[m.kMotorSbo]['pos'] = [1.613, 3.639, 1.597]
  rotors[m.kMotorSbi]['pos'] = [1.613, 1.213, 1.597]
  rotors[m.kMotorPbi]['pos'] = [1.613, -1.213, 1.597]
  rotors[m.kMotorPbo]['pos'] = [1.613, -3.639, 1.597]
  rotors[m.kMotorPto]['pos'] = [1.960, -3.639, -1.216]
  rotors[m.kMotorPti]['pos'] = [1.960, -1.213, -1.216]
  rotors[m.kMotorSti]['pos'] = [1.960, 1.213, -1.216]
  rotors[m.kMotorSto]['pos'] = [1.960, 3.639, -1.216]

  return rotors


from multiprocessing import Process, Queue
import logging
import time

import cv2

from keras import layers

import numpy as np

import tensorflow as tf

from eye_cropper import EyeCropper
import config

logger = logging.getLogger(__name__)

def _is_stale(timestamp):
  """ Checks if an image is stale.
  Args:
    timestamp: The timestamp of the image to check.
  Returns:
    True if the image is stale, false otherwise. """
  if time.time() - timestamp > config.STALE_THRESHOLD:
    print "Dropping stale frame."
    logger.warning("Dropping stale image from %f at %f." % \
                   (timestamp, time.time()))
    return True

  return False

class GazePredictor(object):
  """ Takes in eye images, and uses the model to predict gaze. """

  def __init__(self, model_file, display=False):
    """
    Args:
      model_file: The saved model to load for predictions.
      display: If true, it will enable a debug display that shows the image
               crops. """
    # Initialize landmark and prediction processes.
    self._prediction_process = _CnnProcess(model_file)
    self._landmark_process = _LandmarkProcess(self._prediction_process,
                                              display=display)

  def __del__(self):
    # Make sure internal processes have terminated.
    self._landmark_process.release()
    self._prediction_process.release()

  def predict_gaze(self):
    """ Predicts the user's gaze based on current frames.
    Returns:
      The predicted gaze point, in cm, the sequence number of the
      corresponding frame, and the timestamp. """
    # Wait for new output from the predictor.
    return self._prediction_process.get_output()

  def process_image(self, image, seq_num):
    """ Adds a new image to the prediction pipeline.
    Args:
      image: The image to add.
      seq_num: The sequence number of the image. """
    # Add a new timestamp for this image. This allows us to drop frames that go
    # stale.
    timestamp = time.time()

    self._landmark_process.add_new_input(image, seq_num, timestamp)

class GazePredictorWithCapture(GazePredictor):
  """ Same as a GazePredictor, except it also handles capturing images from the
  camera. """

  def __init__(self, *args, **kwargs):
    super(GazePredictorWithCapture, self).__init__(*args, **kwargs)

    # Initialize capture process.
    self._capture_process = _CaptureProcess(self._landmark_process)

  def __del__(self):
    super(GazePredictorWithCapture, self).__del__()

    self._capture_process.release()

class _CnnProcess(object):
  """ Runs the CNN prediction in a separate process on the GPU, so that it can
  be handled concurrently. """

  def __init__(self, model_file):
    """
    Args:
      model_file: The file to load the predictor model from. """
    self.__model_file = model_file

    # Create the queues.
    self.__input_queue = Queue()
    self.__output_queue = Queue()

    # Fork the predictor process.
    self.__process = Process(target=_CnnProcess.predict_forever, args=(self,))
    self.__process.start()

  @staticmethod
  def _make_eye_pathway(eye_input):
    """ Generates the preprocessing pathway for a single eye.
    Args:
      eye_input: The raw eye input to be preprocessed.
    Returns:
      The input placeholder, and the output node. """
    # Convert to grayscale.
    eye_gray = tf.image.rgb_to_grayscale(eye_input)
    # Crop it properly.
    eye_cropped = tf.image.crop_to_bounding_box(eye_gray, 3, 3, 218, 218)
    # Resize it properly.
    eye_resized = tf.image.resize_images(eye_cropped, (32, 32),
                                          align_corners=True)
    # Normalize it.
    eye_norm = tf.map_fn(lambda frame: \
                         tf.image.per_image_standardization(frame), eye_resized)

    return eye_norm

  def release(self):
    """ Cleans up and terminates internal process. """
    self.__process.terminate()

  def predict_forever(self):
    """ Generates predictions indefinitely. """
    # Create preprocessing layer.
    eye_preproc_layer = layers.Lambda(_CnnProcess._make_eye_pathway)

    # Load the model we trained.
    model = config.NET_ARCH((224, 224, 3),
                            eye_preproc=eye_preproc_layer)
    self.__predictor = model.build()
    self.__predictor.load_weights(self.__model_file)

    while True:
      self.__predict_once()

  def __predict_once(self):
    """ Reads an image from the input queue, processes it, and writes a
    prediction to the output queue. """
    left_eye, right_eye, face, grid, pose, seq_num, timestamp = \
        self.__input_queue.get()
    if seq_num is None:
      # A None tuple means the end of the sequence. Propagate this through the
      # pipeline.
      self.__output_queue.put((None, None, None))
      return
    if left_eye is None:
      # If we have a sequence number but no images, we got a bad detection.
      self.__output_queue.put((None, seq_num, timestamp))
      return
    if _is_stale(timestamp):
      # The image is stale, so indicate that it is invalid.
      self.__output_queue.put((None, seq_num, timestamp))
      return

    # Convert everything to floats.
    left_eye = left_eye.astype(np.float32)
    right_eye = right_eye.astype(np.float32)
    face = face.astype(np.float32)

    # Add the batch dimension.
    left_eye = np.expand_dims(left_eye, axis=0)
    right_eye = np.expand_dims(right_eye, axis=0)
    face = np.expand_dims(face, axis=0)
    grid = np.expand_dims(grid, axis=0)
    pose = np.expand_dims(pose, axis=0)

    # Generate a prediction.
    pred = self.__predictor.predict([left_eye, right_eye, face, grid, pose],
                                    batch_size=1)
    # Remove the batch dimension, and convert to Python floats.
    pred = [float(x) for x in pred[0]]

    self.__output_queue.put((pred, seq_num, timestamp))

  def add_new_input(self, left_eye, right_eye, face, grid, pose, seq_num,
                    timestamp):
    """ Adds a new input to be processed. Will block.
    Args:
      left_eye: The left eye crop.
      right_eye: The right eye crop.
      face: The face crop.
      grid: The face grid.
      pose: The head pose.
      seq_num: The sequence number of the image.
      timestamp: The timestamp of the image.
    """
    self.__input_queue.put((left_eye, right_eye, face, grid, pose, seq_num,
                            timestamp))

  def get_output(self):
    """ Gets an output from the prediction process. Will block.
    Returns:
      The predicted gaze point and the sequence number. """
    return self.__output_queue.get()

class _LandmarkProcess(object):
  """ Reads images from a queue, and runs landmark detection in a separate
  process. """

  def __init__(self, cnn_process, display=False):
    """
    Args:
      cnn_process: The _CnnProcess to send captured images to.
      display: If true, it will enable a debugging display that shows the
               detected crops on-screen. """
    self.__cnn_process = cnn_process
    self.__display = display

    # Create the queues.
    self.__input_queue = Queue(maxsize=1)
    self.__output_queue = Queue()

    # Fork the capture process.
    self.__process = Process(target=_LandmarkProcess.run_forever,
                             args=(self,))
    self.__process.start()

  def release(self):
    """ Cleans up and terminates internal process. """
    self.__process.terminate()

  def __run_once(self):
    """ Reads and crops a single image. It will send it to the predictor
    process when finished. """
    # Get the next input from the queue.
    image, seq_num, timestamp = self.__input_queue.get()
    if image is None:
      # A None tuple means the end of a sequence. Propagate this through the
      # pipeline.
      self.__cnn_process.add_new_input(None, None, None, None, None, None, None)
      return
    if _is_stale(timestamp):
      # Image is stale. Indicate that it is invalid.
      self.__cnn_process.add_new_input(None, None, None, None, None, seq_num,
                                       timestamp)

    # Crop the image.
    left_eye, right_eye, face = self.__cropper.crop_image(image)
    if face is None:
      # We failed to get an image.
      logger.warning("Failed to get good detection for %d." % (seq_num))
      # Send along the sequence number.
      self.__cnn_process.add_new_input(None, None, None, None, None, seq_num,
                                       timestamp)
      return

    # Produce face mask.
    mask = self.__cropper.face_grid()
    # Estimate the head pose.
    head_pose = self.__cropper.estimate_pose()
    head_pose = head_pose[:, 0]

    if self.__display:
      # Show the debugging display.
      mask_sized = cv2.resize(mask, (224, 224))
      mask_sized = np.expand_dims(mask_sized, axis=2)
      mask_sized = np.tile(mask_sized, [1, 1, 3]) * 225.0
      mask_sized = mask_sized.astype(np.uint8)

      combined = np.concatenate((left_eye, right_eye, face, mask_sized), axis=1)
      cv2.imshow("Server Detections", combined)
      cv2.waitKey(1)

    # Send it along.
    self.__cnn_process.add_new_input(left_eye, right_eye, face, mask,
                                     head_pose, seq_num, timestamp)

  def run_forever(self):
    """ Reads and crops images indefinitely. """
    # Eye cropper to use for eye detection.
    self.__cropper = EyeCropper()

    while True:
      self.__run_once()

  def add_new_input(self, image, seq_num, timestamp):
    """ Adds a new input to be processed. Will block.
    Args:
      image: The image to process.
      seq_num: The sequence number of the image.
      timestamp: The timestamp of the image. """
    self.__input_queue.put((image, seq_num, timestamp))

class _CaptureProcess(object):
  """ Process that captures images from the camera, sends them over a queue. """

  def __init__(self, landmark_process, camera=-1):
    """
    Args:
      landmark_process: The landmark detection process to send images to.
      camera: The camera device to capture from. """
    self.__landmark_process = landmark_process
    self.__camera = camera

    # Fork the capture process.
    self.__process = Process(target=_CaptureProcess.run_forever,
                             args=(self,))
    self.__process.start()

  def __run_once(self):
    """ Runs a single iteration of the capture process. """
    ret, image = self.__capture.read()
    if not ret:
      raise RuntimeError("Could not read from camera.")

    # Send the image along.
    timestamp = time.time()
    self.__landmark_process.add_new_input(image, self.__sequence_num, timestamp)

    # Update the sequence number.
    self.__sequence_num += 1
    self.__sequence_num %= 255

  def run_forever(self):
    """ Runs the process indefinitely. """
    # Capture device to use.
    self.__capture = cv2.VideoCapture(self.__camera)
    # Sequence number for images.
    self.__sequence_num = 0

    while True:
      self.__run_once()

  def release(self):
    """ Cleans up and terminates internal process. """
    self.__process.terminate()


#!/usr/bin/env python
import sys
sys.path.append('C:\Program Files (x86)\Keysight\SD1\Libraries\Python')

from BaseDriver import LabberDriver, Error, IdError
import keysightSD1

import numpy as np

class Driver(LabberDriver):
    """ This class implements the Keysight PXI digitizer"""

    def performOpen(self, options={}):
        """Perform the operation of opening the instrument connection"""
        # set time step and resolution
        self.nBit = 16
        self.bitRange = float(2**(self.nBit-1)-1)
        # timeout
        self.timeout_ms = int(1000 * self.dComCfg['Timeout'])
        # get PXI chassis
        self.chassis = int(self.dComCfg.get('PXI chassis', 1))
        # create AWG instance
        self.dig = keysightSD1.SD_AIN()
        AWGPart = self.dig.getProductNameBySlot(
            self.chassis, int(self.comCfg.address))
        self.log('Serial:', self.dig.getSerialNumberBySlot(
            self.chassis, int(self.comCfg.address)))
        if not isinstance(AWGPart, str):
            raise Error('Unit not available')
        # check that model is supported
        dOptionCfg = self.dInstrCfg['options']
        for validId, validName in zip(dOptionCfg['model_id'], dOptionCfg['model_str']):
            if AWGPart.find(validId)>=0:
                # id found, stop searching
                break
        else:
            # loop fell through, raise ID error
            raise IdError(AWGPart, dOptionCfg['model_id'])
        # set model
        self.setModel(validName)
        # sampling rate and number of channles is set by model
        if validName in ('M3102', 'M3302'):
            # 500 MHz models
            self.dt = 2E-9
            self.nCh = 4
        else:
            # assume 100 MHz for all other models
            self.dt = 10E-9
            self.nCh = 4
        # create list of sampled data
        self.lTrace = [np.array([])] * self.nCh
        self.dig.openWithSlot(AWGPart, self.chassis, int(self.comCfg.address))
        # get hardware version - changes numbering of channels
        hw_version = self.dig.getHardwareVersion()
        if hw_version >= 4:
            # KEYSIGHT - channel numbers start with 1
            self.ch_index_zero = 1
        else:
            # SIGNADYNE - channel numbers start with 0
            self.ch_index_zero = 0
        self.log('HW:', hw_version)

    def getHwCh(self, n):
        """Get hardware channel number for channel n. n starts at 0"""
        return n + self.ch_index_zero

    def performClose(self, bError=False, options={}):
        """Perform the close instrument connection operation"""
        # do not check for error if close was called with an error
        try:
            # flush all memory
            for n in range(self.nCh):
                self.log('Close ch:', n, self.dig.DAQflush(self.getHwCh(n)))
            # close instrument
            self.dig.close()
        except:
            # never return error here
            pass

    def performSetValue(self, quant, value, sweepRate=0.0, options={}):
        """Perform the Set Value instrument operation. This function should
        return the actual value set by the instrument"""
        # start with setting local quant value
        quant.setValue(value)
        # check if channel-specific, if so get channel + name
        if quant.name.startswith('Ch') and len(quant.name)>6:
            ch = int(quant.name[2]) - 1
            name = quant.name[6:]
        else:
            ch, name = None, ''
        # proceed depending on command
        if quant.name in ('External Trig Source', 'External Trig Config',
                          'Trig Sync Mode'):
            extSource = int(self.getCmdStringFromValue('External Trig Source'))
            trigBehavior = int(self.getCmdStringFromValue('External Trig Config'))
            sync = int(self.getCmdStringFromValue('Trig Sync Mode'))
            self.dig.DAQtriggerExternalConfig(0, extSource, trigBehavior, sync)
        elif quant.name in ('Trig I/O', ):
            # get direction and sync from index of comboboxes
            direction = int(self.getCmdStringFromValue('Trig I/O'))
            self.dig.triggerIOconfig(direction)
        elif quant.name in ('Analog Trig Channel', 'Analog Trig Config', 'Trig Threshold'):
            # get trig channel
            trigCh = self.getValueIndex('Analog Trig Channel')
            mod = int(self.getCmdStringFromValue('Analog Trig Config'))
            threshold = self.getValue('Trig Threshold')
            self.dig.channelTriggerConfig(self.getHwCh(trigCh), mod, threshold)
        elif name in ('Range', 'Impedance', 'Coupling'):
            # set range, impedance, coupling at once
            rang = self.getRange(ch)
            imp = int(self.getCmdStringFromValue('Ch%d - Impedance' % (ch + 1)))
            coup = int(self.getCmdStringFromValue('Ch%d - Coupling' % (ch + 1)))
            self.dig.channelInputConfig(self.getHwCh(ch), rang, imp, coup)
        return value

    def performGetValue(self, quant, options={}):
        """Perform the Set Value instrument operation. This function should
        return the actual value set by the instrument"""
        # check if channel-specific, if so get channel + name
        if quant.name.startswith('Ch') and len(quant.name) > 6:
            ch = int(quant.name[2]) - 1
            name = quant.name[6:]
        else:
            ch, name = None, ''

        if name == 'Signal':
            if self.isHardwareLoop(options):
                return self.getSignalHardwareLoop(ch, quant, options)
            # get traces if first call
            if self.isFirstCall(options):
                # don't arm if in hardware trig mode
                self.getTraces(bArm=(not self.isHardwareTrig(options)))
            # return correct data
            value = quant.getTraceDict(self.lTrace[ch], dt=self.dt)
        else:
            # for all others, return local value
            value = quant.getValue()

        return value

    def performArm(self, quant_names, options={}):
        """Perform the instrument arm operation"""
        # make sure we are arming for reading traces, if not return
        signal_names = ['Ch%d - Signal' % (n + 1) for n in range(4)]
        signal_arm = [name in signal_names for name in quant_names]
        if not np.any(signal_arm):
            return

        # arm by calling get traces
        if self.isHardwareLoop(options):
            # in hardware looping, number of records is set by the hw loop
            (seq_no, n_seq) = self.getHardwareLoopIndex(options)
            nSample = int(self.getValue('Number of samples'))

            # arm instrument, then report completed to allow client to continue
            self.reportStatus('Digitizer - Waiting for signal')
            self.getTraces(bArm=True, bMeasure=False, n_seq=n_seq)
            self.report_arm_completed()
            # directly start collecting data (digitizer buffer is limited)
            self.getTraces(bArm=False, bMeasure=True, n_seq=n_seq)
            # re-shape data and place in trace buffer
            self.reshaped_traces = []
            for trace in self.lTrace:
                if len(trace) > 0:
                    trace = trace.reshape((n_seq, nSample))
                self.reshaped_traces.append(trace)

        else:
            self.getTraces(bArm=True, bMeasure=False)

    def getTraces(self, bArm=True, bMeasure=True, n_seq=0):
        """Get all active traces"""
        # # test timing
        # import time
        # t0 = time.perf_counter()
        # lT = []

        # find out which traces to get
        lCh = []
        iChMask = 0
        for n in range(self.nCh):
            if self.getValue('Ch%d - Enabled' % (n + 1)):
                lCh.append(n)
                iChMask += 2**n
        # get current settings
        nPts = int(self.getValue('Number of samples'))
        nCyclePerCall = int(self.getValue('Records per Buffer'))
        # in hardware loop mode, ignore records and use number of sequences
        if n_seq > 0:
            nSeg = n_seq
        else:
            nSeg = int(self.getValue('Number of records'))

        nAv = int(self.getValue('Number of averages'))
        # trigger delay is in 1/sample rate
        nTrigDelay = int(round(self.getValue('Trig Delay') / self.dt))

        if bArm:
            # clear old data
            self.dig.DAQflushMultiple(iChMask)
            self.lTrace = [np.array([])] * self.nCh
            # configure trigger for all active channels
            for nCh in lCh:
                # init data
                self.lTrace[nCh] = np.zeros((nSeg * nPts))
                # channel number depens on hardware version
                ch = self.getHwCh(nCh)
                # extra config for trig mode
                if self.getValue('Trig Mode') == 'Digital trigger':
                    extSource = int(self.getCmdStringFromValue('External Trig Source'))
                    trigBehavior = int(self.getCmdStringFromValue('External Trig Config'))
                    sync = int(self.getCmdStringFromValue('Trig Sync Mode'))
                    self.dig.DAQtriggerExternalConfig(ch, extSource, trigBehavior, sync)
                    self.dig.DAQdigitalTriggerConfig(ch, extSource, trigBehavior)
                elif self.getValue('Trig Mode') == 'Analog channel':
                    digitalTriggerMode= 0
                    digitalTriggerSource = 0
                    trigCh = self.getValueIndex('Analog Trig Channel')
                    analogTriggerMask = 2**trigCh
                    self.dig.DAQtriggerConfig(ch, digitalTriggerMode, digitalTriggerSource, analogTriggerMask)
                # config daq and trig mode
                trigMode = int(self.getCmdStringFromValue('Trig Mode'))
                self.dig.DAQconfig(ch, nPts, nSeg*nAv, nTrigDelay, trigMode)
            # start acquiring data
            self.dig.DAQstartMultiple(iChMask)
        # lT.append('Start %.1f ms' % (1000*(time.perf_counter()-t0)))
        #
        # return if not measure
        if not bMeasure:
            return
        # define number of cycles to read at a time
        nCycleTotal = nSeg * nAv
        nCall = int(np.ceil(nCycleTotal / nCyclePerCall))
        lScale = [(self.getRange(ch) / self.bitRange) for ch in range(self.nCh)]
        # keep track of progress in percent
        old_percent = 0

        # proceed depending on segment or not segment
        if nSeg <= 1:
            # non-segmented acquisiton
            for n in range(nCall):
                # number of cycles for this call, could be fewer for last call
                nCycle = min(nCyclePerCall, nCycleTotal - (n * nCyclePerCall))

                # report progress, only report integer percent
                if nCall > 100:
                    new_percent = int(100 * n / nCall)
                    if new_percent > old_percent:
                        old_percent = new_percent
                        self.reportStatus(
                            'Acquiring traces ({}%)'.format(new_percent))

                # capture traces one by one
                for nCh in lCh:
                    # channel number depens on hardware version
                    ch = self.getHwCh(nCh)
                    data = self.DAQread(self.dig, ch, nPts * nCycle,
                                        int(1000 + self.timeout_ms / nCall))
                    # stop if no data
                    if data.size == 0:
                        return

                    # average
                    data = data.reshape((nCycle, nPts)).mean(0)
                    # adjust scaling to account for summing averages
                    scale = lScale[nCh] * (nCycle / nAv)
                    # convert to voltage, add to total average
                    self.lTrace[nCh] += data * scale

                # break if stopped from outside
                if self.isStopped():
                    break
                # lT.append('N: %d, Tot %.1f ms' % (n, 1000 * (time.perf_counter() - t0)))

        else:
            # segmented acquisition, get caLls per segment
            (nCallSeg, extra_call) = divmod(nSeg, nCyclePerCall)
            # pre-calculate list of cycles/call, last call may have more cycles
            if nCallSeg == 0:
                nCallSeg = 1
                lCyclesSeg = [nSeg]
            else:
                lCyclesSeg = [nCyclePerCall] * nCallSeg
                lCyclesSeg[-1] = nCyclePerCall + extra_call
            # pre-calculate scale, should include scaling for averaging
            lScale = np.array(lScale, dtype=float) / nAv

            for n in range(nAv):
                # report progress, only report integer percent
                if nAv > 1:
                    new_percent = int(100 * n / nAv)
                    if new_percent > old_percent:
                        old_percent = new_percent
                        self.reportStatus(
                            'Acquiring traces ({}%)'.format(new_percent))

                count = 0
                # loop over number of calls per segment
                for m, nCycle in enumerate(lCyclesSeg):

                    # capture traces one by one
                    for nCh in lCh:
                        # channel number depens on hardware version
                        ch = self.getHwCh(nCh)
                        data = self.DAQread(self.dig, ch, nPts * nCycle,
                                            int(1000 + self.timeout_ms / nCall))
                        # stop if no data
                        if data.size == 0:
                            return
                        # store all data in one long vector
                        self.lTrace[nCh][count:(count + data.size)] += \
                            data * lScale[nCh]

                    count += data.size

                # break if stopped from outside
                if self.isStopped():
                    break

                # lT.append('N: %d, Tot %.1f ms' % (n, 1000 * (time.perf_counter() - t0)))

        # # log timing info
        # self.log(': '.join(lT))

    def getRange(self, ch):
        """Get channel range, as voltage.  Index start at 0"""
        rang = float(self.getCmdStringFromValue('Ch%d - Range' % (ch + 1)))
        # range depends on impedance
        if self.getValue('Ch%d - Impedance' % (ch + 1)) == 'High':
            rang = rang * 2
            # special case if range is .25, 0.5, or 1, scale to 0.2, .4, .8
            if rang < 1.1:
                rang *= 0.8
        return rang

    def DAQread(self, dig, nDAQ, nPoints, timeOut):
        """Read data diretly to numpy array"""
        if dig._SD_Object__handle > 0:
            if nPoints > 0:
                data = (keysightSD1.c_short * nPoints)()
                nPointsOut = dig._SD_Object__core_dll.SD_AIN_DAQread(dig._SD_Object__handle, nDAQ, data, nPoints, timeOut)
                if nPointsOut > 0:
                    return np.frombuffer(data, dtype=np.int16, count=nPoints)
                else:
                    return np.array([], dtype=np.int16)
            else:
                return keysightSD1.SD_Error.INVALID_VALUE
        else:
            return keysightSD1.SD_Error.MODULE_NOT_OPENED

    def getSignalHardwareLoop(self, ch, quant, options):
        """Get data from round-robin type averaging"""
        (seq_no, n_seq) = self.getHardwareLoopIndex(options)
        # after getting data, pick values to return
        return quant.getTraceDict(self.reshaped_traces[ch][seq_no], dt=self.dt)

if __name__ == '__main__':
    pass


import cv2
import numpy as np

def mask_plate(image_path):
    original = cv2.imread(image_path)
    resized = image_resize(original, width=2000)  # Image resize for optimal OCR results
    cropped = crop_to_rect(resized)  # Used to crop image, now it just masks the image

    blur = cv2.GaussianBlur(cropped, (7, 7), 0)  # Blurring after resize for smoother image
    frame_hsv = cv2.cvtColor(blur, cv2.COLOR_BGR2HSV)
    low_h, low_s, low_v = 12, 0, 0  # Low threshold for yellow mask
    high_h, high_s, high_v = 40, 255, 255  # High threshold for yellow mask
    frame_threshold = cv2.inRange(frame_hsv, (low_h, low_s, low_v), (high_h, high_s, high_v))  # Yellow mask
    blur = cv2.GaussianBlur(frame_threshold, (9, 9), 0)  # Mask blurring to reduce noise
    ret3, th3 = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)  # Thresholding

    number_of_white_pix = np.sum(frame_threshold == 255)  # Extracting only white pixels
    number_of_black_pix = np.sum(frame_threshold == 0)  # Extracting only black pixels

    # This checks whether the plate is yellow
    if number_of_white_pix / (number_of_white_pix + number_of_black_pix) > 0.4:
        blur = cv2.GaussianBlur(th3, (9, 9), 0)
    else:  # B&W plate
        frame_bgr = cv2.cvtColor(frame_hsv, cv2.COLOR_HSV2BGR)
        blur = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2GRAY)

    kernel = np.ones((2, 2), np.uint8)
    eroded = cv2.erode(blur, kernel, iterations=2)  # Erosion to make digits bold
    masked = cv2.bitwise_and(cropped, cropped, mask=eroded)  # Masking plate image
    gray = cv2.cvtColor(masked, cv2.COLOR_BGR2GRAY)  # Masked image to grayscale
    bfilter = cv2.bilateralFilter(gray, 77, 17, 17)  # Noise reduction

    # Additional image processing to make digits bold and reduce noise
    kernel = np.ones((3, 3), np.uint8)
    dilated = cv2.dilate(bfilter, kernel, iterations=4)
    blur = cv2.GaussianBlur(dilated, (3, 3), 0)
    smooth = cv2.addWeighted(blur, 1.5, dilated, -0.5, 0)
    blur = cv2.GaussianBlur(smooth, (9, 9), 0)
    ret3, th3 = cv2.threshold(blur, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    kernel = np.ones((4, 4), np.uint8)
    eroded = cv2.erode(th3, kernel, iterations=5)
    blur = cv2.GaussianBlur(eroded, (21, 21), 0)
    kernel = np.ones((5, 5), np.uint8)
    dilated = cv2.dilate(blur, kernel, iterations=6)
    ret, img = cv2.threshold(dilated, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)  # Apply threshold

    # Saving masked image to file and returning its path
    cv2.imwrite("JPGs\\masked_plate.jpg", img)
    return "JPGs\\masked_plate.jpg"

def image_resize(image, width=None, height=None, inter=cv2.INTER_AREA):
    if width is None:
        width = 800
    # initialize the dimensions of the image to be resized and
    # grab the image size
    dim = None
    (h, w) = image.shape[:2]

    # if both the width and height are None, then return the
    # original image
    if width is None and height is None:
        return image

    # check to see if the width is None
    if width is None:
        # calculate the ratio of the height and construct the
        # dimensions
        r = height / float(h)
        dim = (int(w * r), height)

    # otherwise, the height is None
    else:
        # calculate the ratio of the width and construct the
        # dimensions
        r = width / float(w)
        dim = (width, int(h * r))

    # resize the image
    resized = cv2.resize(image, dim, interpolation=inter)

    # return the resized image
    return resized

def crop_to_rect(image):
    # Turn image to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Apply threshold
    ret, img = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    # Find contours in img.
    cnts = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)[-2]

    # Find the contour with the maximum area.
    sorted_contours = sorted(cnts, key=cv2.contourArea, reverse=True)[:10]
    mask = np.zeros_like(img)  # Create mask where white is what we want, black otherwise
    cv2.drawContours(mask, sorted_contours, 0, 255, -1)  # Draw filled contour in mask
    out = np.zeros_like(img)  # Extract out the object and place into output image
    out[mask == 255] = img[mask == 255]

    # # Now crop | Uncomment this code block and comment the line below it to crop instead
    # (y, x) = np.where(mask == 255)
    # (topy, topx) = (np.min(y), np.min(x))
    # (bottomy, bottomx) = (np.max(y), np.max(x))
    # # out = image[topy:bottomy + 1, topx:bottomx + 1]
    out = cv2.bitwise_and(image, image, mask=mask)  # Masking plate image

    return out


from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.utils as vutils
from torch.autograd import Variable
from model import _netlocalD, _netG
from pre_data import pre_data
from tqdm import tqdm
import numpy as np

os.environ["CUDA_VISIBLE_DEVICES"] = "1"
flag_use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if flag_use_cuda else "cpu")

class Trainer:
    def __init__(self, opt) -> None:
        self.opt = opt
        self.init_model()
        self.init_cfg()

    def init_cfg(self):
        if self.opt.manualSeed is None:
            self.opt.manualSeed = random.randint(1, 10000)
        print("Random Seed: ", self.opt.manualSeed)
        random.seed(self.opt.manualSeed)
        torch.manual_seed(self.opt.manualSeed)
        if flag_use_cuda:
            torch.cuda.manual_seed_all(self.opt.manualSeed)
        cudnn.benchmark = True

        # setup loss & optimizer
        self.criterion = nn.BCELoss()
        self.criterionMSE = nn.MSELoss()
        self.optimizerD = optim.Adam(self.netD.parameters(), lr=self.opt.lr)
        self.optimizerG = optim.Adam(self.netG.parameters(), lr=self.opt.lr)

        #setpu paras
        self.wtl2 = float(self.opt.wtl2)
        self.overlapL2Weight = 10
        self.input_real = torch.FloatTensor(self.opt.batchSize, 3, self.opt.imageSize, self.opt.imageSize)
        self.input_cropped = torch.FloatTensor(self.opt.batchSize, 3, self.opt.imageSize, self.opt.imageSize)
        self.label = torch.FloatTensor(self.opt.batchSize)
        self.real_label = 1
        self.fake_label = 0
        self.real_center = torch.FloatTensor(self.opt.batchSize, 3, int(self.opt.imageSize/2), int(self.opt.imageSize/2))
        if flag_use_cuda:
            self.input_real, self.input_cropped,self.label = self.input_real.to(device),self.input_cropped.to(device), self.label.to(device)
            self.real_center = self.real_center.to(device)
            self.criterion.to(device)
            self.criterionMSE.to(device)
            print("Using %s" % device)
        self.input_real = Variable(self.input_real)
        self.input_cropped = Variable(self.input_cropped)
        self.label = Variable(self.label)
        self.real_center = Variable(self.real_center)

    def init_model(self):
        self.netG = _netG(self.opt)
        # self.netG = _netG_block(self.opt)
        if self.opt.netG != '':
            self.netG.load_state_dict(torch.load(self.opt.netG, map_location=lambda storage, location: storage)['state_dict'])
        self.netD = _netlocalD(self.opt)
        if self.opt.netD != '':
            self.netD.load_state_dict(torch.load(self.opt.netD,map_location=lambda storage, location: storage)['state_dict'])
        if flag_use_cuda:
            self.netD.to(device)
            self.netG.to(device)

    def train(self, dataloader_train, dataloader_valid = None, iter_max = None):
        iter_max = iter_max if iter_max is not None else self.opt.niter
        self.schedulerD = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizerD, iter_max)
        self.schedulerG = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizerG, iter_max)
        for epoch in range(iter_max):
            pbar = tqdm(total = len(dataloader_train))
            for i, data in enumerate(dataloader_train):
                real_cpu, _ = data
                self.real_center_cpu = real_cpu[:,:,int(self.opt.imageSize/4):int(self.opt.imageSize/4)+int(self.opt.imageSize/2),int(self.opt.imageSize/4):int(self.opt.imageSize/4)+int(self.opt.imageSize/2)]
                batch_size = real_cpu.size(0)
                self.input_real.resize_(real_cpu.size()).copy_(real_cpu)
                self.input_cropped.resize_(real_cpu.size()).copy_(real_cpu)
                self.real_center.resize_(self.real_center_cpu.size()).copy_(self.real_center_cpu)
                self.input_cropped[:,0,int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred),int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred)] = 2*117.0/255.0 - 1.0
                self.input_cropped[:,1,int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred),int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred)] = 2*104.0/255.0 - 1.0
                self.input_cropped[:,2,int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred),int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred)] = 2*123.0/255.0 - 1.0

                # train with real
                self.netD.zero_grad()
                self.label.resize_(batch_size).fill_(self.real_label)
                self.label = torch.reshape(self.label, (self.label.shape[0], 1))

                output = self.netD(self.real_center)
                errD_real = self.criterion(output, self.label)
                errD_real.backward()
                D_x = output.mean()

                # train with fake
                fake = self.netG(self.input_cropped)
                self.label.fill_(self.fake_label)
                output = self.netD(fake.detach())
                errD_fake = self.criterion(output, self.label)
                errD_fake.backward()
                D_G_z1 = output.mean()
                errD = errD_real + errD_fake
                self.optimizerD.step()

                ############################
                # (2) Update G network: maximize log(D(G(z)))
                ###########################
                self.netG.zero_grad()
                self.label.fill_(self.real_label)  # fake labels are real for generator cost
                output = self.netD(fake)
                errG_D = self.criterion(output, self.label)
                # errG_D.backward(retain_variables=True)

                errG_norm = self.criterionMSE(fake,self.real_center)
                self.wtl2Matrix = self.real_center.clone()
                self.wtl2Matrix.fill_(self.wtl2*self.overlapL2Weight)
                self.wtl2Matrix[:,:,int(self.opt.overlapPred):int(self.opt.imageSize/2 - self.opt.overlapPred),int(self.opt.overlapPred):int(self.opt.imageSize/2 - self.opt.overlapPred)] = self.wtl2

                errG_l2 = (fake-self.real_center).pow(2)
                errG_l2 = errG_l2 * self.wtl2Matrix
                errG_l2 = errG_l2.mean()

                errG = (1-self.wtl2) * errG_D + self.wtl2 * errG_l2

                errG.backward()

                D_G_z2 = output.mean()
                self.optimizerG.step()
                # print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f / %.4f l_D(x): %.4f l_D(G(z)): %.4f'
                #     % (epoch, iter_max, i, len(dataloader_train),
                #         errD.item(), errG_D.item(),errG_l2.item(), D_x,D_G_z1, ))
                pbar.update(1)
                pbar.set_description("[%d/%d][%d/%d], errG = %.4f"%(epoch, iter_max, i, len(dataloader_train), errG_norm.item()))
            pbar.close()
            self.schedulerD.step()
            self.schedulerG.step()

            if dataloader_valid is not None:
                self.test(dataloader_valid, epoch)

        # do checkpointing
        torch.save({'epoch': 0,
                    'state_dict':self.netG.state_dict()},
                    'result/model/self.netG_streetview.pth' )
        torch.save({'epoch': 0,
                    'state_dict':self.netD.state_dict()},
                    'result/model/netlocalD.pth' )

    def test(self, dataloader_test, epoch, flag_save = True):
        with torch.no_grad():
            errG_list = []
            for i, data in enumerate(dataloader_test):
                real_cpu, _ = data
                self.real_center_cpu = real_cpu[:,:,int(self.opt.imageSize/4):int(self.opt.imageSize/4)+int(self.opt.imageSize/2),int(self.opt.imageSize/4):int(self.opt.imageSize/4)+int(self.opt.imageSize/2)]
                batch_size = real_cpu.size(0)
                self.input_real.resize_(real_cpu.size()).copy_(real_cpu)
                self.input_cropped.resize_(real_cpu.size()).copy_(real_cpu)
                self.real_center.resize_(self.real_center_cpu.size()).copy_(self.real_center_cpu)
                self.input_cropped[:,0,int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred),int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred)] = 2*117.0/255.0 - 1.0
                self.input_cropped[:,1,int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred),int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred)] = 2*104.0/255.0 - 1.0
                self.input_cropped[:,2,int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred),int(self.opt.imageSize/4+self.opt.overlapPred):int(self.opt.imageSize/4+self.opt.imageSize/2-self.opt.overlapPred)] = 2*123.0/255.0 - 1.0

                fake = self.netG(self.input_cropped)
                errG = self.criterionMSE(fake, self.real_center)
                recon_image = self.input_cropped.clone()
                recon_image[:,:,int(self.opt.imageSize/4):int(self.opt.imageSize/4+self.opt.imageSize/2),int(self.opt.imageSize/4):int(self.opt.imageSize/4+self.opt.imageSize/2)] = fake
                errG_list.append(errG.item())

                if flag_save and i < 5:
                    vutils.save_image(real_cpu, 'result/real/real_samples_batch_%03d_%03d.png' % (epoch, i), normalize=True)
                    vutils.save_image(self.input_cropped, 'result/cropped/cropped_samples_batch_%03d_%03d.png' % (epoch, i), normalize=True)
                    vutils.save_image(recon_image, 'result/recon/recon_center_samples_batch_%03d_%03d.png' % (epoch, i), normalize=True)
            print("errG = ", np.mean(errG_list))

    # # custom weights initialization called on self.netG and self.netD
    # def weights_init(self, m):
    #     classname = m.__class__.__name__
    #     if classname.find('Conv') != -1:
    #         m.weight.data.normal_(0.0, 0.02)
    #     elif classname.find('BatchNorm') != -1:
    #         m.weight.data.normal_(1.0, 0.02)
    #         m.bias.data.fill_(0)