#!/usr/bin/env python3
# encoding: utf-8

# Copyright 2019 Thai-Son Nguyen
# Licensed under the Apache License, Version 2.0 (the "License")

import time
import os
import copy
import random
import numpy as np
import math
import argparse

import matplotlib.pyplot as plt

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

from pynn.decoder.seq2seq import beam_search
from pynn.io.kaldi_seq import ScpStreamReader
from pynn.net.seq2seq import Seq2Seq

parser = argparse.ArgumentParser(description='pynn')

parser.add_argument('--model-dic', help='model dictionary', required=True)

parser.add_argument('--dict', help='dictionary file', default=None)
parser.add_argument('--source', help='path to srouce scp', required=True)
parser.add_argument('--downsample', help='concated frames', type=int, default=1)
parser.add_argument('--mean-sub', help='mean subtraction', action='store_true')
parser.add_argument('--len-norm', help='length normalization', action='store_true')

parser.add_argument('--batch-size', help='batch size', type=int, default=32)
parser.add_argument('--space', help='space token', type=str, default='<space>')
parser.add_argument('--beam-size', help='beam size', type=int, default=8)
parser.add_argument('--max-len', help='max len', type=int, default=200)

parser.add_argument('--head', help='head number', type=int, default=0)
parser.add_argument('--win', help='max len', type=int, default=10000)

if __name__ == '__main__':
    args = parser.parse_args()

    dic = None
    if args.dict is not None:
        dic = {}
        fin = open(args.dict, 'r')
        for line in fin:
            tokens = line.split()
            dic[int(tokens[1])] = tokens[0]

    use_gpu = torch.cuda.is_available()
    device = torch.device('cuda' if use_gpu else 'cpu')

    mdic = torch.load(args.model_dic)
    model = Seq2Seq(**mdic['params']).to(device)
    model.load_state_dict(mdic['state'])
    model.eval()

    reader = ScpStreamReader(args.source, mean_sub=args.mean_sub, sort_src=True, downsample=args.downsample)
    reader.initialize()

    space, beam_size, max_len = args.space, args.beam_size, args.max_len
    win = args.win

    with torch.no_grad():
        while True:
            #src_seq, src_mask, tgt_seq = reader.next_batch(1)
            src_seq, src_mask, utts = reader.read_batch_utt(1)
            if len(utts) == 0: break
            utt = utts[0]

            fout = open(utt+'.info', 'w')
            time_len = src_seq.size(1)
            for i in range((time_len-1) // win + 1):
                e = min(time_len, win*(i+1))
                src = src_seq[:, 0:e, :]
                mask = src_mask[:, 0:e]

                src, mask = src.to(device), mask.to(device)
                hypos, scores = beam_search(model, src, mask, device, beam_size, max_len, len_norm=args.len_norm)
                hypo, score = [], []
                for token, s in zip(hypos[0], scores[0]):
                    if token == 2: break
                    hypo.append(token)
                    score.append(s)
                print(hypo)

                tgt = torch.LongTensor([1]+hypo+[2]).to(device).view(1, -1)
                attn = model.attend(src, mask, tgt)[1]
                attn = attn[0]

                #print(attn)
                cs = torch.cumsum(attn[args.head], dim=1)
                sp = cs.le(0.05).sum(dim=1)
                ep = cs.le(0.95).sum(dim=1)

                sp, ep = sp.cpu().numpy(), ep.cpu().numpy()

                fout.write("# End=%d\n" % (e//4))
                for j, token in enumerate(hypo):
                    js, je = sp[j], ep[j]
                    fout.write("%s: %d %d  %f\n" % (dic[token-2], js, je, math.exp(score[j])))
                fout.write("\n")

                attn = attn.view(-1, attn.size(2))
                attn = attn.cpu().numpy()
                fig = plt.figure()
                plt.imshow(attn)

                if dic is not None:
                    hypo = [dic[token-2].replace(space, '') for token in hypo]
                plt.title(' '.join(map(str, hypo)))

                plt.savefig('%s-%d.png' % (utt, i))

            fout.close()


# -*- coding: utf-8 -*-
'''
This script performs transformation on images and dumps the output arrays as npzfiles, which saves time for training.
'''

import argparse
import os
import torch
from PIL import Image
import base64
from io import BytesIO
from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize, RandomResizedCrop
from tqdm import tqdm
import numpy as np

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_dir", type=str, default="../Multimodal_Retrieval/", help="the directory which stores the image tsvfiles")
    parser.add_argument("--image_resolution", type=int, default=224, help="the resolution of transformed images, default to 224*224")
    return parser.parse_args()

def _convert_to_rgb(image):
    return image.convert('RGB')

def build_transform(resolution):
    normalize = Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
    return Compose([
            Resize(resolution, interpolation=Image.BICUBIC),
            CenterCrop(resolution),
            _convert_to_rgb,
            ToTensor(),
            normalize,
        ])

if __name__ == "__main__":
    args = parse_args()
    transform = build_transform(args.image_resolution)
    train_path = os.path.join(args.data_dir, "train_imgs.tsv")
    val_path = os.path.join(args.data_dir, "valid_imgs.tsv")
    test_path = os.path.join(args.data_dir, "test_imgs.tsv")
    for path, split in zip((train_path, val_path, test_path), ("train", "valid", "test")):
        assert os.path.exists(path), "the {} filepath {} not exists!".format(split, path)
        print("begin to transform {} split".format(split))
        image_dict = {}
        with open(path, "r") as fin:
            for line in tqdm(fin):
                img_id, b64 = line.strip().split("\t")
                image = Image.open(BytesIO(base64.urlsafe_b64decode(b64)))
                image_array = transform(image).numpy()
                image_dict[img_id] = image_array
        output_path = "{}.{}.npz".format(path[:-4], args.image_resolution)
        np.savez(output_path, **image_dict)
        print("finished transforming {} images for {} split, the output is saved at {}".format(len(image_dict), split, output_path))
    print("done!")


import tensorflow.compat.v1 as tf
import numpy as np
import os
import glob
from imageio import imread, imsave
import cv2

def preprocess(img):
    return (img / 255. - 0.5) * 2

def deprocess(img):
    return (img + 1) / 2

class beautyGAN:
    def __init__(self, img_size=256):
        tf.disable_eager_execution()
        tf.reset_default_graph()
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())
        saver = tf.train.import_meta_graph(os.path.join('beautyGAN/model', 'model.meta'))
        saver.restore(self.sess, tf.train.latest_checkpoint('beautyGAN/model'))
        graph = tf.get_default_graph()
        self.X = graph.get_tensor_by_name('X:0')
        self.Y = graph.get_tensor_by_name('Y:0')
        self.Xs = graph.get_tensor_by_name('generator/xs:0')
        self.img_size = img_size

    def generate(self, input_img, makeups_path, n=10):
        no_makeup = cv2.resize(imread(input_img), (self.img_size, self.img_size))
        X_img = np.expand_dims(preprocess(no_makeup), 0)
        makeups = glob.glob(os.path.join(makeups_path, '*.png'))
        makeups = np.random.choice(makeups, n, replace=False)
        result = np.ones((2 * self.img_size, (len(makeups) + 1) * self.img_size, 3))
        result[self.img_size: 2 * self.img_size, :self.img_size] = no_makeup / 255.
        for i in range(len(makeups)):
            makeup = cv2.resize(imread(makeups[i]), (self.img_size, self.img_size))
            Y_img = np.expand_dims(preprocess(makeup), 0)
            Xs_ = self.sess.run(self.Xs, feed_dict={self.X: X_img, self.Y: Y_img})
            Xs_ = deprocess(Xs_)
            result[:self.img_size, (i + 1) * self.img_size: (i + 2) * self.img_size] = makeup / 255.
            result[self.img_size: 2 * self.img_size, (i + 1) * self.img_size: (i + 2) * self.img_size] = Xs_[0]

        imsave('result.jpg', result)


#!/usr/bin/env python
from __future__ import division
import os
import re
import argparse
from itertools import izip, compress
from collections import defaultdict
import numpy as np
import pandas as pd
from Bio import SeqIO
import matplotlib.pyplot as plt

def get_parser():
    """
    Create a parser and add arguments
    """
    parser = argparse.ArgumentParser()
    parser.add_argument('-d', '--depth', type=int, default=4, help='minimum read depth')
    parser.add_argument('-p', '--pmtsize', type=int, default=1000, help='promoter size')
    parser.add_argument('gtf', help='GTF file')
    parser.add_argument('fasta', help='reference genome FASTA file')
    parser.add_argument('cgmap', help='CGmap file')
    return parser

def const_gtftree(gtffile):
    """
    Read a GTF file and convert it to a nested dictionary
    """
    gtftree = defaultdict(lambda: defaultdict(list))
    with open(gtffile) as infile:
        for line in infile:
            if not line.startswith('#'):
                gene_id = None
                transcript_id = None
                line = line.strip().split('\t')
                chr = line[0]
                feature = line[2]
                start = int(line[3]) - 1
                end = int(line[4])
                strand = line[6]
                attributes = line[8].split(';')
                if feature == 'exon':
                    for atb in attributes:
                        if 'gene_id' in atb:
                            gene_id = atb.strip().split()[1][1:-1]
                        elif 'transcript_id' in atb:
                            transcript_id = atb.strip().split()[1][1:-1]
                    if gene_id and transcript_id:
                        gtftree[chr][(gene_id, strand)].append((start, end))
    return gtftree

def const_ctxstr(reffile):
    """
    Construct methylation context strings from a reference genome FASTA file
    """
    with open(reffile) as infile:
        fasta = SeqIO.to_dict(SeqIO.parse(infile, 'fasta'))
        for chr in fasta:
            fasta[chr] = str(fasta[chr].seq).upper()
    ctxstr = {}
    for chr in fasta:
        ctxstr[chr] = ['-']*len(fasta[chr])
        cg = [match.start() for match in re.finditer(r'(?=(CG))', fasta[chr])]
        for pos in cg:
            ctxstr[chr][pos] = 'X'
        chg = [match.start() for match in re.finditer(r'(?=(C[ACT]G))', fasta[chr])]
        for pos in chg:
            ctxstr[chr][pos] = 'Y'
        chh = [match.start() for match in re.finditer(r'(?=(C[ACT][ACT]))', fasta[chr])]
        for pos in chh:
            ctxstr[chr][pos] = 'Z'
        rcg = [match.start()-1 for match in re.finditer(r'(?<=(CG))', fasta[chr])]
        for pos in rcg:
            ctxstr[chr][pos] = 'x'
        rchg = [match.start()-1 for match in re.finditer(r'(?<=(C[AGT]G))', fasta[chr])]
        for pos in rchg:
            ctxstr[chr][pos] = 'y'
        rchh = [match.start()-1 for match in re.finditer(r'(?<=([AGT][AGT]G))', fasta[chr])]
        for pos in rchh:
            ctxstr[chr][pos] = 'z'
    for chr in ctxstr:
        ctxstr[chr] = ''.join(ctxstr[chr])
    return ctxstr

def const_cgmap(ctxstr, cgmapfile, readdepth=4):
    """
    Construct lists of methylation levels from a CGmap file for rapid access
    """
    cgmap = {}
    with open(cgmapfile) as infile:
        for chr in ctxstr.keys():
            cgmap[chr] = ['-' for _ in xrange(len(ctxstr[chr]))]
        for line in infile:
            line = line.strip().split()
            chr = line[0]
            pos = int(line[2]) - 1 # Transfer to 0-based
            context = line[3]
            level = float(line[5])
            depth = int(line[7])
            if context in ['CG', 'CHG', 'CHH'] and depth >= readdepth:
                cgmap[chr][pos] = level
    return cgmap

def calc_bulk(ctxstr, cgmap):
    """
    Compute the global methylation level in CG/CHG/CHH
    """
    inv_ctxs = {'X': 'CG', 'Y': 'CHG', 'Z': 'CHH'}
    bulk = defaultdict(list)
    for chr in set(ctxstr) & set(cgmap):
        for tag, mlevel in izip(ctxstr[chr], cgmap[chr]):
            tag = tag.upper()
            if tag in inv_ctxs and mlevel != '-':
                bulk[inv_ctxs[tag]].append(mlevel)
    return bulk

def calc_mlevel(ctxstr, cgmap, gtftree, pmtsize=1000):
    """
    Compute the mean methylation level of promoter/gene/exon/intron/IGN in each gene
    """
    inv_ctxs = {'X': 'CG', 'Y': 'CHG', 'Z': 'CHH'}
    ign = defaultdict(list)
    mtable = defaultdict(lambda: defaultdict(lambda: defaultdict(float)))
    counter = defaultdict(lambda: defaultdict(int))
    for chr in set(ctxstr) & set(cgmap) & set(gtftree):
        mask = [1]*len(cgmap[chr])
        for (gene_id, strand) in gtftree[chr]:
            feature_mlevels = defaultdict(lambda: defaultdict(list))
            gstart = min(gtftree[chr][(gene_id, strand)])[0]
            gend = max(gtftree[chr][(gene_id, strand)])[1]
            mask[gstart:gend] = [0]*(gend - gstart)
            if strand == '+':
                for (pos, (tag, mlevel)) in enumerate(izip(ctxstr[chr][gstart-pmtsize:gstart], cgmap[chr][gstart-pmtsize:gstart])):
                    tag = tag.upper()
                    if tag in inv_ctxs and mlevel != '-':
                        feature_mlevels[inv_ctxs[tag]]['pmt'].append(mlevel)
            elif strand == '-':
                for (pos, (tag, mlevel)) in enumerate(izip(ctxstr[chr][gend:gend+pmtsize], cgmap[chr][gend:gend+pmtsize])):
                    tag = tag.upper()
                    if tag in inv_ctxs and mlevel != '-':
                        feature_mlevels[inv_ctxs[tag]]['pmt'].append(mlevel)
            for (pos, (tag, mlevel)) in enumerate(izip(ctxstr[chr][gstart:gend], cgmap[chr][gstart:gend])):
                tag = tag.upper()
                inexon = False
                if tag in inv_ctxs and mlevel != '-':
                    feature_mlevels[inv_ctxs[tag]]['gene'].append(mlevel)
                    for exon in gtftree[chr][(gene_id, strand)]:
                        if exon[0] <= pos+gstart < exon[1]:
                            feature_mlevels[inv_ctxs[tag]]['exon'].append(mlevel)
                            inexon = True
                            break
                    if not inexon:
                        feature_mlevels[inv_ctxs[tag]]['intron'].append(mlevel)
            for ctx in ['CG', 'CHG', 'CHH']:
                for feature in feature_mlevels[ctx]:
                    counter[ctx][feature] += len(feature_mlevels[ctx][feature])
                    mtable[ctx][gene_id][feature] = np.mean(feature_mlevels[ctx][feature])
        for (pos, (tag, mlevel)) in enumerate(izip(ctxstr[chr], cgmap[chr])):
            tag = tag.upper()
            if (tag in inv_ctxs) and (mask[pos] == 1) and (mlevel != '-'):
                ign[inv_ctxs[tag]].append(mlevel)
    for ctx in ign:
        ign[ctx] = np.mean(ign[ctx])
    cg_table = pd.DataFrame(mtable['CG']).T
    cg_table = cg_table[['pmt', 'gene', 'exon', 'intron']]
    chg_table = pd.DataFrame(mtable['CHG']).T
    chg_table = chg_table[['pmt', 'gene', 'exon', 'intron']]
    chh_table = pd.DataFrame(mtable['CHH']).T
    chh_table = chh_table[['pmt', 'gene', 'exon', 'intron']]
    return ign, cg_table, chg_table, chh_table

def plot_bar(dataframe, bulk, ctx):
    colors = { 'CG': ( 38/255, 173/255,  84/255),
              'CHG': ( 44/255, 180/255, 234/255),
              'CHH': (249/255,  42/255,  54/255)}
    dataframe = dataframe*100
    plt.switch_backend('Agg')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax = dataframe.plot(ax=ax, kind='bar', grid=False, rot=0, color=colors[ctx], ylim=(0, 100))
    ax.set_ylabel('Methylation Level (%)', fontsize='xx-large', fontweight='bold')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_linewidth(2)
    ax.spines['left'].set_linewidth(2)
    #ax.spines['bottom'].set_position(('outward', 5))
    #ax.spines['left'].set_position(('outward', 5))
    ax.tick_params(direction='out', length=6, width=2, labelsize='xx-large', top='off', right='off')
    for label in ax.xaxis.get_ticklabels():
        label.set_fontweight('bold')
    for label in ax.yaxis.get_ticklabels():
        label.set_fontweight('bold')
    ax.set_title(ctx, fontsize='xx-large', weight='bold')
    #ax.axhline(y=np.mean(bulk[ctx])*100, linewidth=2, linestyle='--', color='k')
    fig.tight_layout()
    return ax

def plot_feature_mlevel(bulk, ign, cg_table, chg_table, chh_table):
    cg = cg_table.mean()
    cg = cg.set_value('genome', np.mean(bulk['CG']))
    cg = cg.set_value('IGN', ign['CG'])
    cg = cg[['genome', 'pmt', 'gene', 'exon', 'intron', 'IGN']]
    cg_ax = plot_bar(cg, bulk, 'CG')
    chg = chg_table.mean()
    chg = chg.set_value('genome', np.mean(bulk['CHG']))
    chg = chg.set_value('IGN', ign['CHG'])
    chg = chg[['genome', 'pmt', 'gene', 'exon', 'intron', 'IGN']]
    chg_ax = plot_bar(chg, bulk, 'CHG')
    chh = chh_table.mean()
    chh = chh.set_value('genome', np.mean(bulk['CHH']))
    chh = chh.set_value('IGN', ign['CHH'])
    chh = chh[['genome', 'pmt', 'gene', 'exon', 'intron', 'IGN']]
    chh_ax = plot_bar(chh, bulk, 'CHH')
    return cg_ax, chg_ax, chh_ax

def plot_bulkmean(bulk):
    bulk_mean = {}
    for ctx in ['CG', 'CHG', 'CHH']:
        bulk_mean[ctx] = np.mean(bulk[ctx])
    bulk_mean = pd.Series(bulk_mean)*100
    colors = { 'CG': ( 38/255, 173/255,  84/255),
              'CHG': ( 44/255, 180/255, 234/255),
              'CHH': (249/255,  42/255,  54/255)}
    plt.switch_backend('Agg')
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax = bulk_mean.plot(ax=ax, kind='bar', grid=False, rot=0, color=[colors[ctx] for ctx in ['CG', 'CHG', 'CHH']], ylim=(0, 100))
    ax.set_ylabel('Methylation Level (%)', fontsize='xx-large', fontweight='bold')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_linewidth(2)
    ax.spines['left'].set_linewidth(2)
    ax.tick_params(direction='out', length=6, width=2, labelsize='xx-large', top='off', right='off')
    for label in ax.xaxis.get_ticklabels():
        label.set_fontweight('bold')
    for label in ax.yaxis.get_ticklabels():
        label.set_fontweight('bold')
    fig.tight_layout()
    return ax

def plot_bulkhist(bulk):
    colors = { 'CG': ( 38/255, 173/255,  84/255),
              'CHG': ( 44/255, 180/255, 234/255),
              'CHH': (249/255,  42/255,  54/255)}
    plt.switch_backend('Agg')
    fig = plt.figure(figsize=(8, 3))
    axes = {}
    for i, ctx in enumerate(['CG', 'CHG', 'CHH']):
        if i == 0:
            axes[ctx] = fig.add_axes((0.15, 0.25, 0.25, 0.65))
            #axes[ctx] = fig.add_subplot(131)
            axes[ctx].hist(bulk[ctx], weights=np.repeat(1.0/len(bulk[ctx]), len(bulk[ctx])), color=colors[ctx])
            axes[ctx].spines['top'].set_visible(False)
            axes[ctx].spines['right'].set_visible(False)
            axes[ctx].spines['bottom'].set_linewidth(2)
            axes[ctx].spines['left'].set_linewidth(2)
            axes[ctx].spines['left'].set_position(('outward', 10))
            plt.setp(axes[ctx].get_xticklabels(), visible=False)
            axes[ctx].tick_params(axis='y', direction='out', right='off', length=6, width=2, labelsize='xx-large')
            axes[ctx].tick_params(axis='x', top='off', bottom='off')
            for label in axes[ctx].yaxis.get_ticklabels():
                label.set_fontweight('bold')
            axes[ctx].set_ylabel('Fraction', fontsize='xx-large', fontweight='bold')
        else:
            axes[ctx] = fig.add_axes((0.15 + (0.25 + 0.025) * i, 0.25, 0.25, 0.65))
            #axes[ctx] = fig.add_subplot(1, 3, i+1)
            axes[ctx].hist(bulk[ctx], weights=np.repeat(1.0/len(bulk[ctx]), len(bulk[ctx])), color=colors[ctx])
            axes[ctx].spines['top'].set_visible(False)
            axes[ctx].spines['left'].set_visible(False)
            axes[ctx].spines['right'].set_visible(False)
            axes[ctx].spines['bottom'].set_linewidth(2)
            axes[ctx].spines['left'].set_linewidth(2)
            plt.setp(axes[ctx].get_xticklabels(), visible=False)
            plt.setp(axes[ctx].get_yticklabels(), visible=False)
            axes[ctx].tick_params(top='off', bottom='off', left='off', right='off')
        axes[ctx].set_ylim(0, 1)
        axes[ctx].set_yticks(np.arange(0, 1.2, 0.2))
        axes[ctx].set_xlim(-0.025, 1.025)
        axes[ctx].set_xlabel(ctx, fontsize='xx-large', fontweight='bold')
        fig.suptitle('Methylation Level (0 -> 100%)', x=0.55, y=0.1, fontsize='xx-large', fontweight='bold')
    return fig

def calc_genomewide(ctxstr, cgmap, winsize=200000):
    inv_ctxs = {'X': 'CG', 'Y': 'CHG', 'Z': 'CHH'}
    win_mlevel = defaultdict(list)
    win_x = []
    pos = 0
    if 'chr' in ctxstr.keys()[0].lower():
        chrs = sorted(ctxstr.keys(), key=lambda s: s[3:])
    else:
        chrs = sorted(ctxstr.keys())
    #chrs = map(str, range(1, 23)) + ['X', 'Y']
    for chr in chrs:
        start = 0
        while (start + winsize) <= len(ctxstr[chr]):
            win_x.append(pos+(winsize/2))
            tmp = defaultdict(list)
            for tag, mlevel in izip(ctxstr[chr][start:start+winsize], cgmap[chr][start:start+winsize]):
                tag = tag.upper()
                if tag in inv_ctxs and mlevel != '-':
                    tmp[inv_ctxs[tag]].append(mlevel)
            for ctx in ['CG', 'CHG', 'CHH']:
                win_mlevel[ctx].append(np.mean(tmp[ctx])*100)
            start += winsize
            pos += winsize
    return win_x, win_mlevel

def plot_genomewide(ctxstr, gpos, gmlevel):
    colors = { 'CG': ( 38/255, 173/255,  84/255),
              'CHG': ( 44/255, 180/255, 234/255),
              'CHH': (249/255,  42/255,  54/255)}
    if 'chr' in ctxstr.keys()[0].lower():
        chrs = sorted(ctxstr.keys(), key=lambda s: s[3:])
    else:
        chrs = sorted(ctxstr.keys())
    #chrs = map(str, range(1, 23)) + ['X', 'Y']
    vlines = [0]
    for i, chr in enumerate(chrs):
        vlines.append(vlines[i] + len(ctxstr[chr]))
    plt.switch_backend('Agg')
    fig = plt.figure(figsize=(16, 4.5))
    ax = fig.add_subplot(111)
    ax.plot(gpos, gmlevel['CG'], color=colors['CG'], linewidth=1.5, label='CG')
    ax.plot(gpos, gmlevel['CHG'], color=colors['CHG'], linewidth=1.5, label='CHG')
    ax.plot(gpos, gmlevel['CHH'], color=colors['CHH'], linewidth=1.5, label='CHH')
    ax.set_ylim(0, 100)
    ax.set_xlim(0, vlines[-1])
    for pos in vlines[1:-1]:
        ax.axvline(x=pos, linestyle='--', linewidth=1.5, color='gray')
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    ax.spines['bottom'].set_linewidth(2)
    ax.spines['left'].set_linewidth(2)
    ax.spines['left'].set_position(('outward', 10))
    for label in ax.xaxis.get_ticklabels():
        label.set_fontweight('bold')
    for label in ax.yaxis.get_ticklabels():
        label.set_fontweight('bold')
    ax.tick_params(direction='out', length=6, width=2, labelsize='large', top='off', right='off', bottom='off')
    ax.set_xticks([(vlines[i] + vlines[i+1])/2 for i in xrange(len(vlines) - 1)])
    ax.set_xticklabels(chrs)
    ax.set_xlabel('Chromosome', fontsize='xx-large', fontweight='bold')
    ax.set_ylabel('Methylation Level (%)', fontsize='xx-large', fontweight='bold')
    ax.legend(loc='upper right', fontsize='large', frameon=False)
    fig.tight_layout()
    return ax

def main():
    parser = get_parser()
    args = parser.parse_args()
    root = os.path.splitext(os.path.basename(args.cgmap))[0]
    ctxstr = const_ctxstr(args.fasta)
    cgmap = const_cgmap(ctxstr, args.cgmap, args.depth)
    gtftree = const_gtftree(args.gtf)
    bulk = calc_bulk(ctxstr, cgmap)
    plt.switch_backend('Agg')
    bulk_ax = plot_bulkmean(bulk)
    fig = bulk_ax.get_figure()
    fig.savefig('{}.bulk.mean.png'.format(root), dpi=300)
    plt.close(fig)
    bulk_fig = plot_bulkhist(bulk)
    bulk_fig.savefig('{}.bulk.hist.png'.format(root), dpi=300)
    plt.close(fig)
    ign, cg_table, chg_table, chh_table = calc_mlevel(ctxstr, cgmap, gtftree, args.pmtsize)
    cg_table.to_csv('{}.feature.CG.txt'.format(root), sep='\t', float_format='%.3f')
    chg_table.to_csv('{}.feature.CHG.txt'.format(root), sep='\t', float_format='%.3f')
    chh_table.to_csv('{}.feature.CHH.txt'.format(root), sep='\t', float_format='%.3f')
    cg_ax, chg_ax, chh_ax = plot_feature_mlevel(bulk, ign, cg_table, chg_table, chh_table)
    fig = cg_ax.get_figure()
    fig.savefig('{}.feature.CG.png'.format(root), dpi=300)
    plt.close(fig)
    fig = chg_ax.get_figure()
    fig.savefig('{}.feature.CHG.png'.format(root), dpi=300)
    plt.close(fig)
    fig = chh_ax.get_figure()
    fig.savefig('{}.feature.CHH.png'.format(root), dpi=300)
    plt.close(fig)
    gpos, gmlevel = calc_genomewide(ctxstr, cgmap)
    gax = plot_genomewide(ctxstr, gpos, gmlevel)
    fig = gax.get_figure()
    fig.savefig('{}.genomewide.png'.format(root), dpi=300)
    plt.close(fig)

if __name__ == '__main__':
    main()


"""
classes_multicrystal.py
Class that combines multiple crystal classes with special functions

By Dan Porter, PhD
Diamond
2020

Version 1.1
Last updated: 27/09/21

Version History:
30/03/20 1.0    class MultiCrytal moved from classes_crystal
27/09/21 1.1    Added __get__, __repr__ methods

@author: DGPorter
"""

import numpy as np

from . import functions_general as fg
from . import functions_crystallography as fc
from .classes_plotting import MultiPlotting

__version__ = '1.1'

class MultiCrystal:
    """
    Multi_Crystal class for combining multiple phases
    """
    _scattering_type = 'xray'

    def __init__(self, crystal_list):
        """
        Multi-crystal class
        """
        self.crystal_list = crystal_list
        self.Plot = MultiPlotting(crystal_list)

    def set_scale(self, index, scale=1.0):
        """Set scale of crystal[index]"""
        xtl = self.crystal_list[index]
        xtl.scale = scale

    def orient_set_r(self, rotation):
        """Set rotation matrix in diffractometer frame"""
        for xtl in self.crystal_list:
            xtl.Cell.orientation.set_r(rotation)

    def orient_6circle(self, phi=0, chi=0, eta=0, mu=0):
        """Set rotation matrix using 6-circle diffractometer axes"""
        rotation = fc.diffractometer_rotation(phi, chi, eta, mu)
        self.orient_set_r(rotation)

    def set_labframe(self, lab):
        """Set transformation matrix between diffractometer and lab"""
        for xtl in self.crystal_list:
            xtl.Cell.orientation.set_lab(lab)

    def set_labframe_i16(self):
        """Set lab transformation matrix for beamline I16 at Diamond Light Source"""
        self.set_labframe([[0, 0, 1], [1, 0, 0], [0, 1, 0]])

    def setup_scatter(self, **kwargs):
        """
        Simple way to set scattering parameters, each parameter is internal to xtl (self)

        type        : self._scattering type               :  'xray','neutron','xray magnetic','neutron magnetic','xray resonant'
        energy_kev  : self._energy_kev                    :  radiation energy in keV
        wavelength_a: self._wavelength_a                  :  radiation wavelength in Angstrom
        powder_units: self._powder_units                  :  units to use when displaying/ plotting ['twotheta', 'd',' 'q']
        min_twotheta: self._scattering_min_two_theta      :  minimum detector (two-theta) angle
        max_twotheta: self._scattering_max_two_theta      :  maximum detector (two-theta) angle
        min_theta   : self._scattering_min_theta          :  minimum sample angle = -opening angle
        max_theta   : self._scattering_max_theta          :  maximum sample angle = opening angle
        theta_offset: self._scattering_theta_offset       :  sample offset angle
        specular    : self._scattering_specular_direction : [h,k,l] : reflections normal to sample surface
        parallel    : self._scattering_parallel_direction : [h,k,l] : reflections normal to sample surface
        """

        if 'type' in kwargs.keys():
            self._scattering_type = kwargs['type']

        for xtl in self.crystal_list:
            xtl.Scatter.setup_scatter(**kwargs)

    def print_all_reflections(self, energy_kev=None, print_symmetric=False,
                              min_intensity=None, max_intensity=None, units=None):
        """
        Prints a list of all allowed reflections at this energy
            energy_kev = energy in keV
            print_symmetric = False*/True : omits reflections with the same intensity at the same angle
            min_intensity = None/ 0.01 : omits reflections less than this (remove extinctions)
            max_intensity = None/ 0.01 : omits reflections greater than this (show extinctions only)
            units = None/ 'twotheta'/ 'q'/ 'dspace' : choose scattering angle units to display
        """

        if min_intensity is None: min_intensity = -1
        if max_intensity is None: max_intensity = np.inf

        if energy_kev is None:
            energy_kev = self.crystal_list[0].Scatter._energy_kev
        if units is None:
            units = self.crystal_list[0].Scatter._powder_units

        hkl = np.empty(shape=(0,3))
        tth = np.empty(shape=0)
        inten = np.empty(shape=0)
        name = np.empty(shape=0)
        for xtl in self.crystal_list:
            hkl_xtl = xtl.Cell.all_hkl(energy_kev, xtl.Scatter._scattering_max_twotheta)
            if not print_symmetric:
                hkl_xtl = xtl.Symmetry.remove_symmetric_reflections(hkl_xtl)
            hkl_xtl = xtl.Cell.sort_hkl(hkl_xtl)
            # remove [0,0,0]
            hkl_xtl = hkl_xtl[1:, :]

            tth_xtl = xtl.Cell.tth(hkl_xtl, energy_kev)
            inrange = np.all([tth_xtl < xtl.Scatter._scattering_max_twotheta, tth_xtl > xtl.Scatter._scattering_min_twotheta], axis=0)
            hkl_xtl = hkl_xtl[inrange, :]
            tth_xtl = tth_xtl[inrange]
            inten_xtl = xtl.Scatter.intensity(hkl_xtl)

            name = np.append(name, [xtl.name]*len(tth_xtl))
            hkl = np.vstack((hkl, hkl_xtl))
            tth = np.append(tth, tth_xtl)
            inten = np.append(inten, inten_xtl)

        # Sort reflections from all reflections
        idx = np.argsort(tth)
        name = name[idx]
        hkl = hkl[idx, :]
        tth = tth[idx]
        inten = inten[idx]

        units = units.lower()
        if units in ['d', 'dspc', 'dspace', 'd space', 'd-space', 'dspacing', 'd spacing', 'd-spacing']:
            unit_str = 'd-spacing'
            unit = fc.caldspace(tth, energy_kev)
        elif units in ['q', 'wavevector']:
            unit_str = 'Q'
            unit = fc.calqmag(tth, energy_kev)
        else:
            unit_str = 'TwoTheta'
            unit = tth

        fmt = '(%3.0f,%3.0f,%3.0f) %10.2f  %9.2f   %s\n'
        outstr = ''

        outstr += 'Energy = %6.3f keV\n' % energy_kev
        outstr += 'Radiation: %s\n' % self._scattering_type
        outstr += '( h, k, l)    %10s  Intensity  Crystal\n' % unit_str
        # outstr+= fmt % (hkl[0,0],hkl[0,1],hkl[0,2],tth[0],inten[0]) # hkl(0,0,0)
        count = 0
        for n in range(0, len(tth)):
            if inten[n] < min_intensity: continue
            if inten[n] > max_intensity: continue
            count += 1
            outstr += fmt % (hkl[n, 0], hkl[n, 1], hkl[n, 2], unit[n], inten[n], name[n])
        outstr += 'Reflections: %1.0f\n' % count
        return outstr

    def find_close_reflections(self, HKL, energy_kev=None, max_twotheta=2, max_angle=10):
        """
        Find reflections close to the given one  and return list
        """

        if energy_kev is None:
            energy_kev = self.crystal_list[0].Scatter._energy_kev

        HKL_tth = self.crystal_list[0].Cell.tth(HKL, energy_kev)
        HKL_Q = self.crystal_list[0].Cell.calculateQ(HKL)

        HKL_list = np.empty([0, 3])
        TTH_list = np.empty([0])
        ANGLE_list = np.empty([0])
        I_list = np.empty([0])
        NAMES_list = np.empty([0])
        for xtl in self.crystal_list:
            # xtl._scattering_type = self._scattering_type
            all_HKL = xtl.Cell.all_hkl(energy_kev, xtl._scattering_max_twotheta)
            all_TTH = xtl.Cell.tth(all_HKL, energy_kev)
            dif_TTH = np.abs(all_TTH - HKL_tth)
            all_Q = xtl.Cell.calculateQ(all_HKL)
            all_ANG = np.abs([fg.ang(HKL_Q, Q, 'deg') for Q in all_Q])
            selected = (dif_TTH < max_twotheta) * (all_ANG < max_angle)
            sel_HKL = all_HKL[selected, :]
            sel_TTH = all_TTH[selected]
            sel_ANG = all_ANG[selected]
            sel_INT = xtl.Scatter.intensity(sel_HKL)
            NAMES = np.asarray(xtl.name).repeat(len(sel_INT))

            HKL_list = np.append(HKL_list, sel_HKL, axis=0)
            TTH_list = np.append(TTH_list, sel_TTH)
            ANGLE_list = np.append(ANGLE_list, sel_ANG)
            I_list = np.append(I_list, sel_INT)
            NAMES_list = np.append(NAMES_list, NAMES)

        # Sort by TTH
        index = np.argsort(TTH_list)
        HKL_list = HKL_list[index, :]
        TTH_list = TTH_list[index]
        ANGLE_list = ANGLE_list[index]
        I_list = I_list[index]
        NAMES_list = NAMES_list[index]

        fmt = '(%3.0f,%3.0f,%3.0f) %8.2f %8.2f %9.2f %s\n'

        out = 'Energy = %6.3f keV\n' % energy_kev()
        out += '%s Reflection: (%3.0f,%3.0f,%3.0f)\n' % (self.crystal_list[0].name, HKL[0], HKL[1], HKL[2])
        out += '( h, k, l) TwoTheta Angle    Intensity Crystal\n'
        for n in range(0, len(TTH_list)):
            print(fmt % (
                HKL_list[n, 0], HKL_list[n, 1], HKL_list[n, 2], TTH_list[n], ANGLE_list[n], I_list[n], NAMES_list[n]))
        return out

    def remove(self, index=-1):
        """Remove crystal [index] from crystal list"""
        del(self.crystal_list[index])

    def __add__(self, other):
        if type(other) is MultiCrystal:
            return MultiCrystal(self.crystal_list+other.crystal_list)
        return MultiCrystal(self.crystal_list + [other])

    def __get__(self, item):
        return self.crystal_list[item]

    def info(self):
        """Display information about the contained crystals"""

        out = "Crystals: %d\n" % len(self.crystal_list)
        for n, xtl in enumerate(self.crystal_list):
            out += "%1.0f %r\n" % (n, xtl)
        return out

    def __repr__(self):
        s = ', '.join([xtl.name for xtl in self.crystal_list])
        return "MultiCrystal([%s])" % s

    def __str__(self):
        return self.info()


import numpy as np
import tensorflow as tf
import subprocess
import os
import sys
import transformer.model.transformer as tft
import transformer.model.embedding_layer as tf_embedding
import transformer.model.model_params as tf_params

def setGPU():
    freeGpu = subprocess.check_output('nvidia-smi -q | grep "Minor\|Processes" | grep "None" -B1 | tr -d " " | cut -d ":" -f2 | sed -n "1p"', shell=True)
    freeGpu = freeGpu.decode().strip()
    if len(freeGpu) == 0:
        print('no free GPU!')
        sys.exit(1)
    os.environ['CUDA_VISIBLE_DEVICES'] = freeGpu
    print('got GPU ' + (freeGpu))

    return(int(freeGpu.strip()))

class transformer_gan(object):

    def __init__(self, init_features=16, data_shape=(32,800), gp_scale=10,
                 batch_size=16, max_labels=100, num_chars=276, color_lambda=0,
                 ctc_lambda=1):

        self.num_chars = num_chars
        self.max_labels = max_labels
        self.batch_size = batch_size
        self.init_features = init_features
        self.scale = gp_scale
        self.ctc_lambda = ctc_lambda
        self.color_lambda = color_lambda
        self.final_stage = np.ceil((np.log2(data_shape[0]) - 2)).astype(np.uint8) # number of maxpooling steps to achieve tensor of size 4 or less
        self.params = tf_params.MEDIUM_PARAMS

        self.graph = tf.Graph()
        tf.reset_default_graph()
        self.initialized_vars = []
        with self.graph.as_default():
            with tf.variable_scope("DIS"):
                self.dis_transformer = tft.Transformer(self.params, train=True)
            with tf.variable_scope("GEN"):
                self.gen_transformer = tft.Transformer(self.params, train=True)

            # set i/o
            self.input_images = tf.placeholder(tf.float32, shape=(batch_size, data_shape[0], data_shape[1], 3), name='input_content')
            self.reference_images = tf.placeholder(tf.float32, shape=(batch_size, data_shape[0], data_shape[1], 3))
            self.transcriptions = tf.placeholder(tf.int32, shape=(batch_size, max_labels), name='input_transcriptions')
            self.reference_transcriptions = tf.placeholder(tf.int32, shape=(batch_size, max_labels))

            self.ctc_targets = tf.sparse_placeholder(tf.int32, name='ctc_targets')
            self.ctc_seq_len = tf.placeholder(tf.int32, [None], name='ctc_seq_len')

            self.inference_content = tf.placeholder(tf.float32, shape=(1, data_shape[0], data_shape[1], 3), name='inference_content')
            self.inference_transcriptions = tf.placeholder(tf.int32, shape=(1, max_labels), name='inference_transcriptions')
            self.inference_op = self.generator(self.inference_content, self.inference_transcriptions, train=False)
            self.inference_op = tf.identity(self.inference_op, name='inference_op')

    def build_optimizers(self, session):
        with self.graph.as_default():
            self.d_loss, self.g_loss = self.compute_losses()

            # initialize all vars and add them to optimizers
            self.dis_variables = [var for var in tf.trainable_variables() if var.name.startswith('DIS')]
            self.gen_variables = [var for var in tf.trainable_variables() if var.name.startswith('GEN')]
            self.ocr_variables = [var for var in tf.trainable_variables() if var.name.startswith('OCR')]

            self.dis_optimizer = tf.train.AdamOptimizer(learning_rate=0.0001, beta1=0.0, beta2=0.99, name='dis_Adam').minimize(self.d_loss, var_list=self.dis_variables)
            self.gen_optimizer = tf.train.AdamOptimizer(learning_rate=0.0001, beta1=0.0, beta2=0.99, name='gen_Adam').minimize(self.g_loss, var_list=self.gen_variables)
            self.ocr_optimizer = tf.train.AdamOptimizer(learning_rate=0.0003, name='ocr_Adam').minimize(self.ctc_loss, var_list=self.ocr_variables)

            session.run(tf.global_variables_initializer())
            self.saver = tf.train.Saver()

    def get_embeddings(self, tensor):
        embedding_softmax_layer = tf_embedding.EmbeddingSharedWeights(self.num_chars, self.params["hidden_size"], method="gather")
        return embedding_softmax_layer(tensor)

    def instance_norm(self, tensor):
        num_features = tensor.shape[-1]

        mean = tf.reduce_mean(tensor, axis=[1,2], keepdims=True)
        variance = tf.reduce_mean(tf.square(tensor - mean), axis=[1,2], keepdims=True)
        tensor_norm = (tensor - mean) * tf.rsqrt(variance + 1e-6)

        return tensor_norm

    def discriminator(self, tensor, transcriptions):
        with tf.variable_scope("DIS"):
            for stage in range(self.final_stage):
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None, name='first_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None, name='second_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = tf.layers.max_pooling2d(tensor, 2, 2)
            # final scoring layers
            height = int(tensor.shape[1])
            data_embeddings = tf.layers.conv2d(tensor, self.params['hidden_size'], (height, 1), padding='valid', activation=None, name='final', reuse=tf.AUTO_REUSE)
            data_embeddings = data_embeddings[:,0,:,:]

            transcription_embeddings = self.get_embeddings(transcriptions)
            tensor = self.dis_transformer(transcription_embeddings, data_embeddings, train=True)
            tensor = tf.nn.leaky_relu(tensor)
            score = tf.layers.dense(tensor, 1, name='score', reuse=tf.AUTO_REUSE)
            score = tf.reduce_mean(score, axis=1)
        return score

    def ocr(self, tensor):
        with tf.variable_scope("OCR"):
            for stage in range(2):
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None,
                                          name='first_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None,
                                          name='second_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = tf.layers.max_pooling2d(tensor, 2, 2)
            # final scoring layers
            height = int(tensor.shape[1])
            net = tf.layers.conv2d(tensor, self.params['hidden_size'], (height, 1), padding='valid',
                                       activation=None, name='final', reuse=tf.AUTO_REUSE)
            net = tf.nn.leaky_relu(net)
            net = net[:,0,:,:]
            for i in range(2):
                net = tf.layers.conv1d(
                    net, filters=self.params['hidden_size'], kernel_size=3, padding='same',
                    name='out_{:d}'.format(i), reuse=tf.AUTO_REUSE)
                net = tf.nn.leaky_relu(net)
            logits = tf.layers.conv1d(net, self.num_chars + 1, 3,
                                      padding='same', name='ocr_logits', reuse=tf.AUTO_REUSE)
            print('OCR output', logits.shape)
            return logits

    def generator(self, tensor, transcriptions, train=True):
        tensors_to_append = list()
        with tf.variable_scope("GEN"):
            # content encoder
            for stage in range(self.final_stage):
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None, name='first_down_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None, name='second_down_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensors_to_append.append(tensor)
                tensor = tf.layers.max_pooling2d(tensor, 2, 2)
            height = int(tensor.shape[1])
            width = int(tensor.shape[2])
            data_embeddings = tf.layers.conv2d(tensor, self.params['hidden_size'], (height, 1), padding='valid', activation=None, name='bottleneck', reuse=tf.AUTO_REUSE)
            data_embeddings = data_embeddings[:,0,:,:]

            # text encoder
            transcription_embeddings = self.get_embeddings(transcriptions)
            tensor = self.gen_transformer(transcription_embeddings, data_embeddings, train=train)
            tensor = tf.reshape(tensor, (-1, width, height, tensor.shape[2] // height))
            tensor = tf.transpose(tensor, (0, 2, 1, 3))

            # output decoder
            for stage in range(self.final_stage-1, -1, -1):
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None, name='first_up_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = tf.layers.conv2d(tensor, self.init_features * (2 ** stage), 3, padding='same', activation=None, name='second_up_stage_' + str(stage), reuse=tf.AUTO_REUSE)
                tensor = tf.nn.leaky_relu(tensor)
                tensor = self.instance_norm(tensor)
                tensor = tf.image.resize_images(tensor, (tensors_to_append[stage].shape[1:3]))
                tensor = tf.concat((tensors_to_append[stage], tensor), axis=-1)

            # produce output
            tensor = tf.layers.conv2d(tensor, self.init_features, 3, padding='same', activation=None,
                                      name='prefinal', reuse=tf.AUTO_REUSE)
            tensor = tf.nn.leaky_relu(tensor)
            output = tf.layers.conv2d(tensor, 3, 3, padding='same', activation=None, name='GEN_final', reuse=tf.AUTO_REUSE)
        return output

    def color_consistency_loss(self, input, output, tiles=8):
        height = int(input.shape[1])
        width = int(input.shape[2]) // tiles

        input_avgs = tf.nn.avg_pool(input, (1, height, width, 1), (1, height, width, 1), padding='SAME')
        output_avgs = tf.nn.avg_pool(output, (1, height, width, 1), (1, height, width, 1), padding='SAME')

        return tf.losses.mean_squared_error(input_avgs, output_avgs)

    def compute_losses(self):
        # restoration model output
        gen_out = self.generator(self.input_images, self.transcriptions)

        # discriminator scores assigned to restored output and reference image
        self.dis_gen_score  = self.discriminator(gen_out, self.transcriptions)
        self.dis_real_score = self.discriminator(self.reference_images, self.reference_transcriptions)

        # text recognition model logits read from restored output and reference image
        self.ctc_gen_logits = self.ocr(gen_out)
        self.ctc_gen_loss = tf.reduce_mean(tf.nn.ctc_loss(self.ctc_targets, self.ctc_gen_logits, self.ctc_seq_len,
        ctc_merge_repeated=True, time_major=False))

        # text recognition model training with CTC loss
        self.ctc_logits = self.ocr(self.input_images)
        self.ctc_gen_decoded, self.ctc_gen_log_prob = tf.nn.ctc_greedy_decoder(
            tf.transpose(self.ctc_gen_logits, (1, 0, 2)), self.ctc_seq_len, merge_repeated=True)
        self.ctc_gen_cer = tf.reduce_mean(tf.edit_distance(
            tf.cast(self.ctc_gen_decoded[0], tf.int32), self.ctc_targets))

        self.ctc_decoded, self.ctc_log_prob = tf.nn.ctc_greedy_decoder(tf.transpose(self.ctc_logits, (1, 0, 2)),
                                                                       self.ctc_seq_len, merge_repeated=True)
        self.ctc_cer = tf.reduce_mean(tf.edit_distance(tf.cast(self.ctc_decoded[0], tf.int32), self.ctc_targets))
        self.ctc_loss = tf.reduce_mean(tf.nn.ctc_loss(self.ctc_targets, self.ctc_logits, self.ctc_seq_len,
                                                      ctc_merge_repeated=True, time_major=False))

        # color consistency loss
        self.color_loss = self.color_consistency_loss(self.input_images, gen_out)

        # discriminator and restoraiton model losses for training
        d_loss = - tf.reduce_mean(self.dis_gen_score) + tf.reduce_mean(self.dis_real_score)
        g_loss = tf.reduce_mean(self.dis_gen_score) + self.ctc_lambda * self.ctc_gen_loss + self.color_lambda * self.color_loss

        # add gradient penalty to discriminator loss
        epsilon = tf.random_uniform([], 0.0, 1.0)
        mixed_input = epsilon * self.reference_images + (1 - epsilon) * gen_out
        # transcriptions serve as int32 indices for gather op so lets just mix them by sampling from input and reference transcriptions with epsilon and 1-epsilon probability, respectively
        mixed_weights = tf.random_uniform(self.transcriptions.shape, 0.0, 1.0)
        mixed_transcriptions = tf.zeros(self.transcriptions.shape)
        mixed_transcriptions = self.reference_transcriptions * tf.cast(mixed_weights<=epsilon, tf.int32)
        mixed_transcriptions += self.transcriptions * tf.cast(mixed_weights>epsilon, tf.int32)
        mixed_score = self.discriminator(mixed_input, mixed_transcriptions)
        gp = tf.gradients(mixed_score, mixed_input)[0]
        gp = tf.sqrt(tf.reduce_sum(tf.square(gp), axis=[1,2,3]))
        gp = tf.reduce_mean(tf.square(gp - 1.0) * self.scale)

        return d_loss + gp, g_loss


# coding=utf-8

import math
import os
from collections import OrderedDict
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision

class encoder(nn.Module):
    def __init__(self):
        super(encoder, self).__init__()

    def forward(self, x):
        res = None
        return res

class decoder(nn.Module):
    def __init__(self):
        super(decoder, self).__init__()

    def forward(self, x):
        res = None
        return res

class netG(nn.Module):
    def __init__(self):
        super(netG, self).__init__()

    def forward(self, x):
        res = None
        return res

class netD(nn.Module):
    def __init__(self):
        super(netD, self).__init__()

    def forward(self, x):
        res = None
        return res

def BUILD_NET():
    """
    net interface
    :return:
    """
    netG = NetG()
    netD = NetD()
    return netG, netD


#encoding: utf-8
import numpy as np
import six
from keras import backend as K
from keras.layers import Convolution2D, Dense, Embedding, GRU, Input,Concatenate, Dropout
from keras.models import Model, Sequential
from keras.applications.vgg19 import VGG19
from layers import L2Normalize
from keras.utils import plot_model
from keras import __version__ as keras_ver
from tools import encode_sentences

import pandas
#from glove import Glove

# from translate import Translator
import pdb
try:
    import cPickle as pickle
except ImportError:
    import pickle

import h5py

# def ModelwithMetadata(Model):
#     def __init__(self,metadata_json=None):
#         super(ModelwithMetadata,self).__init__(name='KerasModelwithMetadata'))
#     def save()
#         ModelwithMetadata

def concept_detector(model_file, glove_file, input_length ,data_vocab ,
                     token_count, num_classes,
                     image_only_model =False,dropout_before_final = 0.0,
                     final_act="sigmoid",
                     trainable_early_layers=False):

        #, embedding_weights, gru_weights, init_vocab_map =
    if image_only_model:
        vocab_map = None
        img_enc_weights = None
        if model_file is not None:
            img_enc_weights = load_pretrained_parameters(model_file)
        image_encoder,image_feat_extractor = build_image_encoder(weights=img_enc_weights,
                                                                 embedding_dim=1024, normalize=True,
                                                                 trainable_early_layers=trainable_early_layers)
        if dropout_before_final> 0.0 :
            image_encoder_out = Dropout(dropout_before_final) (image_encoder.outputs[0])
        else:
            image_encoder_out = image_encoder.outputs[0]
        concept_detector_scores = Dense(num_classes, activation=final_act )(image_encoder_out)
        image_encoder.compile(optimizer='nadam', loss='mse')
        end_to_end_model = Model(inputs= [image_feat_extractor.inputs[0] ],
                             outputs = concept_detector_scores)
        plot_model(end_to_end_model, to_file='model.png')
        return end_to_end_model,vocab_map

    image_encoder, sentence_encoder, vocab_map , image_feat_extractor = \
        build_pretrained_models(model_file, glove_file,input_length,
                                data_vocab = data_vocab,token_count=token_count,
                                trainable_early_layers=False)
    captioned_image_descriptor  = Concatenate(axis=-1)([image_encoder.outputs[0],
                                                        sentence_encoder.outputs[0]])
    if dropout_before_final > 0.0 :
        captioned_image_descriptor = Dropout(dropout_before_final) (captioned_image_descriptor)

    concept_detector_scores = Dense(num_classes , activation=final_act)(captioned_image_descriptor)
    for enc in image_encoder, sentence_encoder:
        enc.compile(optimizer='nadam', loss='mse')
    end_to_end_model = Model(inputs= [image_feat_extractor.inputs[0] ,
                                      sentence_encoder.inputs[0]],
                             outputs = concept_detector_scores)
    plot_model(end_to_end_model, to_file='model.png')
    return end_to_end_model,vocab_map

def build_image_feat_extractor():
    return  VGG19(include_top=False, weights='imagenet', input_tensor=None, input_shape=None, pooling="avg")

def build_image_encoder(weights=None, input_dim=4096, embedding_dim=1024, normalize=True, trainable_early_layers=False ):
    #input  = Input(shape=(input_dim,))
    init_model = build_image_feat_extractor()
    for layer in init_model.layers:
        layer.trainable = trainable_early_layers
    #init_model.compile
    x= Dense(
        embedding_dim,
        weights=weights
    )(init_model.outputs[0])
    if normalize:
        x = L2Normalize()(x)
    model = Model(inputs=[init_model.inputs[0]], outputs=[x])

    plot_model(model, to_file='image_model.png')
    return model, init_model

def build_image_encoder_seq_api(weights=None, input_dim=4096, embedding_dim=1024, normalize=True):
    model = build_image_feat_extractor()

    model.add(Dense(
        embedding_dim,
        weights=weights
    ))
    if normalize:
        model.add(L2Normalize())

    return model

def build_sentence_encoder(embedding_weights=None, gru_weights=None, input_length=None, vocab_dim=32198,
        vocab_embedding_dim=300, embedding_dim=1024, normalize=True, finetune_word_embedding=False):
    # NOTE: This gives slightly different results than the original model.
    # I think it's because the original has a different masking scheme.
    model = Sequential([
        Embedding(
            input_dim=vocab_dim, output_dim= vocab_embedding_dim, input_length=input_length,
            weights=embedding_weights, mask_zero=True,   # TODO: masking isn't quite right,
            trainable = finetune_word_embedding
        ),
        GRU(embedding_dim, weights=gru_weights, inner_activation='sigmoid'),
    ])
    if normalize:
        model.add(L2Normalize())
    return model

def build_pretrained_models(model_filename, glove_file,input_length=None,data_vocab = None,
                             token_count=None, normalize=True, trainable_early_layers =False):
    gru_weights= None
    img_enc_weights = None

    if model_filename is not None:
        img_enc_weights = load_pretrained_parameters(model_filename)
        #, embedding_weights, gru_weights, init_vocab_map =

    glove_embedding_mat = None
    #reading weights from original vse(@ryankiros) trained model (coco,flickr8k or flickr30k)
    if token_count is None:
        print("assuming original dict pkl is available for pretrained model")
        embedding_weights, gru_weights, init_vocab_map =load_pretrained_embedding_weights(model_filename)
        token_count = len(init_vocab_map)
        glove_embedding_mat= embedding_weights
    glove_embedding_mat,_ = compute_embedding_matrix(glove_file,data_vocab,None)
    vocab_embed_dim = glove_embedding_mat.shape[1]

    image_encoder,image_feat_extractor = build_image_encoder(weights=img_enc_weights,
                                                             embedding_dim=1024, normalize=normalize,
                                                             trainable_early_layers= trainable_early_layers)
    print ("Word Embedding matrix shape")
    print(glove_embedding_mat.shape)
    #print(glove_embedding_mat[10,:])
    sentence_encoder = build_sentence_encoder(
        embedding_weights=[glove_embedding_mat],
        gru_weights=gru_weights,
        input_length=input_length, vocab_dim=token_count+1,
        vocab_embedding_dim=vocab_embed_dim,
        normalize=normalize)
    return image_encoder, sentence_encoder, data_vocab , image_feat_extractor

def load_pretrained_parameters(filename):
    '''Load up the pre-trained weights from the @ryankiros implementation.
    '''
    params = np.load(filename)

    # image encoder weights
    if params:
        img_enc_weights = [params['ff_image_W'], params['ff_image_b']]
    else:
        img_enc_weights = None
    return img_enc_weights

def load_pretrained_embedding_weights(filename):
    params = np.load(filename)
    vocab_map = np.load('{}.dictionary.pkl'.format(filename))
    # sentence encoder weights
    embedding_weights = [params['Wemb']]
    W_h = params['encoder_Wx']
    U_h = params['encoder_Ux']
    b_h = params['encoder_bx']
    W_r, W_z = np.split(params['encoder_W'], 2, axis=1)
    U_r, U_z = np.split(params['encoder_U'], 2, axis=1)
    b_r, b_z = np.split(params['encoder_b'], 2)
    gru_weights = [
        W_z, U_z, b_z,
        W_r, U_r, b_r,
        W_h, U_h, b_h,
    ]
    return embedding_weights, gru_weights, vocab_map

def compute_embedding_matrix(glove_file,word_index,vocab_embed_dim=None):
    embeddings_index = dict()
    vocab_embed_dim = 0
    avg_embed_vec = None
    if glove_file.endswith('pkl') :
        num_embeds = 0
        avg_embed_vec = None
        embeddings_index=pickle.load(open(glove_file,'rb'))
        #print (embeddings_index.keys())
        for _,avec in embeddings_index.items():
            vocab_embed_dim  = avec.size
            if avg_embed_vec is None:
                avg_embed_vec = avec.copy()
            else:
                avg_embed_vec += avec
            num_embeds +=1
        avg_embed_vec = avg_embed_vec/num_embeds
    elif glove_file.endswith('h5') :
        glove_hdf5 = h5py.File(glove_file,'r')
        embeddings_index = glove_hdf5["embed_dict"]
        avg_embed_vec = embeddings_index["avg_word_vector"].copy()
        vocab_embed_dim = avg_embed_vec.size
        del embeddings_index["avg_word_vector"]

    else:
        f= open(glove_file,'r')
        avg_embed_vec = None
        num_embeds = 0
        for line in f:
            values = line.split(' ')
            word = values[0]
            try:
                coefs = np.asarray(values[1:], dtype='float32')
            except Exception as e:
                print (e)
                print (values[0])
                print (values[1])
                print (values[-1])
                print (line)
            vocab_embed_dim = len(values[1:])
            embeddings_index[word] = coefs
            if avg_embed_vec is None:
                avg_embed_vec = coefs.copy()

            else:
                avg_embed_vec += coefs
            num_embeds +=1
        avg_embed_vec = avg_embed_vec/num_embeds

        f.close()
        f= open(glove_file+'.pkl','wb')
        pickle.dump(embeddings_index, f)
        f.close()
        with h5py.File(glove_file+'.h5','w') as f:
            embed_dict=f.create_group("embed_dict")
            for word, i in word_index.items():
                embed_dict[word] = coefs
            embed_dict["avg_word_vector"] = avg_embed_vec
    embedding_matrix = np.zeros((len(word_index) + 1, vocab_embed_dim))
    fh = open("oov.txt",'w', encoding="utf8")
    for word, i in word_index.items():
        embedding_vector = embeddings_index.get(word)
        if embedding_vector is not None:
            # words not found in embedding index will be all-zeros.
            embedding_matrix[i,:] = embedding_vector
        else:
            #translator= Translator(to_lang="english")
            #translation = translator.translate(word)
            #pdb.set_trace()
            #print( u' %s not found in glove embedding. translating word to english in case foreign word' % word)

            fh.write("%s\n" % word)

            embedding_matrix[i,:] = avg_embed_vec

            #print (translation)
            #pdb.set_trace()
            # embedding_vector = embeddings_index.get(translation)
            # if embedding_vector is not None:
            #     embedding_matrix[i] = embedding_vector
    embedding_matrix[-1,:] =avg_embed_vec
    fh.close()
    return embedding_matrix,vocab_embed_dim


# -*- coding: utf-8 -*-
"""
Created on Fri Mar 31 11:39:31 2017

@author: 12102083
"""

import os
import numpy
import pandas
from sklearn.utils import shuffle
from keras.wrappers.scikit_learn import KerasRegressor
from sklearn.model_selection import cross_val_score
from sklearn.model_selection import KFold
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import time
from sklearn.pipeline import Pipeline

from keras.models import Sequential
from keras.layers import Dense, Activation
from keras.layers import Convolution2D, MaxPooling2D, Flatten
from keras.optimizers import SGD
from keras.models import model_from_json
from keras import backend as K
from keras.models import model_from_config
from keras.utils.np_utils import convert_kernel
import tensorflow as tf
#from tensorflow_serving.session_bundle import exporter
# In[ ]:
model = model_from_json(open(r'pathToKeraModelStructure',).read())
model.load_weights(r'pathToKerasModelWeights')
output_tensor = model.output
tensor = k.get_
tf.train.write_graph(output_tensor, 'TensorGraphs', 'train.pb')

# In[ ]:

import keras
import tensorflow
from keras import backend as K
from tensorflow.contrib.session_bundle import exporter
from keras.models import model_from_config, Sequential
from tensorflow.python.framework.graph_util import convert_variables_to_constants

print("Loading model for exporting to Protocol Buffer format...")
model_path = 'temp.json'
model = model_from_json(open('temp.json',).read())
model.load_weights('temp.h5')

K.set_learning_phase(0)  # all new operations will be in test mode from now on
ops = []
for layer in model.layers:
   print("layer name = " + layer.name)
   if layer.__class__.__name__ in ['Convolution1D', 'Convolution2D', 'Convolution3D', 'AtrousConvolution2D']:
      original_w = K.get_value(layer.W)
      converted_w = convert_kernel(original_w)
      ops.append(tf.assign(layer.W, converted_w).op)
      print("Hi")

K.get_session().run(ops)

sess = K.get_session()

export_path = "TensorGraphs\\simple.pb"  # where to save the exported graph
export_version = 1  # version number (integer)
tf.global_variables_initializer()#initialize_all_variables();
minimal_graph = convert_variables_to_constants(sess, sess.graph_def, ["dense_3_b/Assign"]) # output
tf.train.write_graph(sess.graph_def, './tmp/beginner-export','beginner-graph.pb', as_text=False)

# In[ ]:

import keras
import tensorflow
from keras import backend as K
from tensorflow.contrib.session_bundle import exporter
from keras.models import model_from_config, Sequential

print("Loading model for exporting to Protocol Buffer format...")
model_path = 'temp.h5'
model = keras.models.load_model(model_path)

K.set_learning_phase(0)  # all new operations will be in test mode from now on
sess = K.get_session()

# serialize the model and get its weights, for quick re-building
config = model.get_config()
weights = model.get_weights()

# re-build a model where the learning phase is now hard-coded to 0
new_model = Sequential.from_config(config)
new_model.set_weights(weights)

export_path = "TensorGraphs//simple.pb"  # where to save the exported graph
export_version = 1  # version number (integer)

saver = tensorflow.train.Saver(sharded=True)
model_exporter = exporter.Exporter(saver)
signature = exporter.classification_signature(input_tensor=model.input, scores_tensor=model.output)
model_exporter.init(sess.graph.as_graph_def(), default_graph_signature=signature)
model_exporter.export(export_path, tensorflow.constant(export_version), sess)


import pandas as pd
import numpy as np
from itertools import product
import math

class UserPoints:
    def __init__(self):
        self.survey_weight = [5, 10, 10, 15, 10, 20, 20, 10]

    def a100_answer(self, answer):
        survey_weight = self.survey_weight[0]
        comfort_ls = [2, 2, 2, 3, 5]
        clean_ls = [2, 2, 2, 3, 5]
        comfort_point = comfort_ls[answer-1] * survey_weight
        clean_point = clean_ls[answer - 1] * survey_weight
        return [comfort_point, 0, 0, 0, clean_point]

    def a200_answer(self, answer):
        survey_weight = self.survey_weight[1]
        together_ls = [5, 3, 3, 0]
        together_point = together_ls[answer-1] * survey_weight
        return [0, together_point, 0, 0, 0]

    def a210_answer(self, answer):
        survey_weight = self.survey_weight[2]
        together_ls = [5, 0]
        together_point = together_ls[answer-1] * survey_weight
        return [0, together_point, 0, 0, 0]

    def a300_answer(self, answer):
        survey_weight = self.survey_weight[3]
        comfort_ls = [5, 10, 0]
        fun_ls = [3, 1, 5]
        clean_ls = [5, 10, 0]
        comfort_point = comfort_ls[answer-1] * survey_weight
        fun_point = fun_ls[answer - 1] * survey_weight
        clean_point = clean_ls[answer - 1] * survey_weight
        return [comfort_point, 0, fun_point, 0, clean_point]

    def a410_answer(self, answer):
        survey_weight = self.survey_weight[4]
        together_ls = [3, 5, 3, 3]
        fun_ls = [5, 10, 5, 3]
        together_point = together_ls[answer-1] * survey_weight
        fun_point = fun_ls[answer - 1] * survey_weight
        return [0, together_point, fun_point, 0, 0]

    def a420_answer(self, answer):
        survey_weight = self.survey_weight[5]
        comfort_ls = [2, 2, 2, 5]
        fun_ls = [5, 5, 1, 1]
        healing_ls = [5, 5, 5, 2]
        clean_ls = [2, 2, 2, 5]
        comfort_point = comfort_ls[answer - 1] * survey_weight
        fun_point = fun_ls[answer - 1] * survey_weight
        healing_point = healing_ls[answer - 1] * survey_weight
        clean_point = clean_ls[answer - 1] * survey_weight
        return [comfort_point, 0, fun_point, healing_point, clean_point]

    def a500_answer(self, answer):
        survey_weight = self.survey_weight[6]
        together_ls = [5, 1, 3, 3]
        fun_ls = [5, 1, 3, 2]
        healing_ls = [2, 5, 3, 1]
        together_point = together_ls[answer - 1] * survey_weight
        fun_point = fun_ls[answer - 1] * survey_weight
        healing_point = healing_ls[answer - 1] * survey_weight
        return [0, together_point, fun_point, healing_point, 0]

    def a600_answer(self, answer):
        survey_weight = self.survey_weight[7]
        comfort_ls = [5, 1, 1, 1, 1, 1]
        fun_ls = [1, 5, 5, 5, 5, 5]
        healing_ls = [1, 5, 5, 5, 5, 5]
        comfort_point = comfort_ls[answer - 1] * survey_weight
        fun_point = fun_ls[answer - 1] * survey_weight
        healing_point = healing_ls[answer - 1] * survey_weight
        return [comfort_point, 0, fun_point, healing_point, 0]

    def calc_cat_point(self, answer_ls):
        ls100 = self.a100_answer(answer_ls[0])
        ls200 = self.a200_answer(answer_ls[1])
        ls210 = self.a210_answer(answer_ls[2])
        ls300 = self.a300_answer(answer_ls[3])
        ls410 = self.a410_answer(answer_ls[4])
        ls420 = self.a420_answer(answer_ls[5])
        ls500 = self.a500_answer(answer_ls[6])
        ls600 = self.a600_answer(answer_ls[7])
        result_ls = [ls100[i] + ls200[i] + ls210[i] + ls300[i] + ls410[i] + ls420[i] + ls500[i] + ls600[i] for i in range(len(ls100))]
        return result_ls

    def calc_final_point(self, user_answer_ls):
        result_ls = self.calc_cat_point(user_answer_ls)
        questions = [[i for i in range(1, 6)],
                     [i for i in range(1, 5)],
                     [i for i in range(1, 3)],
                     [i for i in range(1, 4)],
                     [i for i in range(1, 5)],
                     [i for i in range(1, 5)],
                     [i for i in range(1, 5)],
                     [i for i in range(1, 7)]]
        answer_ls = list(product(*questions))
        comfort_ls, together_ls, fun_ls, healing_ls, clean_ls = [], [], [], [], []
        for answer in answer_ls:
            result = self.calc_cat_point(answer)
            comfort_ls.append(result[0])
            together_ls.append(result[1])
            fun_ls.append(result[2])
            healing_ls.append(result[3])
            clean_ls.append(result[4])

        comfort_point = result_ls[0] / max(comfort_ls)*100
        together_point = result_ls[1] / max(together_ls)*100
        fun_point = result_ls[2] / max(fun_ls)*100
        healing_point = result_ls[3] / max(healing_ls)*100
        clean_point = result_ls[4] / max(clean_ls)*100

        return [comfort_point, together_point, fun_point, healing_point, clean_point]

class PolarArea(UserPoints):
    def __init__(self):
        super().__init__()

    def calc_polar_area(self, result_ls):
        # result_ls = self.calc_cat_point(answer_ls)
        comfort_together = result_ls[0] * result_ls[1] * math.sin(360/5) / 2
        together_fun = result_ls[1] * result_ls[2] * math.sin(360 / 5) / 2
        fun_healing = result_ls[2] * result_ls[3] * math.sin(360 / 5) / 2
        healing_clean = result_ls[3] * result_ls[4] * math.sin(360 / 5) / 2
        clean_comfort = result_ls[4] * result_ls[0] * math.sin(360 / 5) / 2
        area = comfort_together + together_fun + fun_healing + healing_clean + clean_comfort
        return area

    def matching_pct(self, camp_point_ls, user_point_ls):
        pct_ls = []
        for i in range(len(camp_point_ls)):

            pct = camp_point_ls[i] / user_point_ls[i]
            if pct >= 1:
                pct_ls.append(1/pct)
            elif pct == 0:
                pass
            else:
                pct_ls.append(pct)
        return np.round(np.average(pct_ls) * 100, 1)


# --------------------------------------------------------
# Tensorflow Faster R-CNN
# Licensed under The MIT License [see LICENSE for details]
# Written by Jiasen Lu, Jianwei Yang, based on code from Ross Girshick

# Modified by Peiliang Li for Stereo RCNN test
# --------------------------------------------------------
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths
import os
import sys
import numpy as np
import argparse
import shutil
import time
import cv2
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import torchvision
import math as m
from roi_data_layer.roidb import combined_roidb
from roi_data_layer.roibatchLoader import roibatchLoader
from model.utils.config import cfg
from model.rpn.bbox_transform import clip_boxes
# from model.roi_layers import nms
from model.rpn.bbox_transform import bbox_transform_inv, kpts_transform_inv, border_transform_inv
from model.utils.net_utils import save_net, load_net, vis_detections
from model.stereo_rcnn.resnet import resnet
from model.utils import kitti_utils
from model.utils import vis_3d_utils as vis_utils
from model.utils import box_estimator as box_estimator
from model.dense_align import dense_align

try:
    xrange          # Python 2
except NameError:
    xrange = range  # Python 3

def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Test the Stereo R-CNN network')

    parser.add_argument('--load_dir', dest='load_dir',
                        help='directory to load models', default="models_stereo",
                        type=str)
    parser.add_argument('--checkepoch', dest='checkepoch',
                        help='checkepoch to load network',
                        default=12, type=int)
    parser.add_argument('--checkpoint', dest='checkpoint',
                        help='checkpoint to load network',
                        default=6477, type=int)

    args = parser.parse_args()
    return args

if __name__ == '__main__':

    args = parse_args()

    np.random.seed(cfg.RNG_SEED)

    cfg.TRAIN.USE_FLIPPED = False
    imdb, roidb, ratio_list, ratio_index = combined_roidb('kitti_val', False)
    print('{:d} roidb entries'.format(len(roidb)))

    # input_dir = args.load_dir + "/"
    # if not os.path.exists(input_dir):
    #     raise Exception('There is no input directory for loading network from ' + input_dir)
    #  load_name = os.path.join(input_dir, 'stereo_rcnn_{}_{}.pth'.format(args.checkepoch, args.checkpoint))
    load_name = '/home/kkb/users/yqh/Stereo-RCNN/models_stereo/stereo_rcnn_12_6477.pth'

    result_dir = args.load_dir + '/result/'
    if os.path.exists(result_dir):
        shutil.rmtree(result_dir)

    # initilize the network here.
    stereoRCNN = resnet(imdb.classes, 101, pretrained=False)
    stereoRCNN.create_architecture()

    print("load checkpoint %s" % (load_name))
    checkpoint = torch.load(load_name)
    stereoRCNN.load_state_dict(checkpoint['model'])
    print('load model successfully!')

    with torch.no_grad():
        # initilize the tensor holder here.
        im_left_data = Variable(torch.FloatTensor(1).cuda())
        im_right_data = Variable(torch.FloatTensor(1).cuda())
        im_info = Variable(torch.FloatTensor(1).cuda())
        num_boxes = Variable(torch.LongTensor(1).cuda())
        gt_boxes_left = Variable(torch.FloatTensor(1).cuda())
        gt_boxes_right = Variable(torch.FloatTensor(1).cuda())
        gt_boxes_merge = Variable(torch.FloatTensor(1).cuda())
        gt_dim_orien = Variable(torch.FloatTensor(1).cuda())
        gt_kpts = Variable(torch.FloatTensor(1).cuda())

        stereoRCNN.cuda()

        eval_thresh = 0.05
        vis_thresh = 0.7

        num_images = len(imdb.image_index)

        dataset = roibatchLoader(roidb, ratio_list, ratio_index, 1, \
                            imdb.num_classes, training=False, normalize = False)
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=1,
                                shuffle=False, num_workers=0,
                                pin_memory=True)

        data_iter = iter(dataloader)
        data = next(data_iter)
        data = next(data_iter)
        for i in data:
            print(i)


# Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://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.
"""A demo script showing how to use the uisrnn package on toy data."""

import numpy as np

import uisrnn

import tensorflow as tf

if tf.test.gpu_device_name():
  print('Default GPU Device:{}'.format(tf.test.gpu_device_name()))
else:
  print("Please install GPU version of TF")

SAVED_MODEL_NAME = '/scratch/hh2263/VCTK/saved_model_vctk_muda_updated_transfer.uisrnn_benchmark'
original_model = '/home/hh2263/Speaker-Diarization/pretrained/saved_model.uisrnn_benchmark'

def diarization_experiment(model_args, training_args, inference_args):
  """Experiment pipeline.

  Load data --> train model --> test model --> output result

  Args:
    model_args: model configurations
    training_args: training configurations
    inference_args: inference configurations
  """

  predicted_labels = []
  test_record = []

  np_load_old = np.load
  # modify the default parameters of np.load
  np.load = lambda *a,**k: np_load_old(*a, allow_pickle=True, **k)

  train_data = np.load('/scratch/hh2263/VCTK/training_data_vctk_muda_updated.npz')

  # restore np.load for future normal usage
  np.load = np_load_old

  train_sequence = train_data['train_sequence']
  train_cluster_id = train_data['train_cluster_id']
  train_sequence_list = [seq.astype(float)+0.00001 for seq in train_sequence]
  train_cluster_id_list = [np.array(cid).astype(str) for cid in train_cluster_id]

  model = uisrnn.UISRNN(model_args)
  model.load(original_model)

  # training
  model.fit(train_sequence_list, train_cluster_id_list, training_args)
  model.save(SAVED_MODEL_NAME)

  '''
  # testing
  # we can also skip training by calling：
  model.load(SAVED_MODEL_NAME)
  for (test_sequence, test_cluster_id) in zip(test_sequences, test_cluster_ids):
    predicted_label = model.predict(test_sequence, inference_args)
    predicted_labels.append(predicted_label)
    accuracy = uisrnn.compute_sequence_match_accuracy(
        test_cluster_id, predicted_label)
    test_record.append((accuracy, len(test_cluster_id)))
    print('Ground truth labels:')
    print(test_cluster_id)
    print('Predicted labels:')
    print(predicted_label)
    print('-' * 80)

  output_string = uisrnn.output_result(model_args, training_args, test_record)

  print('Finished diarization experiment')
  print(output_string)
  '''

def main():
  """The main function."""
  model_args, training_args, inference_args = uisrnn.parse_arguments()
  model_args.observation_dim = 512
  model_args.rnn_depth = 1
  model_args.rnn_hidden_size = 512
  training_args.enforce_cluster_id_uniqueness = False
  training_args.batch_size = 15#30
  training_args.learning_rate = 1e-4
  training_args.train_iteration = 3000
  training_args.num_permutations = 20
  # training_args.grad_max_norm = 5.0
  training_args.learning_rate_half_life = 1000
  diarization_experiment(model_args, training_args, inference_args)

if __name__ == '__main__':
  main()