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

import argparse
import os
import sys

import mxnet as mx
import numpy as np
import sklearn
from mxnet import ndarray as nd

sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'common'))
import face_image

def main(args):
    ctx = []
    cvd = os.environ['CUDA_VISIBLE_DEVICES'].strip()
    if len(cvd) > 0:
        for i in range(len(cvd.split(','))):
            ctx.append(mx.gpu(i))
    if len(ctx) == 0:
        ctx = [mx.cpu()]
        print('use cpu')
    else:
        print('gpu num:', len(ctx))
    ctx_num = len(ctx)
    path_imgrec = os.path.join(args.input, 'train.rec')
    path_imgidx = os.path.join(args.input, 'train.idx')
    imgrec = mx.recordio.MXIndexedRecordIO(path_imgidx, path_imgrec, 'r')  # pylint: disable=redefined-variable-type
    outf = open(os.path.join(args.input, 'c2c'), 'w')
    s = imgrec.read_idx(0)
    header, _ = mx.recordio.unpack(s)
    assert header.flag > 0
    print('header0 label', header.label)
    header0 = (int(header.label[0]), int(header.label[1]))
    # assert(header.flag==1)
    imgidx = range(1, int(header.label[0]))
    id2range = {}
    seq_identity = range(int(header.label[0]), int(header.label[1]))
    for identity in seq_identity:
        s = imgrec.read_idx(identity)
        header, _ = mx.recordio.unpack(s)
        id2range[identity] = (int(header.label[0]), int(header.label[1]))
    print('id2range', len(id2range))
    prop = face_image.load_property(args.input)
    image_size = prop.image_size
    print('image_size', image_size)
    vec = args.model.split(',')
    prefix = vec[0]
    epoch = int(vec[1])
    print('loading', prefix, epoch)
    model = mx.mod.Module.load(prefix, epoch, context=ctx)
    model.bind(data_shapes=[('data', (args.batch_size, 3, image_size[0], image_size[1]))], label_shapes=[('softmax_label', (args.batch_size,))])
    nrof_images = 0
    nrof_removed = 0
    idx = 1
    id2label = {}
    pp = 0
    for _id, v in id2range.iteritems():
        pp += 1
        if pp % 100 == 0:
            print('processing id', pp)
        _list = range(*v)
        ocontents = []
        for i in range(len(_list)):
            _idx = _list[i]
            # print('_idx', _id, _idx)
            s = imgrec.read_idx(_idx)
            ocontents.append(s)
        # continue
        embeddings = None
        headers = [None] * len(ocontents)
        # print(len(ocontents))
        ba = 0
        while True:
            bb = min(ba + args.batch_size, len(ocontents))
            if ba >= bb:
                break
            _batch_size = bb - ba
            _batch_size2 = max(_batch_size, ctx_num)
            data = nd.zeros((_batch_size2, 3, image_size[0], image_size[1]))
            label = nd.zeros((_batch_size2,))
            count = bb - ba
            ii = 0
            for i in range(ba, bb):
                header, img = mx.recordio.unpack(ocontents[i])
                headers[i] = header
                img = mx.image.imdecode(img)
                img = nd.transpose(img, axes=(2, 0, 1))
                data[ii][:] = img
                label[ii][:] = header.label[0]
                ii += 1
            while ii < _batch_size2:
                data[ii][:] = data[0][:]
                label[ii][:] = label[0][:]
                ii += 1
            db = mx.io.DataBatch(data=(data,), label=(label,))
            model.forward(db, is_train=False)
            net_out = model.get_outputs()
            net_out = net_out[0].asnumpy()
            if embeddings is None:
                embeddings = np.zeros((len(ocontents), net_out.shape[1]))
            embeddings[ba:bb, :] = net_out[0:_batch_size, :]
            ba = bb
        embeddings = sklearn.preprocessing.normalize(embeddings)
        emb_mean = np.mean(embeddings, axis=0, keepdims=True)
        emb_mean = sklearn.preprocessing.normalize(emb_mean)
        sim = np.dot(embeddings, emb_mean.T)
        # print(sim.shape)
        sims = sim.flatten()
        assert len(_list) == len(sims)
        assert len(_list) == len(ocontents)
        for i in range(len(ocontents)):
            _sim = sims[i]
            _idx = _list[i]
            _header = headers[i]
            # TODO
            outf.write("%d,%f,%d\n" % (_idx, _sim, int(_header.label[1])))
    outf.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='')
    # general
    parser.add_argument('--input', default='', type=str, help='')
    parser.add_argument('--model', default='../model/softmax,50', help='path to load model.')
    parser.add_argument('--batch-size', default=32, type=int, help='')
    args = parser.parse_args()
    main(args)


#Author : Lewis Mervin lhm30@cam.ac.uk
#Supervisor : Dr. A. Bender
#All rights reserved 2016
#Protein Target Prediction Tool trained on SARs from PubChem (Mined 21/06/16) and ChEMBL21
#Molecular Descriptors : 2048bit Morgan Binary Fingerprints (Rdkit) - ECFP4
#Dependencies : rdkit, sklearn, numpy

#libraries
from rdkit import Chem
from rdkit.Chem import AllChem
import cPickle
import glob
import zipfile
import os
import sys
import math
import numpy as np
from multiprocessing import Pool
import multiprocessing

def introMessage():
	print '=============================================================================================='
	print ' Author: Lewis Mervin\n Email:  lhm30@cam.ac.uk\n Supervisor: Dr. A. Bender'
	print ' Address: Centre For Molecular Informatics, Dept. Chemistry, Lensfield Road, Cambridge CB2 1EW'
	print '==============================================================================================\n'
	return

#calculate 2048bit morgan fingerprints, radius 2
def calcFingerprints(smiles):
	m1 = Chem.MolFromSmiles(smiles)
	fp = AllChem.GetMorganFingerprintAsBitVect(m1,2, nBits=2048)
	binary = fp.ToBitString()
	return list(binary)

#calculate fingerprints for chunked array of smiles
def arrayFP(inp):
	outfp = []
	outsmi = []
	for i in inp:
		try:
			outfp.append(calcFingerprints(i))
			outsmi.append(i)
		except:
			print 'SMILES Parse Error: ' + i
	return outfp,outsmi

#import user query
def importQuery(in_file):
	query = open(in_file).read().splitlines()
	#collect IDs, if present
	if len(query[0].split()) > 1:
		ids = [line.split()[1] for line in query]
		query = [line.split()[0] for line in query]
	else:
		ids = None
	matrix = np.empty((len(query), 2048), dtype=np.uint8)
	smiles_per_core = int(math.ceil(len(query) / N_cores)+1)
	chunked_smiles = [query[x:x+smiles_per_core] for x in xrange(0, len(query), smiles_per_core)]
	pool = Pool(processes=N_cores)  # set up resources
	jobs = pool.imap(arrayFP, chunked_smiles)
	current_end = 0
	processed_smi = []
	for i, result in enumerate(jobs):
		matrix[current_end:current_end+len(result[0]), :] = result[0]
		current_end += len(result[0])
		processed_smi += result[1]
	pool.close()
	pool.join()
	#remove IDs of SMILES parsing errors
	if ids:
		processed_ids = []
		for idx, smi in enumerate(query):
			if smi in processed_smi:
				processed_ids.append(ids[idx])
		ids = processed_ids
	#if IDs weren't present, use SMILES as IDs
	else:
		ids = processed_smi
	return matrix[:current_end], processed_smi, ids

#get info for uniprots
def getUniprotInfo():
	if os.name == 'nt': sep = '\\'
	else: sep = '/'
	model_info = [l.split('\t') for l in open(os.path.dirname(os.path.abspath(__file__)) + sep + 'classes_in_model.txt').read().splitlines()]
	return_dict = {l[0] : l[0:8] for l in model_info}
	return return_dict

#unzip a pkl model
def open_Model(mod):
	if os.name == 'nt': sep = '\\'
	else: sep = '/'
	with zipfile.ZipFile(os.path.dirname(os.path.abspath(__file__)) + sep + 'models' + sep + mod + '.pkl.zip', 'r') as zfile:
		with zfile.open(mod + '.pkl', 'r') as fid:
			clf = cPickle.load(fid)
	return clf

#prediction worker
def doTargetPrediction(pickled_model_name):
	if os.name == 'nt': sep = '\\'
	else: sep = '/'
	mod = pickled_model_name.split(sep)[-1].split('.')[0]
	clf = open_Model(mod)
	probs = clf.predict_proba(querymatrix)[:,1]
	preds = map(int,probs > threshold)
	if sum(preds) > 0:
		return mod,preds
	else: return None

#prediction runner
def performTargetPrediction(models):
	prediction_results = []
	pool = Pool(processes=N_cores, initializer=initPool, initargs=(querymatrix,threshold,))  # set up resources
	jobs = pool.imap_unordered(doTargetPrediction, models)
	for i, result in enumerate(jobs):
		percent = (float(i)/float(len(models)))*100 + 1
		sys.stdout.write(' Performing Classification on Query Molecules: %3d%%\r' % percent)
		sys.stdout.flush()
		if result is not None: prediction_results.append(result)
	pool.close()
	pool.join()
	return prediction_results

#initializer for the pool
def initPool(querymatrix_, threshold_):
	global querymatrix, threshold
	querymatrix = querymatrix_
	threshold = threshold_

#main
if __name__ == '__main__':
	if os.name == 'nt': sep = '\\'
	else: sep = '/'
	input_name = sys.argv[1]
	N_cores = int(sys.argv[2])
	introMessage()
	print ' Predicting Targets for ' + input_name
	print ' Using ' + str(N_cores) + ' Cores'
	try:
		threshold = float(sys.argv[3])
	except ValueError:
		print 'ERROR: Enter a valid float (max 2 decimal places) for threshold'
		quit()
	try:
		desired_organism = sys.argv[4]
	except IndexError:
		desired_organism = None
	models = [modelfile for modelfile in glob.glob(os.path.dirname(os.path.abspath(__file__)) + sep + 'models' + sep + '*.zip')]
	model_info = getUniprotInfo()
	if desired_organism is not None:
		models = [mod for mod in models if model_info[mod.split(sep)[-1].split('.')[0]][4] == desired_organism]
		print ' Predicting for organism : ' + desired_organism
		output_name = input_name + '_out_binary_' +  str(threshold) + '_' + desired_organism[:3] + '.txt'
	else:
		output_name = input_name + '_out_binary_' +  str(threshold) + '.txt'
	print ' Total Number of Classes : ' + str(len(models))
	print ' Using TPR threshold of : ' + str(threshold)
	out_file = open(output_name, 'w')
	querymatrix,smiles,ids = importQuery(input_name)
	print ' Total Number of Query Molecules : ' + str(len(querymatrix))
	prediction_results = performTargetPrediction(models)
	print str(len(prediction_results))
	out_file.write('Uniprot\tPref_Name\tGene ID\tTarget_Class\tOrganism\tPDB_ID\tDisGeNET_Diseases_0.06\tChEMBL_First_Published\t' + '\t'.join(map(str,ids)) + '\n')
	for row in sorted(prediction_results):
		out_file.write('\t'.join(map(str,model_info[row[0]])) + '\t' + '\t'.join(map(str,row[1])) + '\n')
	print '\n Wrote Results to: ' + output_name
	out_file.close()


#coding:utf-8
import pinyin as py
import numpy as np
import textgrid
phonemes = ['b', 'p', 'f', 'm', \
            'd', 't', 'n', 'l', \
            'g', 'k', 'h', \
            'j', 'q', 'x', \
            'zh', 'ch', 'sh', 'r', \
            'z', 'c', 's',\
            'a',  'ai', 'ao',  'an',  'ang', \
            'o',  'ou', 'ong', \
            'e',  'ei', 'en',  'eng', 'er', 'ev', \
            'i',  'ix', 'iii', \
            'ia', 'iao','ian', 'iang','ie', \
            'in', 'ing','io',  'iou', 'iong', \
            'u',  'ua', 'uo',  'uai', 'uei', \
            'uan','uen','uang','ueng', \
            'v',  've', 'van', 'vn', \
            'ng', 'mm', 'nn',\
            'rr', 'sp']
initialOrFinal = ['0', '1', '2', '3']
pitch = ['0','50','51','52','53','54','55','56','57','58','59',\
         '60','61','62','63','64','65','66','67','68','69',\
         '70','71','72','73','74','75','76','77','78','79','80']
frameLength = 20 # 设置帧长为20ms帧移为5ms
frameHop = 5

# one-hot编码
def onehotEncoding(instance, class1):
    temp = [0] * len(class1)
    temp[class1.index(instance)] = 1
    return temp

durationInput = []
# 时长模型输入（以音素为单位）
with open('./myData/train.txt', encoding='utf-8') as f:
    metadatas = [line.strip().split('|') for line in f]
# 音素one-hot编码
# 音素类型（声母、韵母、零声母）one-hot编码
# 拼接[音素，音素类型，所属音符理论时长]
for line in metadatas:
    phoneme = onehotEncoding(line[-1], phonemes)
    soundType = onehotEncoding(line[-2], initialOrFinal)
    time = [float(line[0])]
    durationInput.append(phoneme+soundType+time)
# 时长模型预期输出：每个音素的开始、结束时间（训练时由人工标注信息获得），最开始一般会比乐谱多一个sil
with open('./myData/001.interval', encoding='utf-8') as f:
    i = 0
    j = 0
    durationOutput = []
    for line in f:
        if j != 12:
            j = j+1
            continue
        line = line.split('\n')[0]
        if i == 0:
            startTime = float(line)
            i = i+1
        elif i == 1:
            endTime = float(line)
            i = i+1
        else:
            i = 0
            if j == 12:
                durationOutput.append([startTime, endTime, line.split('"')[1]])
            else:
                if durationOutput[-1][2] != line.split('"')[1]:
                    durationOutput.append([startTime, endTime, line.split('"')[1]])
                else:
                    durationOutput[-1][1] = endTime

# 声学模型输入（以帧为单位），训练时由人工标注获得，reference时由时长模型的结果来获得

# 各帧的音素one-hot编码
# 此帧所属音符的音调（C4,G3）的one-hot编码
# 此帧位置P0：1.在音素中的前比例，2.在音素中的后比例，3.音素在整个语音中的比例，都为[0,1]间的浮点数
# P0的embedding
# 拼接[音素，音调，位置embedding]


import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

from . import bases

from previous_works.HoHoNet.lib.misc import panostretch, post_proc
from ..utils import peaks_finding
from scipy.ndimage.filters import maximum_filter
from shapely.geometry import Polygon

''' Layout (per-column) estimation '''
class LayoutEstimator(nn.Module):
    def __init__(self, emb_dim, bon_weight=1., cor_weight=1., bon_loss='l1', cor_loss='bce', bon_scale=1.,
                 init_weight=0.1, dropout=0., oneconv=True, last_ks=1, last_bias=True,
                 H=512, W=1024, post_force_cuboid=False):
        super(LayoutEstimator, self).__init__()
        self.bon_loss = bon_loss
        self.cor_loss = cor_loss
        self.bon_scale = bon_scale
        self.bon_weight = bon_weight
        self.cor_weight = cor_weight
        self.H = H
        self.W = W
        self.post_force_cuboid = post_force_cuboid

        if oneconv:
            self.pred_bon = nn.Conv1d(emb_dim, 2, last_ks, padding=last_ks//2, bias=last_bias)
            self.pred_cor = nn.Conv1d(emb_dim, 1, last_ks, padding=last_ks//2, bias=last_bias)
            if last_bias:
                nn.init.constant_(self.pred_bon.bias[0], -0.478)
                nn.init.constant_(self.pred_bon.bias[1], 0.425)
                nn.init.constant_(self.pred_cor.bias, -1.)
        else:
            self.pred_bon = nn.Sequential(
                nn.Conv1d(emb_dim, emb_dim, 3, padding=1, bias=False),
                nn.BatchNorm1d(emb_dim),
                nn.ReLU(inplace=True),
                nn.Conv1d(emb_dim, 2, 1),
            )
            self.pred_cor = nn.Sequential(
                nn.Conv1d(emb_dim, emb_dim, 3, padding=1, bias=False),
                nn.BatchNorm1d(emb_dim),
                nn.ReLU(inplace=True),
                nn.Conv1d(emb_dim, 1, 1),
            )
            nn.init.constant_(self.pred_bon[-1].bias[0], -0.478)
            nn.init.constant_(self.pred_bon[-1].bias[1], 0.425)
            nn.init.constant_(self.pred_cor[-1].bias, -1.)
        self.dropout = None
        if dropout > 0:
            self.dropout = nn.Dropout(dropout)

    def forward(self, x_emb):
        x_emb = x_emb['1D']
        if self.dropout is not None:
            x_emb = self.dropout(x_emb)
        pred_bon = self.pred_bon(x_emb)
        pred_cor = self.pred_cor(x_emb)
        return {'bon': pred_bon, 'cor': pred_cor}

    def infer(self, x_emb):
        pred = self(x_emb)
        pred_bon = pred['bon'] / self.bon_scale
        pred_cor = pred['cor']
        H, W = self.H, self.W

        y_bon_ = (pred_bon[0].cpu().numpy() / np.pi + 0.5) * H - 0.5
        y_cor_ = pred_cor[0,0].sigmoid().cpu().numpy()
        # Init floor/ceil plane
        z0 = 50
        _, z1 = post_proc.np_refine_by_fix_z(*y_bon_, z0)

        # Detech wall-wall peaks
        def find_N_peaks(signal, r, min_v, N):
            max_v = maximum_filter(signal, size=r, mode='wrap')
            pk_loc = np.where(max_v == signal)[0]
            pk_loc = pk_loc[signal[pk_loc] > min_v]
            if N is not None:
                order = np.argsort(-signal[pk_loc])
                pk_loc = pk_loc[order[:N]]
                pk_loc = pk_loc[np.argsort(pk_loc)]
            return pk_loc, signal[pk_loc]
        min_v = 0 if self.post_force_cuboid else 0.05
        r = int(round(W * 0.05 / 2))
        N = 4 if self.post_force_cuboid else None
        xs_ = find_N_peaks(y_cor_, r=r, min_v=min_v, N=N)[0]

        # Generate wall-walls
        cor, xy_cor = post_proc.gen_ww(xs_, y_bon_[0], z0, tol=abs(0.16 * z1 / 1.6), force_cuboid=self.post_force_cuboid)
        if not self.post_force_cuboid:
            # Check valid (for fear self-intersection)
            xy2d = np.zeros((len(xy_cor), 2), np.float32)
            for i in range(len(xy_cor)):
                xy2d[i, xy_cor[i]['type']] = xy_cor[i]['val']
                xy2d[i, xy_cor[i-1]['type']] = xy_cor[i-1]['val']
            if not Polygon(xy2d).is_valid:
                import sys
                print(
                    'Fail to generate valid general layout!! '
                    'Generate cuboid as fallback.',
                    file=sys.stderr)
                xs_ = find_N_peaks(y_cor_, r=r, min_v=0, N=4)[0]
                cor, xy_cor = post_proc.gen_ww(xs_, y_bon_[0], z0, tol=abs(0.16 * z1 / 1.6), force_cuboid=True)

        # Expand with btn coory
        cor = np.hstack([cor, post_proc.infer_coory(cor[:, 1], z1 - z0, z0)[:, None]])
        # Collect corner position in equirectangular
        cor_id = np.zeros((len(cor)*2, 2), np.float32)
        for j in range(len(cor)):
            cor_id[j*2] = cor[j, 0], cor[j, 1]
            cor_id[j*2 + 1] = cor[j, 0], cor[j, 2]
        return {'cor_id': cor_id, 'y_bon_': y_bon_, 'y_cor_': y_cor_}

    def compute_losses(self, x_emb, batch):
        gt_bon = batch['bon'] * self.bon_scale
        gt_vot = batch['vot']
        gt_cor = 0.96 ** gt_vot.abs()

        # Forward
        pred = self(x_emb)

        # Compute losses
        losses = {}
        if self.bon_loss == 'l1':
            losses['bon'] = F.l1_loss(pred['bon'], gt_bon)
        elif self.bon_loss == 'l2':
            losses['bon'] = F.mse_loss(pred['bon'], gt_bon)
        else:
            raise NotImplementedError

        if self.cor_loss == 'bce':
            losses['cor'] = F.binary_cross_entropy_with_logits(pred['cor'], gt_cor)
        elif self.cor_loss == 'prfocal':
            g, p = gt_cor, pred['cor']
            pos_mask = (g >= 1-1e-6)
            B, alpha, beta = len(g), 2, 4
            L_pos = -F.logsigmoid(p) * F.sigmoid(-p).pow(alpha)
            L_neg = -F.logsigmoid(-p) * F.sigmoid(p).pow(alpha) * (1-g).pow(beta)
            L = torch.where(pos_mask, L_pos, L_neg).view(B,-1).sum(-1) / pos_mask.float().view(B,-1).sum(-1)
            losses['cor'] = L.mean()
        else:
            raise NotImplementedError

        losses['total.layout'] = self.bon_weight * losses['bon'] + self.cor_weight * losses['cor']
        with torch.no_grad():
            losses['bon.mae'] = F.l1_loss(pred['bon'], gt_bon) / self.bon_scale
            losses['cor.mae'] = F.l1_loss(pred['cor'].sigmoid(), gt_cor)
        return losses


from __future__ import division
import time
import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import cv2
from model.utils import *
import argparse
import os
import os.path as osp
from darknet import Darknet
import pickle as pkl
import random

def arg_parse():
    """
    Parse arguements to the detect module

    """

    parser = argparse.ArgumentParser(description='YOLO v3 Detection Module')
    parser.add_argument("--bs", dest="bs", help="Batch size", default=1)
    parser.add_argument("--confidence", dest="confidence",
                        help="Object Confidence to filter predictions", default=0.5)
    parser.add_argument("--nms_thresh", dest="nms_thresh",
                        help="NMS Threshhold", default=0.4)
    parser.add_argument("--cfg", dest='cfgfile', help="Config file",
                        default="cfg/yolov3.cfg", type=str)
    parser.add_argument("--weights", dest='weightsfile', help="weightsfile",
                        default="yolov3.weights", type=str)
    parser.add_argument("--reso", dest='reso', help="Input resolution of the network. Increase to increase accuracy. Decrease to increase speed",
                        default="416", type=str)
    parser.add_argument("--video", dest="videofile",
                        help="Video file to     run detection on", default="video.avi", type=str)

    return parser.parse_args()

args = arg_parse()
batch_size = int(args.bs)
confidence = float(args.confidence)
nms_thesh = float(args.nms_thresh)
start = 0
CUDA = torch.cuda.is_available()

num_classes = 80
classes = load_classes("data/coco.names")

# Set up the neural network
print("Loading network.....")
model = Darknet(args.cfgfile)
# model.print_network()
model.load_weights(args.weightsfile)
print("Network successfully loaded")

model.net_info["height"] = args.reso
inp_dim = int(model.net_info["height"])
assert inp_dim % 32 == 0
assert inp_dim > 32

# If there's a GPU availible, put the model on GPU
if CUDA:
    model.cuda()

# Set the model in evaluation mode
model.eval()

def write(x, results):
    c1 = tuple(x[1:3].int())
    c2 = tuple(x[3:5].int())
    img = results
    cls = int(x[-1])
    # color = random.choice(colors)
    color = (255, 255, 255)
    label = "{0}".format(classes[cls])
    cv2.rectangle(img, c1, c2, color, 1)
    t_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_PLAIN, 1, 1)[0]
    c2 = c1[0] + t_size[0] + 3, c1[1] + t_size[1] + 4
    cv2.rectangle(img, c1, c2, color, -1)
    cv2.putText(img, label, (c1[0], c1[1] + t_size[1] + 4),
                cv2.FONT_HERSHEY_PLAIN, 1, [0, 0, 0], 1)
    return img

# Detection phase

videofile = args.videofile  # or path to the video file.

# cap = cv2.VideoCapture(videofile)

cap = cv2.VideoCapture(0)

assert cap.isOpened(), 'Cannot capture source'

frames = 0
start = time.time()

while cap.isOpened():
    ret, frame = cap.read()

    if ret:
        img = prep_image(frame, inp_dim)
#        cv2.imshow("a", frame)
        im_dim = frame.shape[1], frame.shape[0]
        im_dim = torch.FloatTensor(im_dim).repeat(1, 2)

        if CUDA:
            im_dim = im_dim.cuda()
            img = img.cuda()

        with torch.no_grad():
            output = model(Variable(img, volatile=True), CUDA)
        output = write_results(
            output, confidence, num_classes, nms_conf=nms_thesh)

        if type(output) == int:
            frames += 1
            print("FPS of the video is {:5.4f}".format(
                frames / (time.time() - start)))
            cv2.imshow("frame", frame)
            key = cv2.waitKey(1)
            if key & 0xFF == ord('q'):
                break
            continue

        im_dim = im_dim.repeat(output.size(0), 1)
        scaling_factor = torch.min(416/im_dim, 1)[0].view(-1, 1)

        output[:, [1, 3]] -= (inp_dim - scaling_factor *
                              im_dim[:, 0].view(-1, 1))/2
        output[:, [2, 4]] -= (inp_dim - scaling_factor *
                              im_dim[:, 1].view(-1, 1))/2

        output[:, 1:5] /= scaling_factor

        for i in range(output.shape[0]):
            output[i, [1, 3]] = torch.clamp(
                output[i, [1, 3]], 0.0, im_dim[i, 0])
            output[i, [2, 4]] = torch.clamp(
                output[i, [2, 4]], 0.0, im_dim[i, 1])

        classes = load_classes('data/coco.names')
        # colors = pkl.load(open("pallete", "rb"))

        list(map(lambda x: write(x, frame), output))

        cv2.imshow("frame", frame)
        key = cv2.waitKey(1)
        if key & 0xFF == ord('q'):
            break
        frames += 1
        print(time.time() - start)
        print("FPS of the video is {:5.2f}".format(
            frames / (time.time() - start)))
    else:
        break


"""
Name: Neural networks file.
Description: This file contains neural network classes.
Version: [release][3.1]
Source url: https://github.com/OPHoperHPO/image-background-remove-tool
Author: Anodev (OPHoperHPO)[https://github.com/OPHoperHPO] .
License: Apache License 2.0
License:
   Copyright 2020 OPHoperHPO

   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.
"""
import logging
import os
import time
from skimage import io, transform
import numpy as np
from PIL import Image
from io import BytesIO

tf = None  # Define temp modules
ndi = None
torch = None
Variable = None
U2NET_DEEP = None
U2NETP_DEEP = None

logger = logging.getLogger(__name__)
MODELS_NAMES = ["u2net", "u2netp", "xception_model", "mobile_net_model"]

# noinspection PyUnresolvedReferences
def model_detect(model_name):
    """Detects which model to use and returns its object"""
    global tf, ndi, torch, Variable, U2NET_DEEP, U2NETP_DEEP
    if model_name in MODELS_NAMES:
        if model_name == "xception_model" or model_name == "mobile_net_model":
            import scipy.ndimage as ndi
            import tensorflow as tf
            return TFSegmentation(model_name)
        elif "u2net" in model_name:
            import torch
            from torch.autograd import Variable
            from libs.u2net import U2NET as U2NET_DEEP
            from libs.u2net import U2NETP as U2NETP_DEEP
            return U2NET(model_name)
    else:
        return False

# noinspection PyUnresolvedReferences
class U2NET:
    """U^2-Net model interface"""

    def __init__(self, name="u2net"):
        if name == 'u2net':  # Load model
            logger.debug("Loading a U2NET model (176.6 mb) with better quality but slower processing.")
            net = U2NET_DEEP()
        elif name == 'u2netp':
            logger.debug("Loading a U2NETp model (4 mb) with lower quality but fast processing.")
            net = U2NETP_DEEP()
        else:
            logger.debug("Loading a Unknown model!")
        try:
            if torch.cuda.is_available():
                net.load_state_dict(torch.load(os.path.join("models", name, name + '.pth')))
                net.cuda()
            else:
                net.load_state_dict(torch.load(os.path.join("models", name, name + '.pth'), map_location="cpu"))
        except FileNotFoundError:
            raise FileNotFoundError("No pre-trained model found! Run setup.sh or setup.bat to download it!")
        net.eval()
        self.__net__ = net  # Define model

    def process_image(self, path):
        """
        Removes background from image and returns PIL RGBA Image.
        :param path: Path to image
        :return: PIL RGBA Image. If an error reading the image is detected, returns False.
        """
        start_time = time.time()  # Time counter
        logger.debug("Load image: {}".format(path))
        image, org_image = self.__load_image__(path)  # Load image
        if image is False or org_image is False:
            return False
        image = image.type(torch.FloatTensor)
        if torch.cuda.is_available():
            image = Variable(image.cuda())
        else:
            image = Variable(image)
        mask, d2, d3, d4, d5, d6, d7 = self.__net__(image)  # Predict mask
        logger.debug("Mask prediction completed")
        # Normalization
        logger.debug("Mask normalization")
        mask = mask[:, 0, :, :]
        mask = self.__normalize__(mask)
        # Prepare mask
        logger.debug("Prepare mask")
        mask = self.__prepare_mask__(mask, org_image.size)
        # Apply mask to image
        logger.debug("Apply mask to image")
        empty = Image.new("RGBA", org_image.size)
        image = Image.composite(org_image, empty, mask)
        logger.debug("Finished! Time spent: {}".format(time.time() - start_time))
        return image

    def __load_image__(self, path: str):
        """
        Loads an image file for other processing
        :param path: Path to image file
        :return: image tensor, original image shape
        """
        image_size = 320  # Size of the input and output image for the model
        try:
            image = io.imread(path)  # Load image
        except IOError:
            logger.error('Cannot retrieve image. Please check file: ' + path)
            return False, False
        pil_image = Image.fromarray(image)
        image = transform.resize(image, (image_size, image_size), mode='constant')  # Resize image
        image = self.__ndrarray2tensor__(image)  # Convert image from numpy arr to tensor
        return image, pil_image

    @staticmethod
    def __ndrarray2tensor__(image: np.ndarray):
        """
        Converts a NumPy array to a tensor
        :param image: Image numpy array
        :return: Image tensor
        """
        tmp_img = np.zeros((image.shape[0], image.shape[1], 3))
        image /= np.max(image)
        if image.shape[2] == 1:
            tmp_img[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
            tmp_img[:, :, 1] = (image[:, :, 0] - 0.485) / 0.229
            tmp_img[:, :, 2] = (image[:, :, 0] - 0.485) / 0.229
        else:
            tmp_img[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
            tmp_img[:, :, 1] = (image[:, :, 1] - 0.456) / 0.224
            tmp_img[:, :, 2] = (image[:, :, 2] - 0.406) / 0.225
        tmp_img = tmp_img.transpose((2, 0, 1))
        tmp_img = np.expand_dims(tmp_img, 0)
        return torch.from_numpy(tmp_img)

    # noinspection PyUnresolvedReferences
    @staticmethod
    def __normalize__(predicted):
        """Normalize the predicted map"""
        ma = torch.max(predicted)
        mi = torch.min(predicted)
        out = (predicted - mi) / (ma - mi)
        return out

    @staticmethod
    def __prepare_mask__(predict, image_size):
        """Prepares mask"""
        predict = predict.squeeze()
        predict_np = predict.cpu().data.numpy()
        mask = Image.fromarray(predict_np * 255).convert("L")
        mask = mask.resize(image_size, resample=Image.BILINEAR)
        return mask


#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division

import onmt
import onmt.markdown
import torch
import argparse
import math
import numpy
import sys
import h5py as h5
import numpy as np
import apex
from onmt.model_factory import build_model, build_language_model, optimize_model
from onmt.inference.fast_translator import FastTranslator
from onmt.inference.stream_translator import StreamTranslator

parser = argparse.ArgumentParser(description='translate.py')
onmt.markdown.add_md_help_argument(parser)

parser.add_argument('-model', required=True,
                    help='Path to model .pt file')

parser.add_argument('-fp16', action='store_true',
                    help='To use floating point 16 in decoding')
parser.add_argument('-gpu', type=int, default=-1,
                    help="Device to run on")
parser.add_argument('-batch_size', type=int, default=1,
                    help='Batch size')

def main():
    opt = parser.parse_args()
    opt.cuda = opt.gpu > -1
    if opt.cuda:
        torch.cuda.set_device(opt.gpu)

    if opt.output == "stdout":
        outF = sys.stdout
    else:
        outF = open(opt.output, 'w')

    model = opt.model
    checkpoint = torch.load(model,
                            map_location=lambda storage, loc: storage)

    model_opt = checkpoint['opt']

    model = build_model(model_opt, checkpoint['dicts'])
    optimize_model(model)
    model.load_state_dict(checkpoint['model'])

    if opt.fp16:
        model = model.half()

    if opt.cuda:
        model = model.cuda()
    else:
        model = model.cpu()

    model.eval()
    if opt.encoder_type == "audio" and opt.asr_format == "scp":
        import kaldiio
        from kaldiio import ReadHelper
        audio_data = iter(ReadHelper('scp:' + opt.src))

    in_file = None
    src_batch = []
    if opt.encoder_type == "audio":
        i = 0
        while True:
            if opt.asr_format == "h5":
                if i == len(in_file):
                    break
                line = np.array(in_file[str(i)])
                i += 1
            elif opt.asr_format == "scp":
                try:
                    name, line = next(audio_data)
                except StopIteration:
                    break

        if opt.stride != 1:
            line = line[0::opt.stride]
        line = torch.from_numpy(line)
        if opt.concat != 1:
            add = (opt.concat - line.size()[0] % opt.concat) % opt.concat
            z = torch.FloatTensor(add, line.size()[1]).zero_()
            line = torch.cat((line, z), 0)
            line = line.reshape((line.size()[0] // opt.concat, line.size()[1] * opt.concat))

        output = model.inference(line)
        # wave_glow to save wav file


#!/usr/bin/env python
"""
Python interface to mp_fit_arb.c.

Hazen 01/18
"""

import ctypes
import numpy
from numpy.ctypeslib import ndpointer

import storm_analysis.sa_library.dao_fit_c as daoFitC
import storm_analysis.sa_library.loadclib as loadclib

import storm_analysis.multi_plane.mp_fit_c as mpFitC

def loadMPFitCArb():
    mp_fit = loadclib.loadCLibrary("storm_analysis.multi_plane", "mp_fit_arb")

    # From multi_plane/mp_fit_arb.c
    mp_fit.mpArbInitializePSFFFTChannel.argtypes = [ctypes.c_void_p,
                                                    ctypes.c_void_p,
                                                    ndpointer(dtype=numpy.float64),
                                                    ctypes.c_int]

    mp_fit.mpArbInitializePupilFnChannel.argtypes = [ctypes.c_void_p,
                                                     ctypes.c_void_p,
                                                     ndpointer(dtype=numpy.float64),
                                                     ctypes.c_double,
                                                     ctypes.c_double,
                                                     ctypes.c_int]

    mp_fit.mpArbInitializeSplineChannel.argtypes = [ctypes.c_void_p,
                                                    ctypes.c_void_p,
                                                    ndpointer(dtype=numpy.float64),
                                                    ctypes.c_int]

    mp_fit.mpArbNewPeaks.argtypes = [ctypes.c_void_p,
                                     ndpointer(dtype=numpy.float64),
                                     ctypes.c_char_p,
                                     ctypes.c_int]

    return mp_fit

class MPFitArb(mpFitC.MPFit):
    """
    Base class for multi-plane fitting with an arbitrary PSF.
    """
    def __init__(self, independent_heights = None, psf_objects = None, **kwds):
        super(MPFitArb, self).__init__(**kwds)

        self.independent_heights = independent_heights
        self.n_channels = len(psf_objects)
        self.psf_objects = psf_objects

        self.clib = loadMPFitCArb()
        mpFitC.addMPFitC(self.clib)

    def getSize(self):
        return self.psf_objects[0].getSize()

    def newPeaks(self, peaks, peaks_type):
        c_peaks = daoFitC.formatPeaksArbitraryPSF(peaks, peaks_type)
        self.clib.mpArbNewPeaks(self.mfit,
                                c_peaks,
                                ctypes.c_char_p(peaks_type.encode()),
                                c_peaks.shape[0])

class MPPSFFnFit(MPFitArb):
    """
    The basic idea is that we are going to use the functionality from PSF FFT
    to do most of the work. We will have one psfFFTFit C structure per image
    plane / channel.
    """
    def __init__(self, **kwds):
        super(MPPSFFnFit, self).__init__(**kwds)

        self.clamp = numpy.array([1.0,  # Height (Note: This is relative to the initial guess).
                                  1.0,  # x position
                                  0.3,  # width in x (Note: Not relevant for this fitter).
                                  1.0,  # y position
                                  0.3,  # width in y (Note: Not relevant for this fitter).
                                  1.0,  # background (Note: This is relative to the initial guess).
                                  0.5 * self.psf_objects[0].getZSize()]) # z position (in FFT size units).

    def rescaleZ(self, z):
        return self.psf_objects[0].rescaleZ(z)

    def setVariance(self, variance, channel):
        super(MPPSFFnFit, self).setVariance(variance, channel)

        # This where the differentation in which type of fitter to use happens.
        zmax = self.psf_objects[0].getZMax() * 1.0e-3
        zmin = self.psf_objects[0].getZMin() * 1.0e-3
        self.clib.mpArbInitializePSFFFTChannel(self.mfit,
                                               self.psf_objects[channel].getCPointer(),
                                               variance,
                                               channel)

    def setWeights(self, weights):
        if weights is None:
            weights = {"bg" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "h" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "x" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "y" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "z" : numpy.ones((2, self.n_channels))/float(self.n_channels)}
            super(MPSFFnFit, self).setWeights(weights, 0.0, 0.0)

        else:
            zmax = self.psf_objects[0].getZMax() * 1.0e-3
            zmin = self.psf_objects[0].getZMin() * 1.0e-3
            z_offset = -0.5*float(self.psf_objects[0].getZSize())
            z_scale = float(weights["bg"].shape[0]-1)/float(self.psf_objects[0].getZSize())
            super(MPPSFFnFit, self).setWeights(weights, z_offset, z_scale)

class MPPupilFnFit(MPFitArb):
    """
    The basic idea is that we are going to use the functionality from PupilFn
    to do most of the work. We will have one pupilFit C structure per image
    plane / channel.
    """
    def __init__(self, **kwds):
        super(MPPupilFnFit, self).__init__(**kwds)

        self.clamp = numpy.array([1.0,  # Height (Note: This is relative to the initial guess).
                                  1.0,  # x position
                                  0.3,  # width in x (Note: Not relevant for this fitter).
                                  1.0,  # y position
                                  0.3,  # width in y (Note: Not relevant for this fitter).
                                  1.0,  # background (Note: This is relative to the initial guess).
                                  1.0]) # z position

    def setVariance(self, variance, channel):
        super(MPPupilFnFit, self).setVariance(variance, channel)

        # This where the differentation in which type of fitter to use happens.
        zmax = self.psf_objects[0].getZMax() * 1.0e-3
        zmin = self.psf_objects[0].getZMin() * 1.0e-3
        self.clib.mpArbInitializePupilFnChannel(self.mfit,
                                                self.psf_objects[channel].getCPointer(),
                                                variance,
                                                zmin,
                                                zmax,
                                                channel)

    def setWeights(self, weights):
        if weights is None:
            weights = {"bg" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "h" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "x" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "y" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "z" : numpy.ones((2, self.n_channels))/float(self.n_channels)}
            super(MPPupilFnFit, self).setWeights(weights, 0.0, 0.0)

        else:
            zmax = self.psf_objects[0].getZMax() * 1.0e-3
            zmin = self.psf_objects[0].getZMin() * 1.0e-3
            z_offset = zmin
            z_scale = float(weights["bg"].shape[0]-1)/(zmax - zmin + 1.0e-12)
            super(MPPupilFnFit, self).setWeights(weights, z_offset, z_scale)

class MPSplineFit(MPFitArb):
    """
    The basic idea is that we are going to use the functionality from Spliner
    to do most of the work. We will have one splineFit C structure per image
    plane / channel.
    """
    def __init__(self, **kwds):
        super(MPSplineFit, self).__init__(**kwds)

        # Clamp parameters.
        #
        if True:
            self.clamp = numpy.array([1.0,  # Height (Note: This is relative to the initial guess).
                                      1.0,  # x position
                                      0.3,  # width in x (Note: Not relevant for this fitter).
                                      1.0,  # y position
                                      0.3,  # width in y (Note: Not relevant for this fitter).
                                      1.0,  # background (Note: This is relative to the initial guess).
                                      0.5 * self.psf_objects[0].getSplineSize()]) # z position (in spline size units).
        else:
            self.clamp = numpy.array([1.0,  # Height (Note: This is relative to the initial guess).
                                      1.0,  # x position
                                      0.3,  # width in x (Note: Not relevant for this fitter).
                                      1.0,  # y position
                                      0.3,  # width in y (Note: Not relevant for this fitter).
                                      1.0,  # background (Note: This is relative to the initial guess).
                                      1.0]) # z position (in spline size units).

    def rescaleZ(self, z):
        return self.psf_objects[0].rescaleZ(z)

    def setVariance(self, variance, channel):
        super(MPSplineFit, self).setVariance(variance, channel)

        # This where the differentation in which type of fitter to use happens.
        self.clib.mpArbInitializeSplineChannel(self.mfit,
                                               self.psf_objects[channel].getCPointer(),
                                               variance,
                                               channel)

    def setWeights(self, weights):
        if weights is None:
            weights = {"bg" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "h" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "x" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "y" : numpy.ones((2, self.n_channels))/float(self.n_channels),
                       "z" : numpy.ones((2, self.n_channels))/float(self.n_channels)}
            super(MPSplineFit, self).setWeights(weights, 0.0, 0.0)

        else:
            assert (weights["bg"].shape[0] == (self.psf_objects[0].getSplineSize() + 1)), "Incorrect shape for weights array."
            super(MPSplineFit, self).setWeights(weights, 0.0, 1.0)


#!/usr/bin/env python3

import os
import numpy as np
import yaml
import pickle
import argparse

def init_color_space(color_path):
    # type: (str) -> None
    """
    Initialization of color space from yaml or pickle.txt file

    :param str color_path: path to file containing the accepted colors
    :return: None
    """
    color_space = np.zeros((256, 256, 256), dtype=np.uint8)
    if color_path.endswith('.yaml'):
        with open(color_path, 'r') as stream:
            try:
                color_values = yaml.safe_load(stream)
            except yaml.YAMLError as exc:
                # TODO: what now??? Handle the error?
                pass
    # pickle-file is stored as '.txt'
    elif color_path.endswith('.txt'):
        try:
            with open(color_path, 'rb') as f:
                color_values = pickle.load(f)
        except pickle.PickleError as exc:
            pass

    # compatibility with colorpicker
    if 'color_values' in color_values.keys():
        color_values = color_values['color_values']['greenField']
    length = len(color_values['red'])
    if length == len(color_values['green']) and \
                    length == len(color_values['blue']):
        # setting colors from yaml file to True in color space
        for x in range(length):
            color_space[color_values['blue'][x],
                        color_values['green'][x],
                        color_values['red'][x]] = 1
    print("Imported color space")
    return color_space

def compare(positive_color_space, negative_color_space):
    mask = np.invert(np.array(negative_color_space, dtype=np.bool))
    binary_color_space = np.logical_and(mask, positive_color_space)
    return np.array(binary_color_space, dtype=np.uint8)

def generate_color_lists(color_space):
    color_space_positions = np.where(color_space == 1)
    color_lists = ( color_space_positions[0].tolist(),
                    color_space_positions[1].tolist(),
                    color_space_positions[2].tolist())
    return color_lists

def save(filename, color_space):
    red, green, blue = generate_color_lists(color_space)

    output_type = "negative_filtered"

    data = dict(
        red = red,
        green = green,
        blue = blue
    )

    filename = '{}_{}.txt'.format(filename, output_type)
    with open(filename, 'wb') as outfile:
        pickle.dump(data, outfile, protocol=2)
        # stores data of colorspace in file as pickle for efficient loading (yaml is too slow)

    print("Output saved to '{}'.".format(filename))

def run(positive_color_space_path, negative_color_space_path, output_path):
    print("Load positive color space '{}'".format(positive_color_space_path))
    positive_color_space = init_color_space(positive_color_space_path)
    print(np.count_nonzero(positive_color_space))
    print("Load negative color space '{}'".format(negative_color_space_path))
    negative_color_space = init_color_space(negative_color_space_path)
    print(np.count_nonzero(negative_color_space))
    print("Filter color spaces")
    filtered_color_space = compare(positive_color_space, negative_color_space)
    print(np.count_nonzero(filtered_color_space))
    print("Finished filtering")
    print("Save color space")
    save(output_path, filtered_color_space)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-p", "--positive", help="Color space where the negative color space is subtracted from.")
    parser.add_argument("-n", "--negative", help="Color space which is subtracted from the positive color space.")
    parser.add_argument("-o", "--output", help="Saves the output in a Pickle file.")
    args = parser.parse_args()
    np.warnings.filterwarnings('ignore')

    if args.positive and os.path.exists(args.positive):
        if args.negative and os.path.exists(args.negative):
            if args.output and os.path.exists(args.output):
                run(args.positive, args.negative, args.output)
            else:
                print("Output path incorrect!")
        else:
            print("Negative color space path incorrect!")
    else:
        print("Positive color space path incorrect!")


from __future__ import print_function
import numpy
import sys
import cpa.sqltools
from dbconnect import *
from properties import Properties
from datamodel import DataModel

db = DBConnect.getInstance()
p = Properties.getInstance()
dm = DataModel.getInstance()

temp_stump_table = "_stump"
temp_score_table = "_scores"
temp_class_table = "_class"
filter_table_prefix = '_filter_'

def translate(weaklearners):
    '''
    Translate weak leaners into a classifier() expression
    (or something more complicated in the future).
    '''
    num_features = len(weaklearners)
    nClasses = len(weaklearners[0][2])

    has_classifier_function = False
    if p.db_type.lower() == 'sqlite':
        thresholds = numpy.array([wl[1] for wl in weaklearners])
        a = numpy.array([wl[2] for wl in weaklearners])
        b = numpy.array([wl[3] for wl in weaklearners])
        db.setup_sqlite_classifier(thresholds, a, b)
        return "classifier(%s)"%(",".join([wl[0] for wl in weaklearners]))

    if p.db_type.lower() == 'mysql':
        # MySQL
        try:
            res = db.execute("SELECT * from mysql.func where name='classifier'")
            has_classifier_function = len(res) > 0
        except:
            has_classifier_function = False
        if has_classifier_function:
            num_stumps = len(weaklearners)
            featurenames = "1," + ",".join([wl[0] for wl in weaklearners])
            thresholds = "0,"+",".join([str(wl[1]) for wl in weaklearners])
            base = [sum([wl[3][k] for wl in weaklearners]) for k in range(nClasses)]
            weights = ",".join([",".join([str(base[k])] + [str(wl[2][k]-wl[3][k]) for wl in weaklearners]) for k in range(nClasses)])
            return "(classifier(%d, %s, %s, %s)+1)"%(num_stumps + 1, featurenames, thresholds, weights)
        else:
            class_scores = ['+'.join(['IF(`%s` > %f, %f, %f)'%(feature, threshold, a[i], b[i]) for (feature, threshold, a, b, ignore) in weaklearners]) for i in range(nClasses)]
            return "CASE GREATEST(%s) %s END"%(",".join(class_scores), "\n".join(["WHEN %s THEN %d"%(score, idx+1) for idx, score in enumerate(class_scores)]))

def FilterObjectsFromClassN(clNum, weaklearners, filterKeys):
    '''
    clNum: 1-based index of the class to retrieve obKeys from
    weaklearners: Weak learners from fastgentleboostingmulticlass.train
    filterKeys: (optional) A specific list of imKeys OR obKeys (NOT BOTH)
        to classify.
        * WARNING: If this list is too long, you may exceed the size limit to
          MySQL queries.
        * Useful when fetching N objects from a particular class. Use the
          DataModel to get batches of random objects, and sift through them
          here until N objects of the desired class have been accumulated.
        * Also useful for classifying a specific image or group of images.
    RETURNS: A list of object keys that fall in the specified class,
        if Properties.area_scoring_column is specified, area sums are also
        reported for each class
    '''

    class_query = translate(weaklearners)

    if filterKeys != []:
        if isinstance(filterKeys, str):
            whereclause = filterKeys + " AND"
        else:
            isImKey = len(filterKeys[0]) == len(image_key_columns())
            if isImKey:
                whereclause = GetWhereClauseForImages(filterKeys) + " AND"
            else:
                whereclause = GetWhereClauseForObjects(filterKeys) + " AND"
    else:
        whereclause = ""

    return db.execute('SELECT '+UniqueObjectClause()+' FROM %s WHERE %s %s=%d '%(p.object_table, whereclause, class_query, clNum))

def object_scores(weaklearners):
    stump_stmnts, score_stmnts, find_max_query, _, _ = \
                  translate(weaklearners)
    db.execute('DROP TABLE IF EXISTS _stump')
    db.execute('DROP TABLE IF EXISTS _scores')
    [db.execute(stump_query) for stump_query in stump_stmnts]
    [db.execute(score_query) for score_query in score_stmnts]
    col_names = db.GetColumnNames('_scores')
    col_types = db.GetColumnTypes('_scores')
    type_mapping = { long: 'i4', float: 'f8' }
    dtype = numpy.dtype([(name, type_mapping[type])
                         for name, type in zip(col_names, col_types)])
    res = db.execute('SELECT * from _scores')
    return numpy.array(map(tuple, res), dtype)

def create_perobject_class_table(classnames, rules):
    ''' Saves object keys and classes to a text file '''
    nClasses = len(classnames)

    if p.class_table is None:
        raise ValueError('"class_table" in properties file is not set.')

    index_cols = UniqueObjectClause()
    class_cols = UniqueObjectClause() + ', class, class_number'
    class_col_defs = object_key_defs() + ', class VARCHAR (%d)'%(max([len(c) for c in classnames])+1) + ', class_number INT'

    # Drop must be explicitly asked for Classifier.ScoreAll
    db.execute('DROP TABLE IF EXISTS %s'%(p.class_table))
    db.execute('CREATE TABLE %s (%s)'%(p.class_table, class_col_defs))
    db.execute('CREATE INDEX idx_%s ON %s (%s)'%(p.class_table, p.class_table, index_cols))

    case_expr = 'CASE %s'%(translate(rules)) + ''.join([" WHEN %d THEN '%s'"%(n+1, classnames[n]) for n in range(nClasses)]) + " END"
    case_expr2 = 'CASE %s'%(translate(rules)) + ''.join([" WHEN %d THEN '%s'"%(n+1, n+1) for n in range(nClasses)]) + " END"
    db.execute('INSERT INTO %s (%s) SELECT %s, %s, %s FROM %s'%(p.class_table, class_cols, index_cols, case_expr, case_expr2, p.object_table))
    db.Commit()

def _objectify(p, field):
    return "%s.%s"%(p.object_table, field)

def _where_clauses(p, dm, filter_name):
    imkeys = dm.GetAllImageKeys(filter_name)
    imkeys.sort()
    stepsize = max(len(imkeys) / 100, 50)
    key_thresholds = imkeys[-1:1:-stepsize]
    key_thresholds.reverse()
    if len(key_thresholds) == 0:
        return ['(1 = 1)']
    if p.table_id:
        # split each table independently
        def splitter():
            yield "(%s = %d) AND (%s <= %d)"%(_objectify(p, p.table_id), key_thresholds[0][0],
                                              _objectify(p, p.image_id), key_thresholds[0][1])
            for lo, hi in zip(key_thresholds[:-1], key_thresholds[1:]):
                if lo[0] == hi[0]:
                    # block within one table
                    yield "(%s = %d) AND (%s > %d) AND (%s <= %d)"%(_objectify(p, p.table_id), lo[0],
                                                                    _objectify(p, p.image_id), lo[1],
                                                                    _objectify(p, p.image_id), hi[1])
                else:
                    # query spans a table boundary
                    yield "(%s >= %d) AND (%s > %d)"%(_objectify(p, p.table_id), lo[0],
                                                     _objectify(p, p.image_id), lo[1])
                    yield "(%s <= %d) AND (%s <= %d)"%(_objectify(p, p.table_id), hi[0],
                                                      _objectify(p, p.image_id), hi[1])
        return list(splitter())
    else:
        return (["(%s <= %d)"%(_objectify(p, p.image_id), key_thresholds[0][0])] +
                ["(%s > %d) AND (%s <= %d)"
                 %(_objectify(p, p.image_id), lo[0], _objectify(p, p.image_id), hi[0])
                 for lo, hi in zip(key_thresholds[:-1], key_thresholds[1:])])

def PerImageCounts(weaklearners, filter_name=None, cb=None):
    '''
    weaklearners: Weak learners from fastgentleboostingmulticlass.train
    filter: name of filter, or None.
    cb: callback function to update with the fraction complete
    RETURNS: A list of lists of imKeys and respective object counts for each class:
        Note that the imKeys are exploded so each row is of the form:
        [TableNumber, ImageNumber, Class1_ObjectCount, Class2_ObjectCount,...]
        where TableNumber is only present if table_id is defined in Properties.
        If p.area_scoring_column is set, then area scores will be appended to
        the object scores.
    '''

    # I'm pretty sure this would be even faster if we were to run two
    # or more parallel threads and split the work between them.
    def do_by_steps(class_query, tables, filter_name, result_clauses):
        filter_clause = '1 = 1'
        join_clause = ''
        if filter_name is not None:
            filter = p._filters[filter_name]
            if isinstance(filter, cpa.sqltools.OldFilter):
                join_table = '(%s) as filter' % str(filter)
            else:
                if p.object_table in tables:
                    join_table = None
                else:
                    join_table = p.object_table
                    filter_clause = str(filter)
            if join_table:
                join_clause = 'JOIN %s USING (%s)' % (join_table, ','.join(image_key_columns()))
        if cb:
            result =  []
            wheres = _where_clauses(p, dm, filter_name)
            num_clauses = len(wheres)

            for idx, where_clause in enumerate(wheres):
                if filter_clause is not None:
                    where_clause += ' AND ' + filter_clause
                result += [db.execute('SELECT %s, %s as class, %s FROM %s '
                                      '%s WHERE %s GROUP BY %s, class'
                                      %(UniqueImageClause(p.object_table),
                                        class_query, result_clauses, tables,
                                        join_clause, where_clause,
                                        UniqueImageClause(p.object_table)),
                                      silent=(idx > 10))]
                cb(min(1, idx/float(num_clauses)))
            return sum(result, [])
        else:
            return db.execute('SELECT %s, %s as class, %s FROM %s %s WHERE %s GROUP BY %s, class'%
                              (imkeys, class_query, result_clauses, tables, join_clause, filter_clause, imkeys))

    if p.area_scoring_column is None:
        result_clauses = 'COUNT(*)'
    else:
        result_clauses = 'COUNT(*), SUM(%s)'%(_objectify(p, p.area_scoring_column))

    class_query = translate(weaklearners)
    results = do_by_steps(class_query, p.object_table, filter_name, result_clauses)

    # convert to dictionary
    counts = {}
    keylen = 1 + len(image_key_columns()) # includes class

    for r in results:
        counts[r[:keylen]] = r[keylen:]

    num_classes = len(weaklearners[0][2])
    # this is clearer than the one-line version
    def get_count(im_key, classnum):
        return counts.get(tuple(list(im_key) + [classnum]), [0])[0]

    def get_area(im_key, classnum):
        return counts.get(tuple(list(im_key) + [classnum]), [0, 0])[1]

    def get_results():
        for imkey in dm.GetImageKeysAndObjectCounts(filter_name):
            if p.area_scoring_column is None:
                yield list(imkey[0]) + [get_count(imkey[0], cl) for cl in range(1, num_classes+1)]
            else:
                yield list(imkey[0]) + [get_count(imkey[0], cl) for cl in range(1, num_classes+1)] + [get_area(imkey[0], cl) for cl in range(1, num_classes+1)]

    return list(get_results())

if __name__ == "__main__":
#    dir = "/Users/ljosa/research/modifier/piyush"
#    p.LoadFile("%s/batch1and2.properties"%(dir,))
#    import cPickle as pickle
#    f = open("%s/batch1and2.boostingclassifier"%(dir,))
#    learners = pickle.load(f)
#    f.close()
#
#    keys_and_counts = PerImageCounts(learners, filter='HRG')
#    import numpy
#    keys_and_counts = numpy.array(keys_and_counts, dtype='i4')
#    print keys_and_counts[0,:]
#    print sum(keys_and_counts[:,-2:])

    from trainingset import TrainingSet
    from StringIO import StringIO
    import fastgentleboostingmulticlass
    from datatable import DataGrid
    import wx
    p = Properties.getInstance()
    db = DBConnect.getInstance()
    dm = DataModel.getInstance()

#    props = '/Volumes/imaging_analysis/2007_10_19_Gilliland_LeukemiaScreens/Screen3_1Apr09_run3/2007_10_19_Gilliland_LeukemiaScreens_Validation_v2_AllBatches_DuplicatesFiltered_FullBarcode_testSinglePlate.properties'
#    ts = '/Volumes/imaging_analysis/2007_10_19_Gilliland_LeukemiaScreens/Screen3_1Apr09_run3/trainingvalidation3b.txt'
    props = '../Properties/nirht_area_test.properties'
    ts = '/Users/afraser/Desktop/MyTrainingSet3.txt'
    nRules = 5
    filter = 'MAPs'
#    props = '/Users/afraser/Desktop/2007_10_19_Gilliland_LeukemiaScreens_Validation_v2_AllBatches_DuplicatesFiltered_FullBarcode.properties'
#    ts = '/Users/afraser/Desktop/trainingvalidation3d.txt'
#    nRules = 50
#    filter = 'afraser_test'

    p.LoadFile(props)
    trainingSet = TrainingSet(p)
    trainingSet.Load(ts)
    output = StringIO()
    print('Training classifier with '+str(nRules)+' rules...')
    weaklearners = fastgentleboostingmulticlass.train(trainingSet.colnames,
                                                      nRules, trainingSet.label_matrix,
                                                      trainingSet.values, output)
    table = PerImageCounts(weaklearners, filter_name=filter)
    table.sort()

    labels = ['table', 'image'] + list(trainingSet.labels) + list(trainingSet.labels)
    print(labels)
    for row in table:
        print(row)
#    app = wx.PySimpleApp()
#    grid = DataGrid(numpy.array(table), labels, key_col_indices=[0,1])
#    grid.Show()
#    app.MainLoop()


from __future__ import print_function
import argparse
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
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
from torch.autograd import Variable
import torch.nn.functional as F
import numpy as np
import time
import math
from sintel_io import depth_read
from dataloader import sintel_listflowfile as lt
# from dataloader import sintel_listflowfile_with_filter as lt
# from dataloader import sintel_listflowfile_without_filter_with_depth as lt
from dataloader import SintelFlowLoader as DA
from models import *
import pickle
from dfd import Dfd_net, psi_to_depth
from edof import EdofNet
# from fuse_net import Fuse
# from fuse_net import UNet
from fuse_net import MaskFuse
from fuse_net import ConfFuse
from local_utils import load_model, depth2disparity
import matplotlib.pyplot as plt
from disparity_mapping import apply_disparity

parser = argparse.ArgumentParser(description='PSMNet')
parser.add_argument('--maxdisp', type=int, default=192,
                    help='maxium disparity')
parser.add_argument('--model', default='stackhourglass',
                    help='select model')
parser.add_argument('--datapath', default='/media/yotamg/bd0eccc9-4cd5-414c-b764-c5a7890f9785/Yotam/Stereo/',
                    help='datapath')
parser.add_argument('--left_imgs', default=None,
                    help='left img train dir name')
parser.add_argument('--right_imgs', default=None,
                    help='left img train dir name')
parser.add_argument('--disp_imgs', default=None,
                    help='left img train dir name')
parser.add_argument('--disp_R_imgs', default=None,
                    help='right img train dir name')
parser.add_argument('--epochs', type=int, default=100,
                    help='number of epochs to train')
parser.add_argument('--focus_L', default='1500',
                    help='Left img focus point')
parser.add_argument('--focus_R', default='700',
                    help='Right img focus point')
# parser.add_argument('--loadmodel', default=None,
parser.add_argument('--loadmodel', default='./checkpoints/checkpoint_50.tar',
# parser.add_argument('--loadmodel', default='./checkpoints/checkpoint_filter_loss_2.6.tar',
# parser.add_argument('--loadmodel', default='./checkpoints/checkpoint_clean_from_scratch_loss_2.1.tar',
# parser.add_argument('--loadmodel', default='./pretrained_model_KITTI2015.tar',
                    help='load model')
parser.add_argument('--loadfuse', default=None,
                    help='load fuse model')
parser.add_argument('--savemodel', default='./checkpoints/',
                    help='save model')
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
parser.add_argument('--clean', action='store_true', default=False,
                    help='random seed (default: 1)')
parser.add_argument('--cont', action='store_true', default=False,
                    help='random seed (default: 1)')
parser.add_argument('--dfd', action='store_true', default=False,
                    help='include dfd net')
parser.add_argument('--dfd_at_end', action='store_true', default=False,
                    help='include dfd net')
parser.add_argument('--right_head', action='store_true', default=False,
                    help='right disp branch')
parser.add_argument('--pred_occlusion', action='store_true', default=False,
                    help='pred occlusion or right depth')

args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)
if args.cuda:
    torch.cuda.manual_seed(args.seed)
# test_imgs  = ['City_R_0116_1100_maskImg.png', 'City_R_0092_1100_maskImg.png', 'City_R_0114_1100_maskImg.png', 'City_R_0190_1100_maskImg.png', 'City_R_0154_1100_maskImg.png', 'City_R_0066_1100_maskImg.png', 'City_R_0202_1100_maskImg.png', 'City_R_0204_1100_maskImg.png', 'City_R_0128_1100_maskImg.png', 'City_R_0042_1100_maskImg.png', 'City_R_0026_1100_maskImg.png', 'City_R_0058_1100_maskImg.png']

# if args.clean:
#     test_imgs = [img.replace('_1100_maskImg.png', '.tif') for img in test_imgs]

# all_left_img, all_right_img, all_left_disp, test_left_img, test_right_img, test_left_disp = lt.dataloader(args.datapath, args.left_imgs, args.right_imgs, args.disp_imgs,filenames=test_imgs, clean=args.clean)
all_left_img, all_right_img, all_left_disp, all_R_disp, test_left_img, test_right_img, test_left_disp, test_R_disp = lt.dataloader(args.datapath, args.left_imgs, args.right_imgs, args.disp_imgs, args.disp_R_imgs, clean=args.clean, focus_L=args.focus_L, focus_R=args.focus_R)

TrainImgLoader = torch.utils.data.DataLoader(
    # DA.myImageFloder(all_left_img, all_right_img, all_left_disp, True,dploader=DA.depth_loader),
    DA.myImageFloder(all_left_img, all_right_img, all_left_disp, all_R_disp, training=True, dploader=depth_read,cont=args.cont),
    batch_size=4, shuffle=True, num_workers=8, drop_last=False)

TestImgLoader = torch.utils.data.DataLoader(
    # DA.myImageFloder(test_left_img, test_right_img, test_left_disp, False, dploader=DA.depth_loader),
    DA.myImageFloder(test_left_img, test_right_img, test_left_disp, test_R_disp, training=False, dploader=depth_read, cont=args.cont),
    batch_size=2, shuffle=False, num_workers=4, drop_last=False)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if args.model == 'stackhourglass':
    model = stackhourglass(args.maxdisp, device=device, dfd_net=args.dfd, dfd_at_end=args.dfd_at_end, right_head=args.right_head)
elif args.model == 'basic':
    model = basic(args.maxdisp)
else:
    print('no model')

if args.cuda:
    model = nn.DataParallel(model)
    model.cuda()

if args.loadmodel is not None:
    state_dict = torch.load(args.loadmodel)
    model.load_state_dict(state_dict['state_dict'], strict=False)
    print("Loading stereo net checkpoint: ", args.loadmodel)

print('Number of model parameters: {}'.format(sum([p.data.nelement() for p in model.parameters()])))

dfd_net_D5 = Dfd_net(mode='segmentation', target_mode='cont', pool=False)
dfd_net_D5 = dfd_net_D5.eval()
dfd_net_D5 = dfd_net_D5.to(device)
load_model(dfd_net_D5, device, model_path='/home/yotamg/PycharmProjects/dfd/trained_models/Net_continuous_dn1500_D5/checkpoint_254.pth.tar')

dfd_net = Dfd_net(mode='segmentation', target_mode='cont', pool=False)
dfd_net = dfd_net.eval()
dfd_net = dfd_net.to(device)
load_model(dfd_net, device, model_path='/home/yotamg/PycharmProjects/dfd/trained_models/Net_continuous_dn1500/checkpoint_257.pth.tar')

# edof_net = EdofNet(max_dilation=4, device=device)
# edof_net = edof_net.eval()
# edof_net = edof_net.to(device)
# load_model(edof_net,device, model_path='/home/yotamg/PycharmProjects/EDOF/trained_models/EdofNet_after_imaging_fix/checkpoint_99.pth.tar')

# fuse_net = Fuse()
# fuse_net = UNet()
# fuse_net = MaskFuse(256,512)
fuse_net = ConfFuse()
fuse_net = fuse_net.train()
fuse_net = fuse_net.to(device)

if args.loadfuse is not None:
    state_dict = torch.load(args.loadfuse)
    fuse_net.load_state_dict(state_dict['state_dict'])
    print ("Loading fuse net checkpoint: ", args.loadfuse)

optimizer = optim.Adam(fuse_net.parameters(), lr=0.1, betas=(0.9, 0.999))

def train(imgL, imgR, disp_true, disp_true_R):
    optimizer.zero_grad()
    model.eval()
    fuse_net.train()
    imgL = Variable(torch.FloatTensor(imgL))
    imgR = Variable(torch.FloatTensor(imgR))

    if args.cuda:
        imgL, imgR,disp_true = imgL.cuda(), imgR.cuda(), disp_true.cuda()

    # ---------
    mask = disp_true < 192
    # ----

    with torch.no_grad():
        mono_L, mono_L_conf = dfd_net_D5(imgL, float(args.focus_L)*1e-3, D=(2.28*1e-3))
        mono_R, mono_R_conf = dfd_net(imgR, float(args.focus_R)*1e-3, D=(5*1e-3))
        if len(mono_L.shape) == 2:
            mono_L = torch.unsqueeze(mono_L,0).to(device)
            mono_R = torch.unsqueeze(mono_R,0).to(device)
        # plt.subplot(211)
        # plt.imshow(((imgR[0] + 1)/2).permute(1,2,0))
        # plt.subplot(212)
        # plt.imshow(mono_R[0])
        # plt.show()
        if args.right_head:
            if args.pred_occlusion:
                output3, occ, stereo_conf = model(imgL, imgR)
            else:
                output3, output_R, stereo_conf = model(imgL, imgR)
        else:
            output3, stereo_conf = model(imgL, imgR)

    disp_from_stereo = depth2disparity(imgR, output3,device)

    mono_R = torch.unsqueeze(mono_R,1)
    mono_RtoL = apply_disparity(mono_R, -disp_from_stereo)
    mono_RtoL = torch.squeeze(mono_RtoL, 1)
    # mono_RtoL = mono_RtoL * torch.abs(occ-1)

    mono_RtoL_conf = apply_disparity(mono_R_conf, -disp_from_stereo)
    # occ = torch.unsqueeze(occ, 1).repeat(1,16,1,1)
    # mono_RtoL_conf = mono_RtoL_conf * torch.abs(occ-1)

    # plt.subplot(221)
    # plt.imshow(mono_RtoL[0],vmin=torch.min(mono_R), vmax= torch.max(mono_R))
    # plt.subplot(222)
    # plt.imshow(mono_R[0][0],vmin=torch.min(mono_R), vmax= torch.max(mono_R))
    # plt.subplot(223)
    # plt.imshow((imgL[0].permute(1,2,0) + 1) / 2)
    # plt.subplot(224)
    # plt.imshow((imgR[0].permute(1,2,0) + 1) / 2)
    # plt.subplot(223)
    # plt.imshow(output3[0])
    # plt.subplot(224)
    # plt.imshow(disp_from_stereo[0])
    plt.show()

    # fuse_output = fuse_net(fuse_input)
    fuse_output = fuse_net(mono_L, mono_RtoL, output3, mono_L_conf, mono_RtoL_conf, stereo_conf)
    fuse_output = torch.squeeze(fuse_output, dim=1)
    # fuse_output = torch.squeeze(fuse_output, dim=3)
    loss = F.smooth_l1_loss(fuse_output, disp_true, size_average=True)
    loss.backward()
    optimizer.step()
    return loss.data

def test(imgL, imgR, disp_true):
    model.eval()
    dfd_net.eval()
    fuse_net.eval()
    imgL = Variable(torch.FloatTensor(imgL))
    imgR = Variable(torch.FloatTensor(imgR))

    if args.cuda:
        imgL, imgR, disp_true = imgL.cuda(), imgR.cuda(), disp_true.cuda()

    with torch.no_grad():
        mono_L, mono_L_conf = dfd_net(imgL, float(args.focus_L) * 1e-3, D=(5.0 * 1e-3))
        mono_R, mono_R_conf = dfd_net(imgR, float(args.focus_R) * 1e-3, D=(2.28 * 1e-3))
        if len(mono_L.shape) == 2:
            mono_L = torch.unsqueeze(mono_L,0)
            mono_R = torch.unsqueeze(mono_R,0)
        output3, occ, stereo_conf = model(imgL, imgR)

        disp_from_stereo = depth2disparity(imgR, output3,device)

        mono_R = torch.unsqueeze(mono_R, 1)
        mono_R = apply_disparity(mono_R, -disp_from_stereo)
        mono_R = torch.squeeze(mono_R, 1)
        # mono_R = mono_R * torch.abs(occ - 1)

        mono_R_conf = apply_disparity(mono_R_conf, -disp_from_stereo)
        occ = torch.unsqueeze(occ, 1).repeat(1, 16, 1, 1)
        # mono_R_conf = mono_R_conf * torch.abs(occ - 1)

        fuse_output = fuse_net(mono_L, mono_R, output3, mono_L_conf, mono_R_conf, stereo_conf)
        # try:
        #     fuse_output = torch.squeeze(fuse_output, dim=3)
        # except:
        #     fuse_output = fuse_output.permute(2,0,1)
    loss = torch.mean(torch.abs(fuse_output - disp_true))  # end-point-error
    rel_loss = torch.mean(torch.abs(fuse_output - disp_true) / disp_true)  # end-point-error

    mask_psi_range = (disp_true > 1.018) & (disp_true < 1.5)
    mask_psi1500_range_from_stereo = (output3 > 1.63) & (output3 < 1.83)
    # mask_psi1500_range_from_stereo = (disp_true > 1.34) & (disp_true < 1.59)
    mask_psi700_range_from_stereo = (output3 > 0.393) & (output3 < 1.018)
    # mask_psi700_range_from_stereo = (disp_true > 0.393) & (disp_true < 1.018)
    output_threshold_fuse = output3.clone()
    output_threshold_fuse[mask_psi1500_range_from_stereo] = mono_R[mask_psi1500_range_from_stereo]
    output_threshold_fuse[mask_psi700_range_from_stereo]  = mono_L[mask_psi700_range_from_stereo]

    mono_rel_loss = torch.mean(torch.abs(mono_L - disp_true) / disp_true)  # end-point-error
    mono_masked_rel_loss = 0
    mono700_masked_rel_loss = 0
    mono1500_masked_rel_loss = 0
    if len(mono_L[mask_psi_range]) != 0:
        mono_masked_rel_loss = torch.mean(torch.abs(mono_L[mask_psi_range] - disp_true[mask_psi_range]) / disp_true[mask_psi_range])  # end-point-error
    if len(mono_R[mask_psi700_range_from_stereo]) != 0:
        mono700_masked_rel_loss = torch.mean(torch.abs(mono_R[mask_psi700_range_from_stereo] - output3[mask_psi700_range_from_stereo]) / output3[mask_psi700_range_from_stereo])  # end-point-error
    if len(mono_L[mask_psi1500_range_from_stereo]) != 0:
        mono1500_masked_rel_loss = torch.mean(torch.abs(mono_L[mask_psi1500_range_from_stereo] - output3[mask_psi1500_range_from_stereo]) / output3[mask_psi1500_range_from_stereo])  # end-point-error
    stereo_rel_loss = torch.mean(torch.abs(output3 - disp_true) / disp_true)  # end-point-error
    threshold_fuse_rel_loss = torch.mean(torch.abs(output_threshold_fuse - disp_true) / disp_true)
    return loss, rel_loss, mono_rel_loss,mono_masked_rel_loss, stereo_rel_loss, threshold_fuse_rel_loss, mono700_masked_rel_loss, mono1500_masked_rel_loss

def adjust_learning_rate(optimizer, epoch):
    lr = 0.001*(1/epoch)
    print(lr)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

def main():
    start_full_time = time.time()

    # ------------- TRAIN ------------------------------------------------------------
    total_train_loss = 0
    epochs = 300
    do_train = True
    for epoch in range(1,epochs):
        if do_train:
            for batch_idx, (imgL, imgR, disp_L, disp_R) in enumerate(TrainImgLoader):
                loss = train(imgL, imgR, disp_L, disp_R)
                print('Iter %d train loss = %.3f' % (batch_idx, loss))
                total_train_loss += loss

            if not os.path.isdir(args.savemodel):
                os.makedirs(args.savemodel)
            savefilename = args.savemodel + '/checkpoint_' + str(epoch) + '.tar'
            torch.save({
                'epoch': epoch,
                'state_dict': fuse_net.state_dict(),
                'train_loss': total_train_loss / len(TrainImgLoader),
            }, savefilename)
            print('total train loss = %.3f' % (total_train_loss / len(TrainImgLoader)))
            total_train_loss  = 0
        # ------------- TEST ------------------------------------------------------------
        if epoch % 1 == 0 or not do_train:
            total_test_loss = 0
            total_test_rel_loss = 0
            total_test_mono_rel_loss = 0
            total_test_mono_masked_rel_loss = 0
            total_test_mono_700_masked_rel_loss = 0
            total_test_mono_700_masked_rel_cnt = 0
            total_test_mono_1500_masked_rel_loss = 0
            total_test_mono_1500_masked_rel_cnt = 0
            total_test_stereo_rel_loss = 0
            total_test_mono_masked_rel_cnt = 0
            total_test_threshold_rel_loss = 0
            for batch_idx, (imgL, imgR, disp_L, disp_R) in enumerate(TestImgLoader):
                test_loss, rel_loss, mono_rel_loss, mono_masked_rel_loss, stereo_rel_loss, threshold_fuse_rel_loss, mono700_masked_rel_loss, mono1500_masked_rel_loss = test(imgL, imgR, disp_L)
                print('Iter %d test loss = %.3f' % (batch_idx, test_loss))
                print('Iter %d test rel loss = %.3f' % (batch_idx, rel_loss))
                print('Iter %d test mono rel loss = %.3f' % (batch_idx, mono_rel_loss))
                print('Iter %d test mono masked rel loss = %.3f' % (batch_idx, mono_masked_rel_loss))
                print('Iter %d test mono 700 masked rel loss = %.3f' % (batch_idx, mono700_masked_rel_loss))
                print('Iter %d test mono 1500 masked rel loss = %.3f' % (batch_idx, mono1500_masked_rel_loss))
                print('Iter %d test stereo rel loss = %.3f' % (batch_idx, stereo_rel_loss))
                print('Iter %d test threshold rel loss = %.3f' % (batch_idx, threshold_fuse_rel_loss))
                # plt.show()
                total_test_loss += test_loss
                total_test_rel_loss += rel_loss
                total_test_mono_rel_loss += mono_rel_loss
                if mono_masked_rel_loss != 0:
                    total_test_mono_masked_rel_loss += mono_masked_rel_loss
                    total_test_mono_masked_rel_cnt += 1
                if mono700_masked_rel_loss != 0:
                    total_test_mono_700_masked_rel_loss += mono700_masked_rel_loss
                    total_test_mono_700_masked_rel_cnt += 1
                if mono1500_masked_rel_loss != 0:
                    total_test_mono_1500_masked_rel_loss += mono1500_masked_rel_loss
                    total_test_mono_1500_masked_rel_cnt += 1
                total_test_stereo_rel_loss += stereo_rel_loss
                total_test_threshold_rel_loss += threshold_fuse_rel_loss

            print('total test loss = %.3f' % (total_test_loss / len(TestImgLoader)))
            print('total test rel loss = %.3f' % (total_test_rel_loss / len(TestImgLoader)))
            print('total test mono rel loss = %.3f' % (total_test_mono_rel_loss / len(TestImgLoader)))
            print('total test mono masked rel loss = %.3f' % (total_test_mono_masked_rel_loss / total_test_mono_masked_rel_cnt))
            print('total test mono 700 masked rel loss = %.3f' % (total_test_mono_700_masked_rel_loss / total_test_mono_700_masked_rel_cnt))
            print('total test mono 1500 masked rel loss = %.3f' % (total_test_mono_1500_masked_rel_loss / total_test_mono_1500_masked_rel_cnt))
            print('total test stereo rel loss = %.3f' % (total_test_stereo_rel_loss / len(TestImgLoader)))
            print('total test threshold fuse rel loss = %.3f' % (total_test_threshold_rel_loss / len(TestImgLoader)))

if __name__ == '__main__':
    main()


import numpy as np
import cv2
import random
import time
import torch
import copy
import math
from torch.utils.data.dataset import Dataset
from utils.vis import vis_keypoints, vis_3d_skeleton
from utils.utils_pose import fliplr_joints, transform_joint_to_other_db
# from config import cfg
from math import floor
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from utils import utils_pose
from utils import utils_tool as ut_t

candidate_sets = ['Human36M', 'ScanAva', 'MPII', 'MSCOCO', 'MuCo', 'MuPoTS', 'SURREAL']

for i in range(len(candidate_sets)):
	exec('from data.' + candidate_sets[i] + '.' + candidate_sets[i] + ' import ' + candidate_sets[i])

class AdpDataset_3d(Dataset):
	'''
	adaptation datasest to make uniform interface for 3D pose. Take in all data as db original, augment, then map to original,  normalize to HM, to tensor
	changed to dict based feed e
	'''

	def __init__(self, db, ref_joints_name, is_train, transform=None, opts={}):
		self.db = db.data       # get data inherer
		self.joint_num = db.joint_num
		self.skeleton = db.skeleton
		self.flip_pairs = db.flip_pairs
		self.joints_have_depth = db.joints_have_depth
		self.if_SYN = db.if_SYN  # ds determined
		self.joints_name = db.joints_name
		self.ref_joints_name = ref_joints_name
		self.joint_num_ref = len(ref_joints_name)   # should be 17  for wht gen
		if not transform:  # if not given, give tensor normalization to overide
			transform = transforms.Compose([
				transforms.ToTensor(),
				transforms.Normalize(mean=opts.pixel_mean, std=opts.pixel_std)])    # maybe image needs to be normalized

		self.transform = transform
		self.is_train = is_train
		self.opts = opts

		if self.is_train:
			self.do_augment = True
		else:
			self.do_augment = False

	def __getitem__(self, index):
		joint_num = self.joint_num      # the ds jt number instead of output
		skeleton = self.skeleton
		flip_pairs = self.flip_pairs
		joints_have_depth = self.joints_have_depth
		if_SYN = self.if_SYN
		data = copy.deepcopy(self.db[index])
		bbox = data['bbox']  # give some default one for this in case not in there
		joint_img = data['joint_img']
		joint_vis = data['joint_vis']
		joint_num_ref = self.joint_num_ref
		# if 'boneLen' in data.keys():
		# 	boneLen = data['boneLen']   # bone length
		# else:
		# 	boneLen = 4.5       # no bone len, 2d doesn't affect at all
		# 1. load image
		cvimg = cv2.imread(data['img_path'], cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION)
		if not isinstance(cvimg, np.ndarray):
			raise IOError("Fail to read %s" % data['img_path'])
		img_height, img_width, img_channels = cvimg.shape

		# 2. get augmentation params
		if self.do_augment:
			scale, rot, do_flip, color_scale, do_occlusion = get_aug_config()
		else:
			scale, rot, do_flip, color_scale, do_occlusion = 1.0, 0.0, False, [1.0, 1.0, 1.0], False

		# 3. crop patch from img and perform data augmentation (flip, rot, color scale, synthetic occlusion)
		img_patch, trans = generate_patch_image(cvimg, bbox, do_flip, scale, rot, do_occlusion,
		                                        input_shape=self.opts.input_shape)  # rgb, 255 still
		for i in range(img_channels):
			img_patch[:, :, i] = np.clip(img_patch[:, :, i] * color_scale[i], 0, 255)

		# 4. generate patch joint ground truth
		# flip joints and apply Affine Transform on joints
		if do_flip:
			joint_img[:, 0] = img_width - joint_img[:, 0] - 1
			for pair in flip_pairs:
				joint_img[pair[0], :], joint_img[pair[1], :] = joint_img[pair[1], :], joint_img[pair[0], :].copy()
				joint_vis[pair[0], :], joint_vis[pair[1], :] = joint_vis[pair[1], :], joint_vis[pair[0], :].copy()

		for i in range(len(joint_img)):  # 2d first for boneLen calculation
			joint_img[i, 0:2] = trans_point2d(joint_img[i, 0:2], trans)  # to pix:patch under input_shape size
		if 'joint_cam' in data.keys() and 'y' == self.opts.if_normBone:
			joint_2d = joint_img[:, 0:2]  # x,y pix:patch
			joint_cam = data['joint_cam']
			# boneLen3d = data['boneLen3d']
			boneLen2d_pix = utils_pose.get_boneLen(joint_2d, skeleton)  # rescaled bone
			boneLen2d_mm = utils_pose.get_boneLen(joint_cam[:, :2], skeleton)
			if not boneLen2d_mm>0:
				boneLen2d_mm = 3700
			rt_pix2mm = boneLen2d_pix / boneLen2d_mm

			joint_img[:, 2] *= rt_pix2mm / self.opts.input_shape[
				0]  # z = z*2d/3d /inp_size/  normalize to image size. Note, people usually use square space, yet for tight BB, sometimes it is possible depth is larger then 2d space. so I keep sensing range x:y:z = 1:1:2 ratio to include distant z. Estimation use standard z
		else:  # fix box norm
			joint_img[:, 2] /= self.opts.bbox_3d_shape[0] / 2.  # to -1,1 z total 2m

		for i in range(len(joint_img)):
			joint_img[i, 2] = (joint_img[i, 2] + 1.0) / 2.  # 0~1 normalize
			joint_vis[i] *= (       # & bitwise operation  , parenthesis ofr priority over &
					(joint_img[i, 0] >= 0) & \
					(joint_img[i, 0] < self.opts.input_shape[1]) & \
					(joint_img[i, 1] >= 0) & \
					(joint_img[i, 1] < self.opts.input_shape[0]) & \
					(joint_img[i, 2] >= 0) & \
					(joint_img[i, 2] < 1)
			)  # nice filtering  all in range

		vis = False
		if vis:
			filename = str(random.randrange(1, 500))
			tmpimg = img_patch.copy().astype(np.uint8)      # rgb
			tmpkps = np.zeros((3, joint_num))
			tmpkps[:2, :] = joint_img[:, :2].transpose(1, 0)
			tmpkps[2, :] = joint_vis[:, 0]
			tmpimg = vis_keypoints(tmpimg, tmpkps, skeleton)    # rgb
			cv2.imwrite(filename + '_gt.jpg', tmpimg)

		vis = False
		if vis:
			vis_3d_skeleton(joint_img, joint_vis, skeleton, filename)
		joint_img_ori = joint_img.copy()        # raw image
		# change coordinates to output space the hmap space (usually 64 or 32 )
		joint_img[:, 0] = joint_img[:, 0] / self.opts.input_shape[1] * self.opts.output_shape[1]    # 64 x 64
		joint_img[:, 1] = joint_img[:, 1] / self.opts.input_shape[0] * self.opts.output_shape[0]
		joint_img[:, 2] = joint_img[:, 2] * self.opts.depth_dim  # 0 ~ depth_dim, from 1 to 64

		# if self.is_train: # ds completeness assumption, can feed all always
		img_patch = self.transform(img_patch)

		if self.ref_joints_name is not None:
			joint_img = transform_joint_to_other_db(joint_img, self.joints_name, self.ref_joints_name)
			joint_vis = transform_joint_to_other_db(joint_vis, self.joints_name,
			                                        self.ref_joints_name)  # invisible to no exist one

		joint_img = joint_img.astype(np.float32)  # x,y,z, hm
		joint_vis = (joint_vis > 0).astype(np.float32)
		joints_have_depth = np.array([joints_have_depth]).astype(np.float32)  # into 1 D array
		if_SYN = np.array([if_SYN]).astype(np.float32)  # bool -> float32,   to a 1D vec, to concate to N x1 (1 class)

		# make gauss_hm # n_stg:4 x n_jt:17 [sz * sz]
		# shp_t = self.opts.output_shape

		# add the evoSkel ds
		joint_evoSkel = np.zeros([16,2])  # the std joint evo , 0 to 8 10 ~ 16  16 joints
		joint_evoSkel[:9] = joint_img[:9, :2]
		joint_evoSkel[9:] = joint_img[10:, :2]
		joint_evoSkel = (joint_evoSkel - joint_evoSkel.mean(axis=0))/joint_evoSkel.std(axis=0)
		joint_evoSkel = joint_evoSkel.flatten()

		wht = np.zeros([joint_num_ref, 8, 8])   # n_jt x 8 x 8  hardwired features space
		# print(joint_img)
		for i, coord in enumerate(joint_img):     # hm all empty?
			coord = np.round(coord / 8).astype(int)
			# print(coord)
			# rst_andS = coord[0] >= 0 & coord[0] < 8 & coord[1] >= 0 & coord[1] < 8  # always false
			# print(rst_andS)
			if coord[0]>=0 and coord[0]<8 and coord[1]>=0 and coord[1]<8:
				wht[i, coord[1], coord[0]] = 1  # all 0 ???
				# print('wht {} {} {} to 1'.format(i, coord[1], coord[0]))
		wht = wht.astype(np.float32)

		# if 'n' == self.opts.if_scale_gs:        # hm gaussian,  directly cv resize to other scale directly
		# 	li2_gs = []
		# 	li_jts_gauss = []  # for jt geenration
		# 	for jt in joint_img:    # first layer
		# 		hm = np.zeros(shp_t)
		# 		hm = ut_t.draw_gaussian(hm, jt[:2], self.opts.gauss_sigma)
		# 		li_jts_gauss.append(hm)
		# 	li2_gs.append(li_jts_gauss)     # list of cv image?
		# 	sz = [shp_t[1], shp_t[0]]  # to size
		#
		# 	for i in range(1, self.opts.n_stg_D):
		# 		li_jts_t = []
		# 		sz = [e // 2 for e in sz]  # down two every time, 32, 16, 8
		# 		for hm in li_jts_gauss: # all from first layers
		# 			li_jts_t.append(transforms.ToTensor()(cv2.resize(hm, tuple(sz))).double())
		# 		li2_gs.append(li_jts_t) # add list of tensor?
		# 	for i, hm in enumerate(li_jts_gauss):
		# 		li_jts_gauss[i] = transforms.ToTensor()(hm) # to tensor
		# else:
		# 	raise ValueError('scale gaussian not implemented yet')

		# return img_patch, joint_img, joint_vis, joints_have_depth # single item
		return {'img_patch': img_patch}, {'joint_hm': joint_img, 'vis': joint_vis, 'if_depth_v': joints_have_depth,
		                                  'if_SYN_v': if_SYN, 'wts_D':wht, 'joint_img': joint_img_ori,
		                                  'joint_hm_es': joint_evoSkel}  # only image transformed to tensor other still np,  li2_gs  4x17[ 1x64x64]      # for regression later

	# else:
	# 	img_patch = self.transform(img_patch)   # optional data feed
	# 	return {'img_patch':img_patch}

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

# helper functions
def get_aug_config():
	scale_factor = 0.25
	rot_factor = 30
	color_factor = 0.2

	scale = np.clip(np.random.randn(), -1.0, 1.0) * scale_factor + 1.0
	rot = np.clip(np.random.randn(), -2.0,
	              2.0) * rot_factor if random.random() <= 0.6 else 0
	do_flip = random.random() <= 0.5
	c_up = 1.0 + color_factor
	c_low = 1.0 - color_factor
	color_scale = [random.uniform(c_low, c_up), random.uniform(c_low, c_up), random.uniform(c_low, c_up)]

	do_occlusion = random.random() <= 0.5

	return scale, rot, do_flip, color_scale, do_occlusion

def generate_patch_image(cvimg, bbox, do_flip, scale, rot, do_occlusion, input_shape=(256, 256)):
	# return skimage RGB,  h,w,c
	img = cvimg.copy()
	img_height, img_width, img_channels = img.shape # h,w,c

	# synthetic occlusion
	if do_occlusion:
		while True:
			area_min = 0.0
			area_max = 0.7
			synth_area = (random.random() * (area_max - area_min) + area_min) * bbox[2] * bbox[3]

			ratio_min = 0.3
			ratio_max = 1 / 0.3
			synth_ratio = (random.random() * (ratio_max - ratio_min) + ratio_min)

			synth_h = math.sqrt(synth_area * synth_ratio)
			synth_w = math.sqrt(synth_area / synth_ratio)
			synth_xmin = random.random() * (bbox[2] - synth_w - 1) + bbox[0]
			synth_ymin = random.random() * (bbox[3] - synth_h - 1) + bbox[1]

			if synth_xmin >= 0 and synth_ymin >= 0 and synth_xmin + synth_w < img_width and synth_ymin + synth_h < img_height:
				xmin = int(synth_xmin)
				ymin = int(synth_ymin)
				w = int(synth_w)
				h = int(synth_h)
				img[ymin:ymin + h, xmin:xmin + w, :] = np.random.rand(h, w, 3) * 255
				break

	bb_c_x = float(bbox[0] + 0.5 * bbox[2])
	bb_c_y = float(bbox[1] + 0.5 * bbox[3])
	bb_width = float(bbox[2])
	bb_height = float(bbox[3])

	if do_flip:
		img = img[:, ::-1, :]
		bb_c_x = img_width - bb_c_x - 1

	trans = gen_trans_from_patch_cv(bb_c_x, bb_c_y, bb_width, bb_height, input_shape[1], input_shape[0], scale, rot,
	                                inv=False)  # is bb aspect needed? yes, otherwise patch distorted
	img_patch = cv2.warpAffine(img, trans, (int(input_shape[1]), int(input_shape[0])), flags=cv2.INTER_LINEAR)  # is there channel requirements

	img_patch = img_patch[:, :, ::-1].copy()
	img_patch = img_patch.astype(np.float32)

	return img_patch, trans

def rotate_2d(pt_2d, rot_rad):
	x = pt_2d[0]
	y = pt_2d[1]
	sn, cs = np.sin(rot_rad), np.cos(rot_rad)
	xx = x * cs - y * sn
	yy = x * sn + y * cs
	return np.array([xx, yy], dtype=np.float32)

def gen_trans_from_patch_cv(c_x, c_y, src_width, src_height, dst_width, dst_height, scale, rot, inv=False):
	# augment size with scale
	src_w = src_width * scale
	src_h = src_height * scale
	src_center = np.array([c_x, c_y], dtype=np.float32)

	# augment rotation
	rot_rad = np.pi * rot / 180
	src_downdir = rotate_2d(np.array([0, src_h * 0.5], dtype=np.float32), rot_rad)
	src_rightdir = rotate_2d(np.array([src_w * 0.5, 0], dtype=np.float32), rot_rad)

	dst_w = dst_width
	dst_h = dst_height
	dst_center = np.array([dst_w * 0.5, dst_h * 0.5], dtype=np.float32)
	dst_downdir = np.array([0, dst_h * 0.5], dtype=np.float32)
	dst_rightdir = np.array([dst_w * 0.5, 0], dtype=np.float32)

	src = np.zeros((3, 2), dtype=np.float32)
	src[0, :] = src_center
	src[1, :] = src_center + src_downdir
	src[2, :] = src_center + src_rightdir

	dst = np.zeros((3, 2), dtype=np.float32)
	dst[0, :] = dst_center
	dst[1, :] = dst_center + dst_downdir
	dst[2, :] = dst_center + dst_rightdir

	if inv:
		trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
	else:
		trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

	return trans

def trans_point2d(pt_2d, trans):
	src_pt = np.array([pt_2d[0], pt_2d[1], 1.]).T
	dst_pt = np.dot(trans, src_pt)
	return dst_pt[0:2]

def genDsLoader(opts, mode=0, if_trainShuffle=True):
	'''
	generate dataset, dataLoader and iterator list (singler)
	:param opts:
	:param mode: 0 for train , 1 for test, 2 for test in loop (with folding)
	:return: datasets_li, loader_li, iterator li
	'''
	trans = transforms.Compose([
		transforms.ToTensor(),
		transforms.Normalize(mean=opts.pixel_mean, std=opts.pixel_std)])
	if 0 == mode:
		adpDs_li, loader_li, iter_li = [], [], []
		for i in range(len(opts.trainset)):
			ds = eval(opts.trainset[i])("train", opts=opts)
			ds_adp = AdpDataset_3d(ds, opts.ref_joints_name, if_trainShuffle, trans, opts=opts) # if aug here
			loader = DataLoader(dataset=ds_adp, batch_size=opts.num_gpus * opts.batch_size // len(opts.trainset),
			                    shuffle=if_trainShuffle, num_workers=opts.n_thread, pin_memory=opts.if_pinMem) # ave the batch to each ds
			itr_loader = iter(loader)       # this is purely ScanAva, why Human36M coco error?!!
			# to list
			adpDs_li.append(ds_adp)
			loader_li.append(loader)
			iter_li.append(itr_loader)
		return adpDs_li, loader_li, iter_li
	elif 1 == mode:
		ds = eval(opts.testset)("test", opts=opts)
		ds_adp = AdpDataset_3d(ds, opts.ref_joints_name, False, trans, opts=opts)
		loader = DataLoader(dataset=ds_adp, batch_size=opts.num_gpus * opts.batch_size // len(opts.trainset),
		                    shuffle=False, num_workers=opts.n_thread, pin_memory=opts.if_pinMem)
		itr_loader = iter(loader)
		return ds_adp, loader, itr_loader
	elif 2 == mode:
		ds = eval(opts.testset)("testInLoop", opts=opts)
		ds_adp = AdpDataset_3d(ds, opts.ref_joints_name, False, trans, opts=opts)  # few test in loop
		loader = DataLoader(dataset=ds_adp, batch_size=opts.num_gpus * opts.batch_size // len(opts.trainset),
		                    shuffle=False, num_workers=opts.n_thread, pin_memory=opts.if_pinMem)
		itr_loader = iter(loader)
		return ds_adp, loader, itr_loader
	else:
		print('no such mode defined')
		return -1

def genLoaderFromDs(ds, trans=None, opts={}, if_train=False):
	'''
	directly translate input ds input loader and adp_ds. Single ds input.
	:param ds:
	:return:
	'''
	if not trans:  # if not given, give tensor normalization to overide
		trans = transforms.Compose([
			transforms.ToTensor(),
			transforms.Normalize(mean=opts.pixel_mean, std=opts.pixel_std)])    # no trans give default trans
	# if 'test' in ds.data_split:
	# 	is_train = False
	# else:
	# 	is_train = True
	ds_adp = AdpDataset_3d(ds, opts.ref_joints_name, if_train, trans, opts=opts)    # false train
	loader = DataLoader(dataset=ds_adp, batch_size=opts.num_gpus * opts.batch_size, shuffle=if_train, num_workers=opts.n_thread, pin_memory=opts.if_pinMem)
	iterator = iter(loader)     # shuffle false

	return ds_adp, loader, iterator