text:
import numpy as np
from quantum_optimal_control.helper_functions.grape_functions import c_to_r_mat
from quantum_optimal_control.helper_functions.grape_functions import c_to_r_vec
from quantum_optimal_control.helper_functions.grape_functions import get_state_index
import scipy.linalg as la
from scipy.special import factorial

from quantum_optimal_control.helper_functions.data_management import H5File

class SystemParameters:

    def __init__(self,H0,Hops,Hnames,U,U0,total_time,steps,states_concerned_list,dressed_info,maxA, draw,initial_guess, show_plots,Unitary_error,state_transfer,no_scaling,reg_coeffs, save, file_path, Taylor_terms,use_gpu,use_inter_vecs,sparse_H,
                sparse_U,sparse_K):
        # Input variable
        self.sparse_U = sparse_U
        self.sparse_H = sparse_H
        self.sparse_K = sparse_K
        self.use_inter_vecs = use_inter_vecs
        self.use_gpu = use_gpu
        self.Taylor_terms = Taylor_terms
        self.dressed_info = dressed_info
        self.reg_coeffs = reg_coeffs
        self.file_path = file_path
        self.state_transfer = state_transfer
        self.no_scaling = no_scaling
        self.save = save
        self.H0_c = H0
        self.ops_c = Hops
        self.ops_max_amp = maxA
        self.Hnames = Hnames
        self.Hnames_original = Hnames #because we might rearrange them later if we have different timescales
        self.total_time = total_time
        self.steps = steps
        self.show_plots = show_plots
        self.Unitary_error= Unitary_error

        if initial_guess is not None:
            # transform initial_guess to its corresponding base value
            self.u0 = initial_guess
            self.u0_base = np.zeros_like(self.u0)
            for ii in range (len(self.u0_base)):
                self.u0_base[ii]= self.u0[ii]/self.ops_max_amp[ii]
                if max(self.u0_base[ii])> 1.0:
                    raise ValueError('Initial guess has strength > max_amp for op %d' % (ii) )
            self.u0_base = np.arcsin(self.u0_base) #because we take the sin of weights later

        else:
            self.u0 =[]
        self.states_concerned_list = states_concerned_list

        self.is_dressed = False
        self.U0_c = U0
        self.initial_unitary = c_to_r_mat(U0) #CtoRMat is converting complex matrices to their equivalent real (double the size) matrices
        if self.state_transfer == False:
            self.target_unitary = c_to_r_mat(U)
        else:
            self.target_vectors=[]

            for target_vector_c in U:
                self.target_vector = c_to_r_vec(target_vector_c)
                self.target_vectors.append(self.target_vector)

        if draw is not None:
            self.draw_list = draw[0]
            self.draw_names = draw[1]
        else:
            self.draw_list = []
            self.draw_names = []

        if dressed_info !=None:
            self.v_c = dressed_info['eigenvectors']
            self.dressed_id = dressed_info['dressed_id']
            self.w_c = dressed_info['eigenvalues']
            self.is_dressed = dressed_info['is_dressed']
            self.H0_diag=np.diag(self.w_c)

        self.init_system()
        self.init_vectors()
        self.init_operators()
        self.init_one_minus_gaussian_envelope()
        self.init_guess()

    def approx_expm(self,M,exp_t, scaling_terms):
        #approximate the exp at the beginning to estimate the number of taylor terms and scaling and squaring needed
        U=np.identity(len(M),dtype=M.dtype)
        Mt=np.identity(len(M),dtype=M.dtype)
        factorial=1.0 #for factorials

        for ii in xrange(1,exp_t):
            factorial*=ii
            Mt=np.dot(Mt,M)
            U+=Mt/((2.**float(ii*scaling_terms))*factorial) #scaling by 2**scaling_terms

        for ii in xrange(scaling_terms):
            U=np.dot(U,U) #squaring scaling times

        return U

    def approx_exp(self,M,exp_t, scaling_terms):
        # the scaling and squaring of matrix exponential with taylor expansions
        U=1.0
        Mt=1.0
        factorial=1.0 #for factorials

        for ii in xrange(1,exp_t):
            factorial*=ii
            Mt=M*Mt
            U+=Mt/((2.**float(ii*scaling_terms))*factorial) #scaling by 2**scaling_terms

        for ii in xrange(scaling_terms):
            U=np.dot(U,U) #squaring scaling times

        return U

    def Choose_exp_terms(self, d):
        #given our hamiltonians and a number of scaling/squaring, we determine the number of Taylor terms

        exp_t = 20 #maximum

        H=self.H0_c
        U_f = self.U0_c
        for ii in range (len(self.ops_c)):
            H = H + self.ops_max_amp[ii]*self.ops_c[ii]
        if d == 0:
            self.scaling = max(int(2*np.log2(np.max(np.abs(-(0+1j) * self.dt*H)))),0)

        else:
            self.scaling += d

        if self.state_transfer or self.no_scaling:
            self.scaling =0
        while True:

            if len(self.H0_c) < 10:
                for ii in range (self.steps):
                    U_f = np.dot(U_f,self.approx_expm((0-1j)*self.dt*H, exp_t, self.scaling))
                Metric = np.abs(np.trace(np.dot(np.conjugate(np.transpose(U_f)), U_f)))/(self.state_num)
            else:
                max_term = np.max(np.abs(-(0+1j) * self.dt*H))

                Metric = 1 + self.steps *np.abs((self.approx_exp(max_term, exp_t, self.scaling) - np.exp(max_term))/np.exp(max_term))

            if exp_t == 3:
                break
            if np.abs(Metric - 1.0) < self.Unitary_error:
                exp_t = exp_t-1
            else:
                break

        return exp_t

    def init_system(self):
        self.dt = float(self.total_time)/self.steps
        self.state_num= len(self.H0_c)

    def init_vectors(self):
        # initialized vectors used for propagation
        self.initial_vectors=[]
        self.initial_vectors_c=[]

        for state in self.states_concerned_list:

            if self.state_transfer:
                self.initial_vector_c = np.array(state)
            else:
                if self.is_dressed:
                    self.initial_vector_c= self.v_c[:,get_state_index(state,self.dressed_id)]
                else:
                    self.initial_vector_c=np.zeros(self.state_num)
                    self.initial_vector_c[state]=1

            self.initial_vectors_c.append(self.initial_vector_c)
            self.initial_vector = c_to_r_vec(self.initial_vector_c)

            self.initial_vectors.append(self.initial_vector)

        if self.save:
            with H5File(self.file_path) as hf:
                hf.add('initial_vectors_c',data=np.array(self.initial_vectors_c))

    def init_operators(self):
        # Create operator matrix in numpy array

        self.ops=[]
        for op_c in self.ops_c:
            op = c_to_r_mat(-1j*self.dt*op_c)
            self.ops.append(op)

        self.ops_len = len(self.ops)

        self.H0 = c_to_r_mat(-1j*self.dt*self.H0_c)
        self.identity_c = np.identity(self.state_num)
        self.identity = c_to_r_mat(self.identity_c)

        if self.Taylor_terms is None:
            self.exps =[]
            self.scalings = []
            if self.state_transfer or self.no_scaling:
                comparisons = 1
            else:
                comparisons = 6
            d = 0
            while comparisons >0:

                self.exp_terms = self.Choose_exp_terms(d)
                self.exps.append(self.exp_terms)
                self.scalings.append(self.scaling)
                comparisons = comparisons -1
                d = d+1
            self.complexities = np.add(self.exps,self.scalings)
            a = np.argmin(self.complexities)

            self.exp_terms = self.exps[a]
            self.scaling = self.scalings[a]
        else:
            self.exp_terms = self.Taylor_terms[0]
            self.scaling = self.Taylor_terms[1]

        if self.save:
            with H5File(self.file_path) as hf:
                hf.add('taylor_terms',data=self.exp_terms)
                hf.add('taylor_scaling',data=self.scaling)

        print( "Using "+ str(self.exp_terms) + " Taylor terms and "+ str(self.scaling)+" Scaling & Squaring terms")

        i_array = np.eye(2*self.state_num)
        op_matrix_I=i_array.tolist()

        self.H_ops = []
        for op in self.ops:
            self.H_ops.append(op)
        self.matrix_list = [self.H0]
        for ii in range(self.ops_len):
            self.matrix_list = self.matrix_list + [self.H_ops[ii]]
        self.matrix_list = self.matrix_list + [op_matrix_I]

        self.matrix_list = np.array(self.matrix_list)

    def init_one_minus_gaussian_envelope(self):
        # Generating the Gaussian envelope that pulses should obey
        one_minus_gauss = []
        offset = 0.0
        overall_offset = 0.01
        opsnum=self.ops_len
        for ii in range(opsnum):
            constraint_shape = np.ones(self.steps)- self.gaussian(np.linspace(-2,2,self.steps)) - offset
            constraint_shape = constraint_shape * (constraint_shape>0)
            constraint_shape = constraint_shape + overall_offset* np.ones(self.steps)
            one_minus_gauss.append(constraint_shape)

        self.one_minus_gauss = np.array(one_minus_gauss)

    def gaussian(self,x, mu = 0. , sig = 1. ):
        return np.exp(-np.power(x - mu, 2.) / (2 * np.power(sig, 2.)))

    def init_guess(self):
        # initail guess for control field
        if self.u0 != []:

            self.ops_weight_base = np.reshape(self.u0_base, [self.ops_len,self.steps])
        else:
            initial_mean = 0
            index = 0

            initial_stddev = (1./np.sqrt(self.steps))
            self.ops_weight_base = np.random.normal(initial_mean, initial_stddev, [self.ops_len ,self.steps])

        self.raw_shape = np.shape(self.ops_weight_base)

max_stars_repo_name: rayonde/yarn

max_issues_repo_name: rayonde/yarn

max_forks_repo_name: rayonde/yarn

reason: import numpy,import scipy,from scipy


text:
# -*- coding: utf-8 -*-
import cv2
import numpy as np

print('HDR Processing...')

# Loading exposure images into a list
img_fn = ["IMX378_3968x3008_HDR_L.ppm", "IMX378_3968x3008_HDR_S.ppm"]
img_list = [cv2.imread(fn) for fn in img_fn]

# Exposure fusion using Mertens
merge_mertens = cv2.createMergeMertens()
res_mertens = merge_mertens.process(img_list)

# Convert datatype to 8-bit and save
res_mertens_8bit = np.clip(res_mertens*255, 0, 255).astype('uint8')
cv2.imwrite("IMX378_3968x3008_HDR_fusion.ppm", res_mertens_8bit)

ext: py

max_stars_repo_name: centuryarks/Sample

max_issues_repo_name: centuryarks/Sample

max_forks_repo_name: centuryarks/Sample

reason: import numpy


text:
import numpy as np
import cv2
import os
import torch
from torch.nn import functional as F

class TTAOp:
    def __init__(self, mode=None):
        self.mode = mode

    def __call__(self, model, batch):
        with torch.no_grad():
            forwarded = torch.from_numpy(self.forward(batch.numpy())).cuda()
            return self.backward(self.to_numpy(model.forward(forwarded)))

    def forward(self, img):
        raise NotImplementedError

    def backward(self, img):
        raise NotImplementedError

    def to_numpy(self, batch):
        if isinstance(batch, list):
            batch = batch[0]

        if self.mode is None:
            pass
        elif self.mode.lower() == 'sigmoid':
            batch = torch.sigmoid(batch)
        elif self.mode.lower() == 'softmax':
            batch = torch.softmax(batch, dim=1)
        else:
            raise TypeError('Mode must be either None, "sigmoid" or "softmax".')

        data = batch.data.cpu().numpy()
        return data

class BasicTTAOp(TTAOp):
    @staticmethod
    def op(img):
        raise NotImplementedError

    def forward(self, img):
        return self.op(img)

    def backward(self, img):
        return img

class Nothing(BasicTTAOp):
    @staticmethod
    def op(img):
        return img

class HFlip(BasicTTAOp):
    @staticmethod
    def op(img):
        return np.ascontiguousarray(np.flip(img, axis=2))

class VFlip(BasicTTAOp):
    @staticmethod
    def op(img):
        return np.ascontiguousarray(np.flip(img, axis=3))

class Transpose(BasicTTAOp):
    @staticmethod
    def op(img):
        return np.ascontiguousarray(img.transpose(0, 1, 3, 2))

class ScaleTTA(TTAOp):
    def __init__(self, scale, **kwargs):
        super(ScaleTTA, self).__init__(**kwargs)
        self.scale = scale

    def __call__(self, model, batch):
        with torch.no_grad():
            self.im_size = batch.shape[-2:]
            forwarded = self.forward(batch).cuda()
            return self.to_numpy([self.backward(model.forward(forwarded)[0])])

    @staticmethod
    def op(img, scale):
        return F.upsample(img, scale_factor=scale,
                          mode='bilinear', align_corners=True)

    def forward(self, img):
        return self.op(img, self.scale)

    def backward(self, img):
        return F.upsample(img, size=self.im_size, mode='bilinear', align_corners=True)

class ResizeTTA(TTAOp):
    def __init__(self, tgt_size, pad, **kwargs):
        super(ResizeTTA, self).__init__(**kwargs)
        self.tgt_size = tgt_size
        self.pad = pad

    def __call__(self, model, batch):
        with torch.no_grad():
            self.im_size = batch.shape[-2:]
            forwarded = self.forward(batch).cuda()
            return self.to_numpy(self.backward(model.forward(forwarded)))

    @staticmethod
    def op(img, tgt_size, pad):
        if pad:
            H, W = img.shape[2:]
            _max = max(H, W)
            padH = _max - H
            padW = _max - W
            img = np.pad(img,
                         [(0, 0), (0, 0),
                          (padH // 2, padH // 2),
                          (padW // 2, padW // 2)],
                         mode='constant')

        return F.upsample(img, size=(tgt_size, tgt_size),
                          mode='bilinear', align_corners=True)

    def forward(self, img):
        return self.op(img, self.tgt_size, self.pad)

    def backward(self, img):
        return img

def chain_op(data, operations):
    for op in operations:
        data = op.op(data)
    return data

class ChainedTTA(TTAOp):
    @property
    def operations(self):
        raise NotImplementedError

    def forward(self, img):
        return chain_op(img, self.operations)

    def backward(self, img):
        return chain_op(img, reversed(self.operations))

class HVFlip(ChainedTTA):
    @property
    def operations(self):
        return [HFlip, VFlip]

class TransposeHFlip(ChainedTTA):
    @property
    def operations(self):
        return [Transpose, HFlip]

class TransposeVFlip(ChainedTTA):
    @property
    def operations(self):
        return [Transpose, VFlip]

class TransposeHVFlip(ChainedTTA):
    @property
    def operations(self):
        return [Transpose, HFlip, VFlip]

# TODO: Fix and test ScaleVFlip TTA
class ScaleVFlip(ChainedTTA):
    def __init__(self, scale, **kwargs):
        super(ScaleVFlip, self).__init__(**kwargs)
        self.scale = scale

    @property
    def operations(self):
        return [VFlip, ScaleTTA]

if __name__ == '__main__':
    # transforms = [Nothing, HFlip, VFlip, Transpose, HVFlip, TransposeHFlip, TransposeVFlip, TransposeHVFlip, DoubleResize]
    transforms = [DoubleResize]

    class Model(object):
        def predict(self, x):
            return x

    root = './Data/Train/images/'
    imgs = os.listdir(root)[:2]
    imgs = [cv2.imread(os.path.join(root, im)) / 255. for im in imgs]
    data = torch.from_numpy(np.moveaxis(np.stack((imgs)).astype(np.float32), -1, 1))
    model = Model()
    for cls in transforms:
        TTA = cls(mode=None)
        ret = TTA(model, data)
    assert np.allclose(ret, data)


text:
#!usr/bin/env python
# -*- coding: utf-8 -*-

# Usage: python [files]
#
# Image Warper:
#       Extract SIFT features
#       Match features based on Lowe's ratio test
#       Warp image using Perspective Transformation
#
# Author: yatbear <sapphirejyt@gmail.com>
#         2016-03-14

import cv2
import numpy as np
from collections import OrderedDict

class Warper(object):

    def __init__(self, rtthres=0.7, min_match_count=4):
        '''
        Args:
            rtthres: threshold for Lowe's ratio test
            min_match_count: minimun number of matches
        '''
        self.RATIO_TEST_THRES = rtthres
        self.MIN_MATCH_COUNT = min_match_count

    def get_gray_img(self, img_path):
        ''' Read a gray image from the given path. '''
        gray = cv2.imread(img_path, 0)
        return gray

    def resize_img(self, src, dst):
        '''
        Resize the source image to the same size as the destination image.
        Args:
            src: 2-D matrix
            dst: 2-D matrix
        Returns:
            resized source image matrix
        '''
        h, w = dst.shape
        img = cv2.resize(src, (w, h))
        return img

    def extract_features(self, img):
        '''
        Extract SIFT featrues from input image.
        Returns:
            kp: a list of KeyPoint objects
            desc: a list of SIFT descriptors of the key points
        '''
        sift = cv2.SIFT()
        kp, desc = sift.detectAndCompute(img, None)
        return kp, desc

    def get_matches(self, desc1, desc2):
        '''
        Match features with FLANN.
        Select good matches based on Lowe's ratio test.
        Args:
            desc1: a list of SIFT descriptors of image 1
            desc2: a list of SIFT descriptors of image 2
        Returns:
            a list of DMatch objects
        '''
        FLANN_INDEX_KDTREE = 0
        index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
        search_params = dict(checks=50)
        flann = cv2.FlannBasedMatcher(index_params, search_params)
        matches = flann.knnMatch(desc1, desc2, k=2)

        # Filter matches using ratio test
        def select_good_matches(matches, thres):
            good = list()
            for (m1, m2) in matches:
                if m1.distance < thres * m2.distance:
                    good.append(m1)
            return good

        good = select_good_matches(matches, self.RATIO_TEST_THRES)
        return good if len(good) >= self.MIN_MATCH_COUNT else matches

    def display_matches(self, img1, img2, kp1, kp2, matches):
        '''
        Display feature matches.
        Args:
            img1: gray image 1
            img2: gray image 2
            kp1: key points of image 1
            kp2: key points of image 2
            matches: a list of DMatch objects
        '''
        # Concatenate two gray images side by side
        img3 = np.concatenate((img1, img2), axis=1)
        n = img1.shape[1] # width of the left image

        # Draw matching pairs
        for m in matches:
            (x1, y1) = kp1[m.queryIdx].pt # left
            (x2, y2) = kp2[m.trainIdx].pt # right
            # Convert to integers
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            x2 += n # offset by the width of the image on the left
            # Mark SIFT features
            cv2.circle(img3, (x1, y1), 6, (0, 0, 255), thickness=2)
            cv2.circle(img3, (x2, y2), 6, (0, 0, 255), thickness=2)
            # Draw a line between two matched points
            cv2.line(img3, (x1, y1), (x2, y2), (0, 255, 0))

        # Resize for display
        img3 = cv2.resize(img3, (0, 0), fx=0.8, fy=0.8)
        self.display_img(img3, 'SIFT feature matching')

    def display_img(self, img, title=''):
        ''' Display the given image. '''
        cv2.imshow(title, img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

    def warp(self, img1, img2, kp1, kp2, matches):
        '''
        Warp image 1 towards image 2 using perspective transformation.
        Return:
            Warped image
        '''
        # Get points of interest for matching
        src_pts = np.float32([kp1[m.queryIdx].pt for m in matches]).reshape(-1, 1, 2)
        dst_pts = np.float32([kp2[m.trainIdx].pt for m in matches]).reshape(-1, 1, 2)
        # Find a perspective transformation
        H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 1.0)
        h, w = img1.shape
        # Warp the image
        warped = cv2.warpPerspective(img1, H, (w, h))
        return warped

    def save_img(self, img, filename='test.png'):
        ''' Write image to file. '''
        cv2.imwrite(filename, img)

    def driver(self, path1, path2, show_matches=False):
        '''
        Driver function.
        Args:
            path1: path to image file 1
            path2: path to image file 2
            show_matches: if display SIFT matching results
        '''
        # Read images
        img1 = self.get_gray_img(path1)
        img2 = self.get_gray_img(path2)
        img1 = self.resize_img(img1, img2)

        # Pad zeros to image borders
        pad_width = 65
        npad = ((pad_width, pad_width), (pad_width, pad_width))
        img1 = np.pad(img1, pad_width=npad, mode='constant', constant_values=0)
        img2 = np.pad(img2, pad_width=npad, mode='constant', constant_values=0)

        # Calculate SIFT features
        kp1, desc1 = self.extract_features(img1)
        kp2, desc2 = self.extract_features(img2)

        # Match SIFT features
        matches = self.get_matches(desc1, desc2)
        if show_matches:
            self.display_matches(img1, img2, kp1, kp2, matches)

        # Projective Transformation
        warped = self.warp(img1, img2, kp1, kp2, matches)
        warped = self.resize_img(warped, img2)
        img3 = np.concatenate((warped, img2), axis=1)
        self.display_img(img3, 'Warped')

        # Stitch the warped image and image 2 together
        img = cv2.addWeighted(warped, 0.5, img2, 0.5, 0)
        self.display_img(img, 'Stitched')
        self.save_img(img, 'result.png')

def main():
    path1, path2 = 'graf1.png', 'graf2.png'
    warper = Warper(rtthres=0.7, min_match_count=4)
    warper.driver(path1, path2, show_matches=True)

if __name__ == '__main__':
    main()

max_stars_repo_name: yatbear/bear

max_issues_repo_name: yatbear/bear

max_forks_repo_name: yatbear/bear


text:
#!/usr/bin/env pythonw
'''
popup for 2D XRD mask file

'''

import os
import numpy as np

import wx

from wxmplot.imagepanel import ImagePanel
from .ImageControlsFrame import ImageToolboxFrame

from larch.io import tifffile

# HAS_pyFAI = False
# try:
#     import pyFAI
#     import pyFAI.calibrant
#     from pyFAI.calibration import Calibration
#     HAS_pyFAI = True
# except ImportError:
#     pass

# HAS_fabio = False
# try:
#     import fabio
#     HAS_fabio = True
# except ImportError:
#     pass

###################################

class MaskToolsPopup(wx.Frame):

    def __init__(self,parent):

        self.frame = wx.Frame.__init__(self, parent, title='Create mask',size=(800,600))

        self.parent = parent
        self.statusbar = self.CreateStatusBar(2,wx.CAPTION )

        try:
            self.raw_img = parent.plt_img ## raw_img or flp_img or plt_img mkak 2016.10.28
        except:
            self.loadIMAGE()

        self.setDefaults()

        self.Init()
        self.Show()

#        wx.Window.GetEffectiveMinSize
#        wx.GetBestSize(self)

    def Init(self):

        self.panel = wx.Panel(self)

        self.MainSizer()

        self.framebox = wx.BoxSizer(wx.VERTICAL)
        self.framebox.Add(self.mainbox, flag=wx.ALL|wx.EXPAND, border=10)

        ###########################
        ## Pack all together in self.panel
        self.panel.SetSizer(self.framebox)

    def setDefaults(self):

        self.area_list = []

    def MainSizer(self):

        self.mainbox = wx.BoxSizer(wx.VERTICAL)

        ###########################
        ## -----> Main Panel
        self.hmain = wx.BoxSizer(wx.HORIZONTAL)

        self.ImageSizer()
        self.DrawNewSizer()

        self.hmain.Add(self.imagebox,proportion=1,flag=wx.ALL|wx.EXPAND, border=10)
        self.hmain.Add(self.toolbox, flag=wx.ALL, border=10)

        self.mainbox.Add(self.hmain, flag=wx.ALL|wx.EXPAND, border=10)

    def DrawNewSizer(self):

        self.toolbox = wx.BoxSizer(wx.VERTICAL)

        ###########################
        ## Directions
        nwbx = wx.StaticBox(self.panel,label='Drawing Tools', size=(100, 50))
        drawbox = wx.StaticBoxSizer(nwbx,wx.VERTICAL)

        ###########################
        ## Drawing tools
        hbox_shp = wx.BoxSizer(wx.HORIZONTAL)
        self.txt_shp = wx.StaticText(self.panel, label='DRAWING SHAPE')
        shapes = ['circle','pixel','polygon','square']

        self.ch_shp = wx.Choice(self.panel,choices=shapes)

        self.ch_shp.Bind(wx.EVT_CHOICE,self.ShapeChoice)

        hbox_shp.Add(self.txt_shp, flag=wx.RIGHT, border=8)
        hbox_shp.Add(self.ch_shp, flag=wx.EXPAND, border=8)

        drawbox.Add(hbox_shp, flag=wx.ALL|wx.EXPAND, border=10)

        ###########################
        ## Drawn Areas
        vbox_areas = wx.BoxSizer(wx.VERTICAL)
        hbox_areas = wx.BoxSizer(wx.HORIZONTAL)

        self.slct_area = wx.ListBox(self.panel, 26, wx.DefaultPosition, (170, 130), self.area_list, wx.LB_SINGLE)

        self.btn_SHWarea = wx.Button(self.panel,label='SHOW')
        self.btn_DELarea = wx.Button(self.panel,label='DELETE')

        self.btn_SHWarea.Bind(wx.EVT_BUTTON,self.showAREA)
        self.btn_DELarea.Bind(wx.EVT_BUTTON,self.deleteAREA)

        hbox_areas.Add(self.btn_SHWarea, flag=wx.BOTTOM, border=10)
        hbox_areas.Add(self.btn_DELarea, flag=wx.BOTTOM, border=10)

        self.btn_clear = wx.Button(self.panel,label='CLEAR ALL')

        self.btn_clear.Bind(wx.EVT_BUTTON,self.clearMask)

        vbox_areas.Add(self.slct_area, flag=wx.ALL|wx.EXPAND, border=10)
        vbox_areas.Add(hbox_areas, flag=wx.ALL|wx.EXPAND, border=10)
        vbox_areas.Add(self.btn_clear, flag=wx.ALL|wx.EXPAND, border=10)
        drawbox.Add(vbox_areas, flag=wx.ALL|wx.EXPAND, border=10)

        self.toolbox.Add(drawbox, flag=wx.ALL, border=10)

        self.btn_save = wx.Button(self.panel,label='SAVE MASK')
        self.btn_save.Bind(wx.EVT_BUTTON,self.saveMask)
        self.toolbox.Add(self.btn_save, flag=wx.ALL, border=10)

        self.btn_SHWarea.Disable()
        self.btn_DELarea.Disable()
        self.btn_clear.Disable()
        self.btn_save.Disable()

    def startIMAGE(self):

        self.loadIMAGE()

    def loadIMAGE(self,event=None):
        wildcards = 'XRD image (*.edf,*.tif,*.tiff)|*.tif;*.tiff;*.edf|All files (*.*)|*.*'
        dlg = wx.FileDialog(self, message='Choose XRD image',
                           defaultDir=os.getcwd(),
                           wildcard=wildcards, style=wx.FD_OPEN)

        path, read = None, False
        if dlg.ShowModal() == wx.ID_OK:
            read = True
            path = dlg.GetPath().replace('\\', '/')
        dlg.Destroy()

        if read:
            try:
#                 self.raw_img = plt.imread(path)
                self.raw_img = tifffile.imread(path)
                #self.raw_img = fabio.open(path).data
            except:
                print('Image not properly opened.')
                self.raw_img = np.zeros((1024,1024))
        else:
            print('No image selected.')
            self.raw_img = np.zeros((1024,1024))

    def ImageSizer(self):
        '''
        Image Panel
        '''
        self.imagebox = wx.BoxSizer(wx.VERTICAL)

        self.plot2Dimage()

        imagetools = wx.BoxSizer(wx.HORIZONTAL)

        self.btn_image = wx.Button(self.panel,label='IMAGE TOOLS')
        self.btn_load = wx.Button(self.panel,label='CHANGE IMAGE')

        self.btn_image.Bind(wx.EVT_BUTTON,self.onImageTools)
        self.btn_load.Bind(wx.EVT_BUTTON,self.loadIMAGE)

        imagetools.Add(self.btn_load, flag=wx.ALL, border=10)
        imagetools.Add(self.btn_image, flag=wx.ALL, border=10)

        self.imagebox.Add(self.plot2Dimg,proportion=1,flag=wx.ALL|wx.EXPAND, border=10)
        self.imagebox.Add(imagetools, flag=wx.ALL, border=10)

    def write_message(self, s, panel=0):
        """write a message to the Status Bar"""
        self.SetStatusText(s, panel)

    def onImageTools(self,event=None):

        self.toolbox = ImageToolboxFrame(self.plot2Dimg,self.raw_img)

    def plot2Dimage(self):

        self.plot2Dimg = ImagePanel(self.panel)#,size=(300, 300))
        self.plot2Dimg.messenger = self.write_message

        self.Bind(wx.EVT_PAINT, self.OnPaint)

        self.plot2Dimg.display(self.raw_img)
        self.AutoContrast()

        self.plot2Dimg.redraw()

    def AutoContrast(self):

        self.minINT = int(np.min(self.raw_img))
        self.maxINT = int(np.max(self.raw_img)/15) # /15 scales image to viewable
        if self.maxINT == self.minINT:
            self.minINT = self.minINT-50
            self.maxINT = self.minINT+100

        self.minCURRENT = self.minINT
        self.maxCURRENT = self.maxINT
        if self.maxCURRENT > self.maxINT:
            self.maxCURRENT = self.maxINT

        self.plot2Dimg.conf.auto_intensity = False
        self.plot2Dimg.conf.int_lo[0] = self.minCURRENT
        self.plot2Dimg.conf.int_hi[0] = self.maxCURRENT
#         self.plot2Dimg.conf.int_lo['int'] = self.minCURRENT
#         self.plot2Dimg.conf.int_hi['int'] = self.maxCURRENT

        ## vertical flip default
        self.plot2Dimg.conf.flip_ud = True
        self.plot2Dimg.conf.flip_lr = False

        self.plot2Dimg.redraw()

    def ShapeChoice(self,event=None):

        print('The shape you choose: %s' %  self.ch_shp.GetString(self.ch_shp.GetSelection()))

        print()
        print('Not implemented: ShapeChoice')
        self.addAREA()

    def OnPaint(self, event=None):

        print('Shape : %s' %  self.ch_shp.GetString(self.ch_shp.GetSelection()))
#         dc = wx.PaintDC(self)
#         dc.Clear()
#         dc.SetPen(wx.Pen(wx.BLACK, 4))
#         dc.DrawLine(0, 0, 50, 50)

    def clearMask(self,event=None):

        print('Clearing the mask...')
        ## provide a warning message?
        self.area_list = []
        self.slct_area.Set(self.area_list)

        self.btn_SHWarea.Disable()
        self.btn_DELarea.Disable()
        self.btn_clear.Disable()
        self.btn_save.Disable()

    def saveMask(self,event=None):

        print('This will trigger the saving of a mask.')

    def addAREA(self,event=None):
        area_name = 'area %i (%s)' % (len(self.area_list),self.ch_shp.GetString(self.ch_shp.GetSelection()))
        self.area_list.append(area_name)
        self.slct_area.Set(self.area_list)

        if len(self.area_list) > 0:
            self.btn_SHWarea.Enable()
            self.btn_DELarea.Enable()
            self.btn_clear.Enable()
            self.btn_save.Enable()

    def showAREA(self,event=None):

        if len(self.area_list) > 0:
            area_str = self.slct_area.GetString(self.slct_area.GetSelection())
            str_msg = 'Displaying: %s' % area_str
            self.write_message(str_msg,panel=0)

            ## show area on map, image

    def deleteAREA(self,event=None):

        if len(self.area_list) > 0:
            area_str = self.slct_area.GetString(self.slct_area.GetSelection())
            str_msg = 'Deleting: %s' % area_str
            self.write_message(str_msg,panel=0)

            self.area_list.remove(self.slct_area.GetString(self.slct_area.GetSelection()))
            self.slct_area.Set(self.area_list)

        if len(self.area_list) == 0:
            self.btn_SHWarea.Disable()
            self.btn_DELarea.Disable()
            self.btn_clear.Disable()
            self.btn_save.Disable()

# class diFFit_XRDmask(wx.App):
#     def __init__(self):
#         wx.App.__init__(self)
#
#     def run(self):
#         self.MainLoop()
#
#     def createApp(self):
#         frame = MaskToolsPopup(None)
#         frame.Show()
#         self.SetTopWindow(frame)
#
#     def OnInit(self):
#         self.createApp()
#         return True
#
# class DebugViewer(diFFit_XRDmask):
#     def __init__(self, **kws):
#         diFFit_XRDmask.__init__(self, **kws)
#
#     def OnInit(self):
#         #self.Init()
#         self.createApp()
#         #self.ShowInspectionTool()
#         return True
#
# if __name__ == '__main__':
#     diFFit_XRDmask().run()


text:
"""
Quality control of raw Neuropixel electrophysiology data.
"""
from pathlib import Path
import logging
import shutil

import numpy as np
import pandas as pd
from scipy import signal

import alf.io
from brainbox.core import Bunch
from brainbox.metrics.single_units import spike_sorting_metrics
from ibllib.ephys import sync_probes
from ibllib.io import spikeglx, raw_data_loaders
import ibllib.dsp as dsp
from ibllib.io.extractors import ephys_fpga, training_wheel
from ibllib.misc import print_progress
from phylib.io import model

_logger = logging.getLogger('ibllib')

RMS_WIN_LENGTH_SECS = 3
WELCH_WIN_LENGTH_SAMPLES = 1024
NCH_WAVEFORMS = 32  # number of channels to be saved in templates.waveforms and channels.waveforms

def rmsmap(fbin):
    """
    Computes RMS map in time domain and spectra for each channel of Neuropixel probe

    :param fbin: binary file in spike glx format (will look for attached metatdata)
    :type fbin: str or pathlib.Path
    :return: a dictionary with amplitudes in channeltime space, channelfrequency space, time
     and frequency scales
    """
    if not isinstance(fbin, spikeglx.Reader):
        sglx = spikeglx.Reader(fbin)
    rms_win_length_samples = 2 ** np.ceil(np.log2(sglx.fs * RMS_WIN_LENGTH_SECS))
    # the window generator will generates window indices
    wingen = dsp.WindowGenerator(ns=sglx.ns, nswin=rms_win_length_samples, overlap=0)
    # pre-allocate output dictionary of numpy arrays
    win = {'TRMS': np.zeros((wingen.nwin, sglx.nc)),
           'nsamples': np.zeros((wingen.nwin,)),
           'fscale': dsp.fscale(WELCH_WIN_LENGTH_SAMPLES, 1 / sglx.fs, one_sided=True),
           'tscale': wingen.tscale(fs=sglx.fs)}
    win['spectral_density'] = np.zeros((len(win['fscale']), sglx.nc))
    # loop through the whole session
    for first, last in wingen.firstlast:
        D = sglx.read_samples(first_sample=first, last_sample=last)[0].transpose()
        # remove low frequency noise below 1 Hz
        D = dsp.hp(D, 1 / sglx.fs, [0, 1])
        iw = wingen.iw
        win['TRMS'][iw, :] = dsp.rms(D)
        win['nsamples'][iw] = D.shape[1]
        # the last window may be smaller than what is needed for welch
        if last - first < WELCH_WIN_LENGTH_SAMPLES:
            continue
        # compute a smoothed spectrum using welch method
        _, w = signal.welch(D, fs=sglx.fs, window='hanning', nperseg=WELCH_WIN_LENGTH_SAMPLES,
                            detrend='constant', return_onesided=True, scaling='density', axis=-1)
        win['spectral_density'] += w.T
        # print at least every 20 windows
        if (iw % min(20, max(int(np.floor(wingen.nwin / 75)), 1))) == 0:
            print_progress(iw, wingen.nwin)
    return win

def extract_rmsmap(fbin, out_folder=None, overwrite=False):
    """
    Wrapper for rmsmap that outputs _ibl_ephysRmsMap and _ibl_ephysSpectra ALF files

    :param fbin: binary file in spike glx format (will look for attached metatdata)
    :param out_folder: folder in which to store output ALF files. Default uses the folder in which
     the `fbin` file lives.
    :param overwrite: do not re-extract if all ALF files already exist
    :param label: string or list of strings that will be appended to the filename before extension
    :return: None
    """
    _logger.info(f"Computing QC for {fbin}")
    sglx = spikeglx.Reader(fbin)
    # check if output ALF files exist already:
    if out_folder is None:
        out_folder = Path(fbin).parent
    else:
        out_folder = Path(out_folder)
    alf_object_time = f'ephysTimeRms{sglx.type.upper()}'
    alf_object_freq = f'ephysSpectralDensity{sglx.type.upper()}'
    files_time = list(out_folder.glob(f"_iblqc_{alf_object_time}*"))
    files_freq = list(out_folder.glob(f"_iblqc_{alf_object_freq}*"))
    if (len(files_time) == 2 == len(files_freq)) and not overwrite:
        _logger.warning(f'{fbin.name} QC already exists, skipping. Use overwrite option.')
        return files_time + files_freq
    # crunch numbers
    rms = rmsmap(fbin)
    # output ALF files, single precision with the optional label as suffix before extension
    if not out_folder.exists():
        out_folder.mkdir()
    tdict = {'rms': rms['TRMS'].astype(np.single), 'timestamps': rms['tscale'].astype(np.single)}
    fdict = {'power': rms['spectral_density'].astype(np.single),
             'freqs': rms['fscale'].astype(np.single)}
    out_time = alf.io.save_object_npy(
        out_folder, object=alf_object_time, dico=tdict, namespace='iblqc')
    out_freq = alf.io.save_object_npy(
        out_folder, object=alf_object_freq, dico=fdict, namespace='iblqc')
    return out_time + out_freq

def raw_qc_session(session_path, overwrite=False):
    """
    Wrapper that exectutes QC from a session folder and outputs the results whithin the same folder
    as the original raw data.
    :param session_path: path of the session (Subject/yyyy-mm-dd/number
    :param overwrite: bool (False) Force means overwriting an existing QC file
    :return: None
    """
    efiles = spikeglx.glob_ephys_files(session_path)
    qc_files = []
    for efile in efiles:
        if efile.get('ap') and efile.ap.exists():
            qc_files.extend(extract_rmsmap(efile.ap, out_folder=None, overwrite=overwrite))
        if efile.get('lf') and efile.lf.exists():
            qc_files.extend(extract_rmsmap(efile.lf, out_folder=None, overwrite=overwrite))
    return qc_files

def validate_ttl_test(ses_path, display=False):
    """
    For a mock session on the Ephys Choice world task, check the sync channels for all
    device properly connected and perform a synchronization if dual probes to check that
    all channels are recorded properly
    :param ses_path: session path
    :param display: show the probe synchronization plot if several probes
    :return: True if tests pass, errors otherwise
    """

    def _single_test(assertion, str_ok, str_ko):
        if assertion:
            _logger.info(str_ok)
            return True
        else:
            _logger.error(str_ko)
            return False

    EXPECTED_RATES_HZ = {'left_camera': 60, 'right_camera': 150, 'body_camera': 30}
    SYNC_RATE_HZ = 1
    MIN_TRIALS_NB = 6

    ok = True
    ses_path = Path(ses_path)
    if not ses_path.exists():
        return False

    # get the synchronization fronts (from the raw binary if necessary)
    ephys_fpga.extract_sync(session_path=ses_path, overwrite=False)
    rawsync, sync_map = ephys_fpga.get_main_probe_sync(ses_path)
    last_time = rawsync['times'][-1]

    # get upgoing fronts for each
    sync = Bunch({})
    for k in sync_map:
        fronts = ephys_fpga._get_sync_fronts(rawsync, sync_map[k])
        sync[k] = fronts['times'][fronts['polarities'] == 1]
    wheel = ephys_fpga.extract_wheel_sync(rawsync, chmap=sync_map)

    frame_rates = {'right_camera': np.round(1 / np.median(np.diff(sync.right_camera))),
                   'left_camera': np.round(1 / np.median(np.diff(sync.left_camera))),
                   'body_camera': np.round(1 / np.median(np.diff(sync.body_camera)))}

    # check the camera frame rates
    for lab in frame_rates:
        expect = EXPECTED_RATES_HZ[lab]
        ok &= _single_test(assertion=abs((1 - frame_rates[lab] / expect)) < 0.1,
                           str_ok=f'PASS: {lab} frame rate: {frame_rates[lab]} = {expect} Hz',
                           str_ko=f'FAILED: {lab} frame rate: {frame_rates[lab]} != {expect} Hz')

    # check that the wheel has a minimum rate of activity on both channels
    re_test = abs(1 - sync.rotary_encoder_1.size / sync.rotary_encoder_0.size) < 0.1
    re_test &= len(wheel[1]) / last_time > 5
    ok &= _single_test(assertion=re_test,
                       str_ok="PASS: Rotary encoder", str_ko="FAILED: Rotary encoder")
    # check that the frame 2 ttls has a minimum rate of activity
    ok &= _single_test(assertion=len(sync.frame2ttl) / last_time > 0.2,
                       str_ok="PASS: Frame2TTL", str_ko="FAILED: Frame2TTL")
    # the audio has to have at least one event per trial
    ok &= _single_test(assertion=len(sync.bpod) > len(sync.audio) > MIN_TRIALS_NB,
                       str_ok="PASS: audio", str_ko="FAILED: audio")
    # the bpod has to have at least twice the amount of min trial pulses
    ok &= _single_test(assertion=len(sync.bpod) > MIN_TRIALS_NB * 2,
                       str_ok="PASS: Bpod", str_ko="FAILED: Bpod")
    try:
        # note: tried to depend as little as possible on the extraction code but for the valve...
        behaviour = ephys_fpga.extract_behaviour_sync(rawsync, chmap=sync_map)
        res = behaviour.valveOpen_times.size > 1
    except AssertionError:
        res = False
    # check that the reward valve is actionned at least once
    ok &= _single_test(assertion=res,
                       str_ok="PASS: Valve open", str_ko="FAILED: Valve open not detected")
    _logger.info('ALL CHECKS PASSED !')

    # the imec sync is for 3B Probes only
    if sync.get('imec_sync') is not None:
        ok &= _single_test(assertion=np.all(1 - SYNC_RATE_HZ * np.diff(sync.imec_sync) < 0.1),
                           str_ok="PASS: imec sync", str_ko="FAILED: imec sync")

    # second step is to test that we can make the sync. Assertions are whithin the synch code
    if sync.get('imec_sync') is not None:
        sync_result, _ = sync_probes.version3B(ses_path, display=display)
    else:
        sync_result, _ = sync_probes.version3A(ses_path, display=display)

    ok &= _single_test(assertion=sync_result, str_ok="PASS: synchronisation",
                       str_ko="FAILED: probe synchronizations threshold exceeded")

    if not ok:
        raise ValueError('FAILED TTL test')
    return ok

def spike_sorting_metrics_ks2(ks2_path=None, m=None, save=True):
    """
    Given a path containing kilosort 2 output, compute quality metrics and optionally save them
    to a clusters_metric.csv file
    :param ks2_path:
    :param save
    :return:
    """

    # ensure that either a ks2_path or a phylib `TemplateModel` object with unit info is given
    assert not(ks2_path is None and m is None), 'Must either specify a path to a ks2 output ' \
                                                'directory, or a phylib `TemplateModel` object'
    # create phylib `TemplateModel` if not given
    m = phy_model_from_ks2_path(ks2_path) if None else m
    c, drift = spike_sorting_metrics(m.spike_times, m.spike_clusters, m.amplitudes, m.depths)
    #  include the ks2 cluster contamination if `cluster_ContamPct` file exists
    file_contamination = ks2_path.joinpath('cluster_ContamPct.tsv')
    if file_contamination.exists():
        contam = pd.read_csv(file_contamination, sep='\t')
        contam.rename(columns={'ContamPct': 'ks2_contamination_pct'}, inplace=True)
        c = c.set_index('cluster_id', drop=False).join(contam.set_index('cluster_id'))

    #  include the ks2 cluster labels if `cluster_KSLabel` file exists
    file_labels = ks2_path.joinpath('cluster_KSLabel.tsv')
    if file_labels.exists():
        ks2_labels = pd.read_csv(file_labels, sep='\t')
        ks2_labels.rename(columns={'KSLabel': 'ks2_label'}, inplace=True)
        c = c.set_index('cluster_id', drop=False).join(ks2_labels.set_index('cluster_id'))

    if save:
        #  the file name contains the label of the probe (directory name in this case)
        c.to_csv(ks2_path.joinpath('cluster_metrics.csv'))

    return c

def phy_model_from_ks2_path(ks2_path, bin_path, bin_file=None):
    if not bin_file:
        bin_file = next(bin_path.rglob('*.ap.*bin'), None)
    meta_file = next(bin_path.rglob('*.ap.meta'), None)
    if meta_file and meta_file.exists():
        meta = spikeglx.read_meta_data(meta_file)
        fs = spikeglx._get_fs_from_meta(meta)
        nch = (spikeglx._get_nchannels_from_meta(meta) -
               len(spikeglx._get_sync_trace_indices_from_meta(meta)))
    else:
        fs = 30000
        nch = 384
    m = model.TemplateModel(dir_path=ks2_path,
                            dat_path=bin_file,  # this assumes the raw data is in the same folder
                            sample_rate=fs,
                            n_channels_dat=nch,
                            n_closest_channels=NCH_WAVEFORMS)
    m.depths = m.get_depths()
    return m

# Make a bunch gathering all trial QC
def qc_fpga_task(fpga_trials, alf_trials):
    """
    :fpga_task is the dictionary output of
    ibllib.io.extractors.ephys_fpga.extract_behaviour_sync
    : bpod_trials is the dictionary output of ibllib.io.extractors.ephys_trials.extract_all
    : alf_trials is the ALF _ibl_trials object after extraction (alf.io.load_object)
    :return: qc_session, qc_trials, True means QC passes while False indicates a failure
    """

    GOCUE_STIMON_DELAY = 0.01  # -> 0.1
    FEEDBACK_STIMFREEZE_DELAY = 0.01  # -> 0.1
    VALVE_STIM_OFF_DELAY = 1
    VALVE_STIM_OFF_JITTER = 0.1
    ITI_IN_STIM_OFF_JITTER = 0.1
    ERROR_STIM_OFF_DELAY = 2
    ERROR_STIM_OFF_JITTER = 0.1
    RESPONSE_FEEDBACK_DELAY = 0.0005

    def strictly_after(t0, t1, threshold):
        """ returns isafter, iswithinthreshold"""
        return (t1 - t0) > 0, np.abs((t1 - t0)) <= threshold

    ntrials = fpga_trials['stimOn_times'].size
    qc_trials = Bunch({})

    """
    First Check consistency of the dataset: whithin each trial, all events happen after trial
    start should not be NaNs and increasing. This is not a QC but an assertion.
    """
    status = True
    for k in ['response_times', 'stimOn_times', 'response_times',
              'goCueTrigger_times', 'goCue_times', 'feedback_times']:
        if k.endswith('_bpod'):
            tstart = alf_trials['intervals_bpod'][:, 0]
        else:
            tstart = alf_trials['intervals'][:, 0]
        selection = ~np.isnan(alf_trials[k])
        status &= np.all(alf_trials[k][selection] - tstart[selection] > 0)
        status &= np.all(np.diff(alf_trials[k][selection]) > 0)
    assert status

    """
    This part of the function uses only fpga_trials information
    """
    # check number of feedbacks: should always be one
    qc_trials['n_feedback'] = (np.uint32(~np.isnan(fpga_trials['valveOpen_times'])) +
                               np.uint32(~np.isnan(fpga_trials['errorCue_times'])))

    # check for non-Nans
    qc_trials['stimOn_times_nan'] = ~np.isnan(fpga_trials['stimOn_times'])
    qc_trials['goCue_times_nan'] = ~np.isnan(fpga_trials['goCue_times'])

    # stimOn before goCue
    qc_trials['stimOn_times_before_goCue_times'], qc_trials['stimOn_times_goCue_times_delay'] =\
        strictly_after(fpga_trials['stimOn_times'], fpga_trials['goCue_times'], GOCUE_STIMON_DELAY)

    # stimFreeze before feedback
    qc_trials['stim_freeze_before_feedback'], qc_trials['stim_freeze_feedback_delay'] = \
        strictly_after(fpga_trials['stimFreeze_times'], fpga_trials['feedback_times'],
                       FEEDBACK_STIMFREEZE_DELAY)

    # stimOff 1 sec after valve, with 0.1 as acceptable jitter
    qc_trials['stimOff_delay_valve'] = np.less(
        np.abs(
            fpga_trials['stimOff_times'] - fpga_trials['valveOpen_times'] - VALVE_STIM_OFF_DELAY
        ),
        VALVE_STIM_OFF_JITTER, out=np.ones(ntrials, dtype=bool),
        where=~np.isnan(fpga_trials['valveOpen_times']))

    # iti_in whithin 0.01 sec of stimOff
    qc_trials['iti_in_delay_stim_off'] = \
        np.abs(fpga_trials['stimOff_times'] - fpga_trials['itiIn_times']) < ITI_IN_STIM_OFF_JITTER

    # stimOff 2 secs after errorCue_times with jitter
    # noise off happens 2 secs after stimm, with 0.1 as acceptable jitter
    qc_trials['stimOff_delay_noise'] = np.less(
        np.abs(
            fpga_trials['stimOff_times'] - fpga_trials['errorCue_times'] - ERROR_STIM_OFF_DELAY
        ),
        ERROR_STIM_OFF_JITTER, out=np.ones(ntrials, dtype=bool),
        where=~np.isnan(fpga_trials['errorCue_times']))

    """
    This part uses only alf_trials information
    """
    # TEST  Response times (from session start) should be increasing continuously
    #       Note: RT are not durations but time stamps from session start
    #       1. check for non-Nans
    qc_trials['response_times_nan'] = ~np.isnan(alf_trials['response_times'])
    #       2. check for positive increase
    qc_trials['response_times_increase'] = \
        np.diff(np.append([0], alf_trials['response_times'])) > 0
    # TEST  Response times (from goCue) should be positive
    qc_trials['response_times_goCue_times_diff'] = \
        alf_trials['response_times'] - alf_trials['goCue_times'] > 0
    # TEST  1. Response_times should be before feedback
    qc_trials['response_before_feedback'] = \
        alf_trials['feedback_times'] - alf_trials['response_times'] > 0
    #       2. Delay between wheel reaches threshold (response time) and
    #       feedback is 100us, acceptable jitter 500 us
    qc_trials['response_feedback_delay'] = \
        alf_trials['feedback_times'] - alf_trials['response_times'] < RESPONSE_FEEDBACK_DELAY

    # Test output at session level
    qc_session = {k: np.all(qc_trials[k]) for k in qc_trials}

    return qc_session, qc_trials

def _qc_from_path(sess_path, display=True):
    WHEEL = False
    sess_path = Path(sess_path)
    temp_alf_folder = sess_path.joinpath('fpga_test', 'alf')
    temp_alf_folder.mkdir(parents=True, exist_ok=True)

    raw_trials = raw_data_loaders.load_data(sess_path)
    tmax = raw_trials[-1]['behavior_data']['States timestamps']['exit_state'][0][-1] + 60

    sync, chmap = ephys_fpga.get_main_probe_sync(sess_path, bin_exists=False)
    _ = ephys_fpga.extract_all(sess_path, output_path=temp_alf_folder, save=True)
    # check that the output is complete
    fpga_trials = ephys_fpga.extract_behaviour_sync(sync, output_path=temp_alf_folder, tmax=tmax,
                                                    chmap=chmap, save=True, display=display)
    # align with the bpod
    bpod2fpga = ephys_fpga.align_with_bpod(temp_alf_folder.parent)
    alf_trials = alf.io.load_object(temp_alf_folder, 'trials')
    shutil.rmtree(temp_alf_folder)
    # do the QC
    qcs, qct = qc_fpga_task(fpga_trials, alf_trials)

    # do the wheel part
    if WHEEL:
        bpod_wheel = training_wheel.get_wheel_data(sess_path, save=False)
        fpga_wheel = ephys_fpga.extract_wheel_sync(sync, chmap=chmap, save=False)

        if display:
            import matplotlib.pyplot as plt
            t0 = max(np.min(bpod2fpga(bpod_wheel['re_ts'])), np.min(fpga_wheel['re_ts']))
            dy = np.interp(t0, fpga_wheel['re_ts'], fpga_wheel['re_pos']) - np.interp(
                t0, bpod2fpga(bpod_wheel['re_ts']), bpod_wheel['re_pos'])

            fix, axes = plt.subplots(nrows=2, sharex='all', sharey='all')
            # axes[0].plot(t, pos), axes[0].title.set_text('Extracted')
            axes[0].plot(bpod2fpga(bpod_wheel['re_ts']), bpod_wheel['re_pos'] + dy)
            axes[0].plot(fpga_wheel['re_ts'], fpga_wheel['re_pos'])
            axes[0].title.set_text('FPGA')
            axes[1].plot(bpod2fpga(bpod_wheel['re_ts']), bpod_wheel['re_pos'] + dy)
            axes[1].title.set_text('Bpod')

    return alf.io.dataframe({**fpga_trials, **alf_trials, **qct})

max_stars_repo_name: nbonacchi/ibllib

max_issues_repo_name: nbonacchi/ibllib

max_forks_repo_name: nbonacchi/ibllib

reason: import numpy,from scipy


text:
"""
Convert one annotation format to another.

Supported Conversions:
    KITTI -> YOLO
    KITTI -> VOC
    VOC   -> KITTI
    VOC   -> YOLO
    YOLO  -> KITTI
    YOLO  -> VOC

Usage:
    python convert.py
"""

from __future__ import print_function
import os, sys, shutil, glob, argparse
import cv2, random
import numpy as np
import os.path as osp
from PIL import Image, ImageOps

max_stars_repo_name: eweill/datasets

max_issues_repo_name: eweill/datasets

max_forks_repo_name: eweill/datasets


text:
import pandas as pd
import numpy as np
import configparser
import os
import argparse

def read_values_fromfile_as_string (file , joinsuffix=','):
    with open(file) as file:
        rows = [item.strip()for item in file.readlines()]

    return joinsuffix.join(rows)

if __name__=='__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument('-excel_file', action='store',
                        dest='excel',
                        help='excel file with experiments definition', type=str, required=True)

    parser.add_argument('-sheet', action='store',
                        dest='sheet',
                        help='sheet number to use in the excel file', type=str, required=True)

    parser.add_argument('-experiment', action='store',
                        dest='experiment',
                        help=' the number of the experiment to consider', type=int, required=True)

    parser.add_argument('-conf_file', action='store',
                        dest='conf_file',
                        help='conf file to update', type=str, required=True)

    parser.add_argument('-experiment_attribute', action='store',
                        dest='experiment_atr',
                        help='sheet name in the excel', type=str, default='experiment')

    parser.add_argument('-multivalue_attribute_suffix', action='store',
                        dest='multivalue_atr_suffix',
                        help='sheet name in the excel', type=str, default='*')

    parser.add_argument('-join_multivalue', action='store',
                        dest='join_multivalue',
                        help='suffix used to join multivalue', type=str, default=',')

    parser.add_argument('-join_section', action='store',
                        dest='join_section',
                        help='suffix used to join value with section', type=str, default='@')

    args = parser.parse_args()

    excel_file = args.excel

    sheet_name = args.sheet

    expt_name = args.experiment

    conf_file = args.conf_file

    experiment_attribute = args.experiment_atr

    multivariable_attribute_suffix = args.multivalue_atr_suffix

    join_multivalue = args.join_multivalue

    join_section = args.join_section

    dataframe_dict = pd.read_excel(excel_file,sheet_name=None)

    assert sheet_name in dataframe_dict.keys(), 'the specified sheet %s is not in the excel file'%(sheet_name)

    dataframe = pd.read_excel(excel_file,sheet_name=sheet_name)

    assert experiment_attribute in dataframe.columns, 'there is not experiment attribute in the dataframe'

    assert not dataframe[
        experiment_attribute].isnull().values.any(), 'there is at least one value in experiment column which is null'

    dataframe_experiment = dataframe[dataframe[experiment_attribute]==expt_name]

    assert not dataframe_experiment.empty, 'the data frame selected with the experiment is empty'

    cols=[item for item in dataframe_experiment.columns if item!=experiment_attribute]
    config = configparser.ConfigParser()
    for row in dataframe_experiment.index:
        for col in cols:
            if len(col.split('.'))>0:
                col_conf = col.split('.')[0]
            else:
                col_conf = col
            value_section = dataframe_experiment.at[row, col]
            if not pd.isna(value_section):
                value_section = str(value_section)
                assert join_section in value_section, 'The value %s in row %s nd column %s is missing the section'%(str(value_section),str(row),str(col))
                assert len(value_section.split(join_section))==2, 'The value %s in row %s nd column %s is not well formatted ' % (str(value_section), str(row), str(col))
                value, section = value_section.split(join_section)
                if multivariable_attribute_suffix in col_conf:
                    col_conf = col_conf.split(multivariable_attribute_suffix)[1]
                    if os.path.exists(value):
                        value = read_values_fromfile_as_string(value,joinsuffix=join_multivalue)
                    else:
                        value=None
                if section not in config:
                    config[section] = {}
                config[section][col_conf] = str(value)
with open(conf_file, 'w') as file:
    config.write(file)


text:
import pytest
import brainlit
from brainlit.algorithms.generate_fragments import tube_seg
import numpy as np
from brainlit.utils.session import NeuroglancerSession
from brainlit.utils.Neuron_trace import NeuronTrace
from skimage import draw
from pathlib import Path

top_level = Path(__file__).parents[4] / "data"
input = (top_level / "data_octree").as_posix()
url = (top_level / "test_upload").as_uri()
url_seg = url + "_segments"
url = url + "/serial"

def test_pairwise():

    """
    For a given iterable array [A1,A2,...,An], test if the function can return a zipped list [(A1,A2),(A2,A3),...,(An-1,An)]

    The list should contain n-1 pairs(2-element tuple)

    The first element of all the tuples are [A1,A2,...,An-1]

    The second element of all the tubles are [A2,A3,...,An]
    """
    n = np.random.randint(4, 9)
    iterable = np.random.randint(10, size=(n, 3))
    pair = list(tube_seg.pairwise(iterable))

    """
    Verify:

    I. the number of pairs
    II. the first element of each pair
    III. the second elemnt of each pair
    """
    assert len(pair) == n - 1
    assert (iterable[:-1, :] == [a[0] for a in pair]).all()
    assert (iterable[1:, :] == [b[1] for b in pair]).all()

def test_draw_sphere():

    """
    Test if the function maps all the points located within the given radius of the given center to 1, otherwise 0

    The output array should have the same size of input image and binary values

    Distance between a point and the given center:
             <= radius (if the point has value 1)
             >  radius (if the point has value 0)
    """
    ngl_session = NeuroglancerSession(url=url, url_segments=url_seg)
    img, _, _ = ngl_session.pull_vertex_list(2, [4], expand=True)
    shape = img.shape
    center = [
        np.random.randint(shape[0]),
        np.random.randint(shape[1]),
        np.random.randint(shape[2]),
    ]
    radius = np.random.randint(1, 4)
    sphere = tube_seg.draw_sphere(shape, center, radius)
    coords = np.where(sphere < 1)
    d_bg = min(np.sum((np.array(coords).T - center) ** 2, axis=1))
    coords = np.where(sphere > 0)
    d_s = max(np.sum((np.array(coords).T - center) ** 2, axis=1))

    """
    Verify:

    I. the size of output array
    II. if the output is binary-valued
    III. minimum distance between 0-valued points and the center is greater than radius
    IV. maximum distance between 1-valued points and the center is less than or equal to radius
    """
    assert sphere.shape == shape
    assert np.unique(sphere).all() in [0, 1]
    assert d_bg > radius ** 2
    assert d_s <= radius ** 2

def test_draw_tube_spheres():

    """
    Test if the function maps all the points within the radius of a segment line (defined by 2 given points) to 1, otherwise 0

    The output array should have the same size of input image and binary values

    Distance between a point and the segment line:
             <= radius (if the point has value 1)
             >  radius (if the point has value 0)
    """
    ngl_session = NeuroglancerSession(url=url, url_segments=url_seg)
    img, _, _ = ngl_session.pull_vertex_list(2, [4], expand=True)
    shape = img.shape
    vertex0 = [
        np.random.randint(shape[0] / 2),
        np.random.randint(shape[1]),
        np.random.randint(shape[2]),
    ]
    vertex1 = [
        np.random.randint(shape[0] / 2, shape[0]),
        np.random.randint(shape[1]),
        np.random.randint(shape[2]),
    ]
    radius = np.random.randint(1, 4)
    labels = tube_seg.draw_tube_from_spheres(img, vertex0, vertex1, radius)
    line = draw.line_nd(vertex0, vertex1, endpoint=True)
    coords = np.where(labels < 1)
    d_bg = max(shape)
    for pt in np.array(coords).T:
        distance_min = min(np.sum((np.array(line).T - pt) ** 2, axis=1))
        d_bg = min(distance_min, d_bg)
    coords = np.where(labels > 0)
    d_tube = 0
    for pt in np.array(coords).T:
        distance_min = min(np.sum((np.array(line).T - pt) ** 2, axis=1))
        d_tube = max(distance_min, d_tube)

    """
    Verify:

    I. the size of output array
    II. if the output is binary-valued
    III. minimum distance between 0-valued points and the segment line is greater than radius
    IV. maximum distance between 1-valued points and the segment line is less than or equal to radius
    """
    assert labels.shape == shape
    assert np.unique(labels).all() in [0, 1]
    assert d_bg > radius ** 2
    assert d_tube <= radius ** 2

def test_draw_tube_edt():

    """
    Test if the function maps all the points within the radius of a segment line (defined by 2 given points) to 1, otherwise 0

    The output array should have the same size of input image and binary values

    Distance between a point and the segment line:
             <= radius (if the point has value 1)
             >  radius (if the point has value 0)
    """
    ngl_session = NeuroglancerSession(url=url, url_segments=url_seg)
    img, _, _ = ngl_session.pull_vertex_list(2, [4], expand=True)
    shape = img.shape
    vertex0 = [
        np.random.randint(shape[0] / 2),
        np.random.randint(shape[1]),
        np.random.randint(shape[2]),
    ]
    vertex1 = [
        np.random.randint(shape[0] / 2, shape[0]),
        np.random.randint(shape[1]),
        np.random.randint(shape[2]),
    ]
    radius = np.random.randint(1, 4)
    labels = tube_seg.draw_tube_from_edt(img, vertex0, vertex1, radius)
    line = draw.line_nd(vertex0, vertex1, endpoint=True)
    coords = np.where(labels < 1)
    d_bg = max(shape)
    for pt in np.array(coords).T:
        distance_min = min(np.sum((np.array(line).T - pt) ** 2, axis=1))
        d_bg = min(distance_min, d_bg)
    coords = np.where(labels > 0)
    d_tube = 0
    for pt in np.array(coords).T:
        distance_min = min(np.sum((np.array(line).T - pt) ** 2, axis=1))
        d_tube = max(distance_min, d_tube)

    """
    Verify:

    I. the size of output array
    II. if the output is binary-valued
    III. minimum distance between 0-valued points and the segment line is greater than radius
    IV. maximum distance between 1-valued points and the segment line is less than or equal to radius
    """
    assert labels.shape == shape
    assert np.unique(labels).all() in [0, 1]
    assert d_bg > radius ** 2
    assert d_tube <= radius ** 2

def test_tubes_seg():

    """
    Test if the function maps all the points within the radius of polyline (defined by given vertices) to 1, otherwise 0

    The output array should have the same size of input image and binary values

    Distance between a point and the polyline:
             <= radius (if the point has value 1)
             >  radius (if the point has value 0)
    """
    ngl_session = NeuroglancerSession(url=url, url_segments=url_seg)
    img, _, _ = ngl_session.pull_vertex_list(2, [4], expand=True)
    shape = img.shape
    vertices = np.random.randint(min(shape), size=(4, 3))
    radius = np.random.randint(1, 4)
    labels = tube_seg.tubes_seg(img, vertices, radius)
    point = np.empty((3, 0), dtype=int)
    for i in range(3):
        lines = draw.line_nd(vertices[i], vertices[i + 1], endpoint=True)
        point = np.concatenate((point, np.array(lines)), axis=1)
    coords = np.where(labels < 1)
    d_bg = max(shape)
    for pt in np.array(coords).T:
        distance_min = min(np.sum((point.T - pt) ** 2, axis=1))
        d_bg = min(distance_min, d_bg)
    coords = np.where(labels > 0)
    d_tube = 0
    for pt in np.array(coords).T:
        distance_min = min(np.sum((point.T - pt) ** 2, axis=1))
        d_tube = max(distance_min, d_tube)

    """
    Verify:

    I. the size of output array
    II. if the output is binary-valued
    III. minimum distance between 0-valued points and the polyline is greater than radius
    IV. maximum distance between 1-valued points and the polyline is less than or equal to radius
    """
    assert labels.shape == shape
    assert np.unique(labels).all() in [0, 1]
    assert d_bg > radius ** 2
    assert d_tube <= radius ** 2

def test_tubes_from_paths_bad_inputs():
    """Tests that the tubes_from_paths method raises errors when given bad inputs."""
    sess = NeuroglancerSession(url, 0, url_seg)
    img, bbox, verts = sess.pull_voxel(2, 300, radius=5)  # A valid bbox with data.
    G_paths = sess.get_segments(2, bbox)
    G = G_paths[0]
    paths = G_paths[1]  # valid paths
    bbox = bbox.to_list()
    size = np.subtract(bbox[3:], bbox[:3])
    with pytest.raises(TypeError):
        tube_seg.tubes_from_paths("asdf", paths)
    with pytest.raises(ValueError):
        tube_seg.tubes_from_paths((-1, -1, -1), paths)
    with pytest.raises(TypeError):
        tube_seg.tubes_from_paths(size, "asdf")
    with pytest.raises(TypeError):
        tube_seg.tubes_from_paths(size, [[0, 0, "asdf"]])
    with pytest.raises(TypeError):
        tube_seg.tubes_from_paths(size, paths, radius="asdf")
    with pytest.raises(ValueError):
        tube_seg.tubes_from_paths(size, paths, radius=-1)

def test_tubes_from_paths():
    """Tests that, given valid paths, valid tubes are created."""
    sess = NeuroglancerSession(url, 0, url_seg)
    img, bbox, verts = sess.pull_voxel(2, 300, radius=5)  # A valid bbox with data.
    G_paths = sess.get_segments(2, bbox)
    bbox = bbox.to_list()
    paths = G_paths[1]  # valid paths
    size = np.subtract(bbox[3:], bbox[:3])
    tubes = tube_seg.tubes_from_paths(size, paths)
    assert (tubes != 0).any()

def test_tubes_exact():
    """Tests that exact pixels are filled in."""
    img = np.zeros((10, 10, 10))
    verts = [[5, 5, 0], [5, 5, 10]]
    tubes = tube_seg.tubes_from_paths(img.shape, [verts])
    assert tubes.shape == img.shape
    assert (tubes[5, 5, :] == 1).all()
    for i in range(10):  # set middle column to zero
        tubes[5, 5, i] = 0
    assert (tubes == 0).all()  # now everything should be zero

max_stars_repo_name: neurodata/brainl

max_issues_repo_name: neurodata/brainl

max_forks_repo_name: neurodata/brainl


text:
# Copyright 2016 the GPflow authors.
#
# 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.from __future__ import print_function

from __future__ import print_function
import GPflow
import tensorflow as tf
import numpy as np
import unittest

class TestOptimize(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        rng = np.random.RandomState(0)

        class Quadratic(GPflow.model.Model):
            def __init__(self):
                GPflow.model.Model.__init__(self)
                self.x = GPflow.param.Param(rng.randn(10))

            def build_likelihood(self):
                return -tf.reduce_sum(tf.square(self.x))
        self.m = Quadratic()

    def test_adam(self):
        o = tf.train.AdamOptimizer()
        self.m.optimize(o, maxiter=5000)
        self.assertTrue(self.m.x.value.max() < 1e-2)

    def test_lbfgsb(self):
        self.m.optimize(disp=False)
        self.assertTrue(self.m.x.value.max() < 1e-6)

    def test_feval_counter(self):
        self.m._compile()
        self.m.num_fevals = 0
        for _ in range(10):
            self.m._objective(self.m.get_free_state())
        self.assertTrue(self.m.num_fevals == 10)

class TestNeedsRecompile(unittest.TestCase):
    def setUp(self):
        self.m = GPflow.model.Model()
        self.m.p = GPflow.param.Param(1.0)

    def test_fix(self):
        self.m._needs_recompile = False
        self.m.p.fixed = True
        self.assertTrue(self.m._needs_recompile)

    def test_replace_param(self):
        self.m._needs_recompile = False
        new_p = GPflow.param.Param(3.0)
        self.m.p = new_p
        self.assertTrue(self.m._needs_recompile)

    def test_set_prior(self):
        self.m._needs_recompile = False
        self.m.p.prior = GPflow.priors.Gaussian(0, 1)
        self.assertTrue(self.m._needs_recompile)

    def test_set_transform(self):
        self.m._needs_recompile = False
        self.m.p.transform = GPflow.transforms.Identity()
        self.assertTrue(self.m._needs_recompile)

    def test_replacement(self):
        m = GPflow.model.Model()
        m.p = GPflow.param.Parameterized()
        m.p.p = GPflow.param.Param(1.0)
        m._needs_recompile = False
        # replace Parameterized
        new_p = GPflow.param.Parameterized()
        new_p.p = GPflow.param.Param(1.0)
        m.p = new_p
        self.assertTrue(m._needs_recompile is True)

class KeyboardRaiser:
    """
    This wraps a function and makes it raise a KeyboardInterrupt after some number of calls
    """
    def __init__(self, iters_to_raise, f):
        self.iters_to_raise, self.f = iters_to_raise, f
        self.count = 0

    def __call__(self, *a, **kw):
        self.count += 1
        if self.count >= self.iters_to_raise:
            raise KeyboardInterrupt
        return self.f(*a, **kw)

class TestKeyboardCatching(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        X = np.random.randn(1000, 3)
        Y = np.random.randn(1000, 3)
        Z = np.random.randn(100, 3)
        self.m = GPflow.sgpr.SGPR(X, Y, Z=Z, kern=GPflow.kernels.RBF(3))

    def test_optimize_np(self):
        x0 = self.m.get_free_state()
        self.m._compile()
        self.m._objective = KeyboardRaiser(15, self.m._objective)
        self.m.optimize(disp=0, maxiter=10000, ftol=0, gtol=0)
        x1 = self.m.get_free_state()
        self.assertFalse(np.allclose(x0, x1))

    def test_optimize_tf(self):
        x0 = self.m.get_free_state()
        callback = KeyboardRaiser(5, lambda x: None)
        o = tf.train.AdamOptimizer()
        self.m.optimize(o, maxiter=15, callback=callback)
        x1 = self.m.get_free_state()
        self.assertFalse(np.allclose(x0, x1))

class TestLikelihoodAutoflow(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        X = np.random.randn(1000, 3)
        Y = np.random.randn(1000, 3)
        Z = np.random.randn(100, 3)
        self.m = GPflow.sgpr.SGPR(X, Y, Z=Z, kern=GPflow.kernels.RBF(3))

    def test_lik_and_prior(self):
        l0 = self.m.compute_log_likelihood()
        p0 = self.m.compute_log_prior()
        self.m.kern.variance.prior = GPflow.priors.Gamma(1.4, 1.6)
        l1 = self.m.compute_log_likelihood()
        p1 = self.m.compute_log_prior()

        self.assertTrue(p0 == 0.0)
        self.assertFalse(p0 == p1)
        self.assertTrue(l0 == l1)

class TestName(unittest.TestCase):
    def test_name(self):
        m = GPflow.model.Model(name='foo')
        assert m.name == 'foo'

if __name__ == "__main__":
    unittest.main()

max_stars_repo_name: dalnoguer/my-GPflow

max_issues_repo_name: dalnoguer/my-GPflow

max_forks_repo_name: dalnoguer/my-GPflow


text:
import datetime
import matplotlib.pyplot as plt
import numpy as np
import os

import tensorflow as tf
from keras import backend as K
from keras.engine import Layer, InputSpec
from keras.initializers import RandomNormal
from keras.layers import Input, Dropout, Concatenate
from keras.layers import BatchNormalization, Activation, GaussianNoise
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import UpSampling3D, Conv3D
from keras.models import Model
from keras.optimizers import Adam
from keras.utils import multi_gpu_model
from keras_contrib.layers.normalization.instancenormalization import InstanceNormalization

class My3dResize(Layer):
    def __init__(self, sizes, nn=False, **kwargs):
        super(My3dResize, self).__init__(**kwargs)
        self.sizes = sizes
        self.nn = nn

    def compute_output_shape(self, input_shape):
        output_shape = list(input_shape)
        output_shape[1] = output_shape[1]*self.sizes[0]
        output_shape[2] = output_shape[2]*self.sizes[1]
        output_shape[3] = output_shape[3]*self.sizes[2]
        return tuple(output_shape)

    def call(self, inputs):
        input_shape = inputs.get_shape().as_list()
        output_shape = input_shape.copy()
        output_shape[1] = output_shape[1]*self.sizes[0]
        output_shape[2] = output_shape[2]*self.sizes[1]
        output_shape[3] = output_shape[3]*self.sizes[2]

        # resize rows and columns

        u = tf.reshape(inputs, shape=[-1]+input_shape[1:-2]+[input_shape[-2]*input_shape[-1]])
        if self.nn:
            u = tf.image.resize_images(u, size=output_shape[1:3], method=tf.image.ResizeMethod.BILINEAR)
        else:
            u = tf.image.resize_images(u, size=output_shape[1:3], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        u = tf.reshape(u, shape=[-1]+output_shape[1:-2]+input_shape[-2:])

        # repeat depth-wise
        outputs = K.repeat_elements(u, self.sizes[2], axis=3)
        return outputs

## modified from https://github.com/eriklindernoren/Keras-GAN/tree/master/pix2pix

class My3dPix2Pix():
    def __init__(self, data_loader, savepath='result/pilot', L_weights=(1,100), opt='adam', lrs=(0.0002,0.0), \
                 randomshift=0.1, resoutput=0.0, dropout=0.0, smoothlabel=True, gennoise=0,\
                 fmloss=False, coordconv=False, resizeconv=False, multigpu=None):
        '''
        PatchGAN's receptive field = 70x70x4
        '''
        # Configure data loader
        self.data_loader = data_loader
        self.savepath = savepath
        self.coordconv = coordconv
        self.resizeconv = resizeconv
        self.opt = opt
        self.lr_ini, self.lr_decay = lrs
        self.smoothlabel = smoothlabel
        self.gennoise = gennoise
        self.resoutput=resoutput
        self.dropout = dropout
        self.fmloss = fmloss    # discarded
        self.randomshift=randomshift

        self.reswindow = (self.data_loader.window2[0][0]*self.resoutput,\
                          self.data_loader.window2[0][1]-0.5*(1-self.resoutput)*self.data_loader.window2[0][0])

        ###  from data_loader
        # Input shape --- x,y >=64, depth >=16.
        self.img_shape = tuple(data_loader.img_shape) + (len(data_loader.window1),)
        self.img_rows, self.img_cols, self.depth = data_loader.img_shape
        self.channels = len(data_loader.window1)

        # Calculate output shape of D (PatchGAN)
        patch = max(int(self.img_rows / 2**3), 1)
        dpatch = max(int(self.depth / 2**5), 2)
        self.disc_patch = (patch, patch, dpatch, 1)

        # Number of filters in the first layer of G and D
        self.gf = 64
        self.df = 64

        optimizer = Adam(0.0002, 0.5)
        if self.opt=='adam':
            optimizer = Adam(self.lr_ini, 0.5)

        # Build and compile the discriminator
        self.discriminator, self.discriminator_feat = self.build_discriminator()
        self.discriminator.compile(loss='binary_crossentropy',
            optimizer=optimizer, loss_weights=[0.5],
            metrics=['accuracy'])
        self.discriminator_feat.compile(loss='mae',
            optimizer=optimizer)

        #-------------------------
        # Construct Computational
        #   Graph of Generator
        #-------------------------

        # Build the generator
        self.generator = self.build_generator()

        # Input images and their conditioning images
        img_A = Input(shape=self.img_shape)
        img_B = Input(shape=self.img_shape)

        # By conditioning on B generate a fake version of A
        fake_A = self.generator(img_B)

        # For the combined model we will only train the generator
        self.discriminator.trainable = False
        self.discriminator_feat.trainable = False

        # Discriminators determines validity of translated images / condition pairs
        if self.fmloss:
            # discarded
            '''
            valid = self.discriminator([fake_A, img_B])
            valid_feat = self.discriminator_feat([fake_A, img_B])

            self.combined = Model(inputs=[img_A, img_B], outputs=[valid, fake_A, valid_feat])
            if multigpu is not None:
                self.combined = multi_gpu_model(self.combined, gpus=multigpu)

            self.combined.compile(loss=['binary_crossentropy', 'mae', 'mae'],
                                  loss_weights=list(L_weights),
                                  optimizer=optimizer)
            '''
            valid = self.discriminator([fake_A, img_B])

            self.combined = Model(inputs=[img_A, img_B], outputs=[valid, fake_A])
            if multigpu is not None:
                self.combined = multi_gpu_model(self.combined, gpus=multigpu)

            self.combined.compile(loss=['binary_crossentropy', ssim_mae_loss],
                                  loss_weights=list(L_weights),
                                  optimizer=optimizer)

        else:
            valid = self.discriminator([fake_A, img_B])

            self.combined = Model(inputs=[img_A, img_B], outputs=[valid, fake_A])
            if multigpu is not None:
                self.combined = multi_gpu_model(self.combined, gpus=multigpu)

            self.combined.compile(loss=['binary_crossentropy', 'mae'],
                                  loss_weights=list(L_weights),
                                  optimizer=optimizer)

    def build_generator(self):
        """U-Net Generator"""

        def conv3d(layer_input, filters, kernel_size=(4,4,2), strides=(2,2,2), bn=True):
            """Layers used during downsampling"""
            init = RandomNormal(stddev=0.02)
            d = Conv3D(filters, kernel_size=kernel_size, strides=strides, padding='same', kernel_initializer=init)(layer_input)
            d = LeakyReLU(alpha=0.2)(d)
            if bn:
                #d = BatchNormalization(momentum=0.8)(d)
                d = InstanceNormalization()(d)
            return d

        def deconv3d(layer_input, skip_input, filters, kernel_size=(4,4,2), strides=(2,2,2), dropout_rate=0, bn=True):
            """Layers used during upsampling"""
            if self.resizeconv:
                u = My3dResize(strides)(layer_input)
            else:
                u = UpSampling3D(size=strides)(layer_input)
            init = RandomNormal(stddev=0.02)
            u = Conv3D(filters, kernel_size=kernel_size, strides=1, padding='same', kernel_initializer=init, activation='relu')(u)
            if dropout_rate:
                u = Dropout(dropout_rate)(u)
            if bn:
                #u = BatchNormalization(momentum=0.8)(u)
                u = InstanceNormalization()(u)
            u = Concatenate()([u, skip_input])
            u = Activation('relu')(u)
            return u

        # Image input
        d00 = Input(shape=self.img_shape)

        if self.coordconv:
            d0 = CoordinateChannel3D()(d00)
        else:
            d0 = d00

        n_layers = 7
        encoders = []
        decoders = []

        # Downsampling
        for i in range(n_layers):
            z = 1
            if i < self.depth.bit_length()-1: z = 2
            if i==0:
                encoders.append(conv3d(d0, self.gf, kernel_size=(4,4,z), strides=(2,2,z), bn=False))
            else:
                encoders.append(conv3d(encoders[-1], self.gf*(2**min(i,3)), kernel_size=(4,4,z), strides=(2,2,z)))

        # Upsampling
        for i in range(n_layers-1):
            z = 1
            if i+self.depth.bit_length()>n_layers: z = 2
            if i==0:
                decoders.append(deconv3d(encoders[-(i+1)], encoders[-(i+2)], self.gf*(2**min(n_layers-2-i,3)),
                                         kernel_size=(4,4,z), strides=(2,2,z)))
            else:
                decoders.append(deconv3d(decoders[-1], encoders[-(i+2)], self.gf*(2**min(n_layers-2-i,3)),
                                         kernel_size=(4,4,z), strides=(2,2,z), dropout_rate=self.dropout))

        if self.resizeconv:
            u7 = My3dResize((2,2,2))(decoders[-1])
        else:
            u7 = UpSampling3D(size=2)(decoders[-1])
        init = RandomNormal(stddev=0.02)
        output_img = Conv3D(self.channels, kernel_size=(4,4,2), strides=1, padding='same', kernel_initializer=init, activation='tanh')(u7)

        return Model(d00, output_img)

    def build_discriminator(self):
        def d_layer(layer_input, filters, kernel_size=(4,4,2), strides=(2,2,2), bn=True):
            """Discriminator layer"""
            init = RandomNormal(stddev=0.02)
            d = Conv3D(filters, kernel_size=kernel_size, strides=strides, padding='same', kernel_initializer=init)(layer_input)
            d = LeakyReLU(alpha=0.2)(d)
            if bn:
                #d = BatchNormalization(momentum=0.8)(d)
                d = InstanceNormalization()(d)
            return d

        img_A = Input(shape=self.img_shape)
        img_B = Input(shape=self.img_shape)

        # Concatenate image and conditioning image by channels to produce input
        combined_imgs = Concatenate(axis=-1)([img_A, img_B])

        ## testing
        #combined_imgs = GaussianNoise(0.02)(combined_imgs)

        if self.coordconv:
            combined_imgs = CoordinateChannel3D()(combined_imgs)

        n_d_layer = 4
        d_layers = []
        df = self.df
        dout = self.depth
        for i in range(n_d_layer):
            if i==0:
                d_layers.append(d_layer(combined_imgs, df, kernel_size=(4,4,2), strides=(2,2,2)))
            else:
                z = 2
                s = 2
                if dout == 2:
                    z = 1
                if i == n_d_layer-1:
                    s = 1
                df = min(df*2, self.df*8)
                d_layers.append(d_layer(d_layers[-1], df, kernel_size=(4,4,z), strides=(s,s,z)))
            dout = max(2, int(0.5*dout))

        init = RandomNormal(stddev=0.02)
        features = Conv3D(1, kernel_size=(4,4,1), strides=1, padding='same', kernel_initializer=init)(d_layers[-1])
        validity = Activation('sigmoid')(features)

        return Model([img_A, img_B], validity), Model([img_A, img_B], features)

    def convert_resoutput(self, imgs_A, imgs_B):
        ## imgs 5D tensor (batch, rows, cols, depth, channels)
        if self.resoutput:
            new_A = []
            for i in range(imgs_B.shape[0]):
                a = rWND(255*(0.5*imgs_A[i]+0.5),self.data_loader.window2)
                b = rWND(255*(0.5*imgs_B[i]+0.5),self.data_loader.window1)
                c = a-b
                #c[b>100] = 0
                #c = WND(c,self.reswindow)
                M = np.max(c)
                m = np.min(c)
                if M==m:
                    c = np.zeros(c.shape)
                else:
                    c = 2.*(c-m)/(M-m) - 1
                new_A.append(c)
            imgs_A = np.array(new_A)
        return imgs_A

    def invert_resoutput(self, fake_A, imgs_B):
        if self.resoutput:
            '''
            new_A = []
            for i in range(fake_A.shape[0]):
                a = rWND(255*(0.5*fake_A[i]+0.5),self.reswindow)
                b = rWND(255*(0.5*imgs_B[i]+0.5),self.data_loader.window1)
                c = a+b
                c = WND(c,self.data_loader.window2)
                new_A.append(c/127.5 - 1.)
            fake_A = np.array(new_A)
            '''
            pass
        return fake_A

    def generate_noise(self, mode, imgs_A, imgs_B):
        if mode==0:
            pass
        elif mode==1:
            r = np.random.uniform(0,1)
            if r>=0.5:
                row,col,dep,ch= imgs_A.shape[1:]
                mean = 0
                sigma = 0.05
                gauss = np.random.normal(mean,sigma,(row,col,dep,ch))
                gauss = gauss.reshape(row,col,dep,ch)
                gauss = np.expand_dims(gauss, axis=0)
                imgs_A = imgs_A + gauss
                imgs_B = imgs_B + gauss
            else:
                pass
        elif mode==2:
            r = np.random.uniform(0,1)
            r2 = np.random.uniform(0,1)
            if r>=0.5:
                row,col,dep,ch= imgs_A.shape[1:]
                mean = 0
                sigma = 0.05
                gauss = np.random.normal(mean,sigma,(row,col,dep,ch))
                gauss = gauss.reshape(row,col,dep,ch)
                gauss = np.expand_dims(gauss, axis=0)
                imgs_A = imgs_A + gauss
                imgs_B = imgs_B + gauss
            else:
                pass
            if r2>=0.75:
                imgs_A = gaussian_filter(imgs_A, sigma=0.05)
                imgs_B = gaussian_filter(imgs_B, sigma=0.05)
            else:
                pass
        else:
            pass

        return imgs_A, imgs_B

    def train(self, epochs, batch_size=1, sample_interval=200, model_interval=-1):

        start_time = datetime.datetime.now()

        # Adversarial loss ground truths
        patch_shape = (batch_size,) + self.disc_patch
        if self.smoothlabel:
            valid = np.ones(patch_shape) - 0.1*np.random.rand(*patch_shape)
        else:
            valid = np.ones(patch_shape)

        fake = np.zeros(patch_shape)

        # log file
        f = open(os.path.join(self.savepath, 'log.txt'), 'w')
        f.close()

        for epoch in range(epochs):
            if epoch<10:
                lr = self.lr_ini
            else:
                lr = self.lr_ini/(1+(epoch-10)*self.lr_decay)
            K.set_value(self.discriminator.optimizer.lr, lr)
            K.set_value(self.combined.optimizer.lr, lr)

            for batch_i, (imgs_A, imgs_B) in enumerate(self.data_loader.load_batch(batch_size=batch_size)):

                # ---------------------
                #  Train Discriminator
                # ---------------------
                imgs_A = self.convert_resoutput(imgs_A, imgs_B)

                # randomshift
                non = lambda s: s if s<0 else None
                mom = lambda s: max(0,s)

                shift_A = np.full(imgs_A.shape, -1.)
                shift_B = np.full(imgs_A.shape, -1.)
                sx = int(self.randomshift*shift_A.shape[2])
                sy = int(self.randomshift*shift_A.shape[1])
                for i in range(shift_A.shape[0]):
                    ox = np.random.randint(2*sx+1, size=1)[0] - sx
                    oy = np.random.randint(2*sy+1, size=1)[0] - sy
                    shift_A[i,mom(oy):non(oy),mom(ox):non(ox),:,:] = imgs_A[i,mom(-oy):non(-oy),mom(-ox):non(-ox),:,:]
                    shift_B[i,mom(oy):non(oy),mom(ox):non(ox),:,:] = imgs_B[i,mom(-oy):non(-oy),mom(-ox):non(-ox),:,:]
                imgs_A = shift_A
                imgs_B = shift_B

                imgs_A, imgs_B = self.generate_noise(self.gennoise, imgs_A, imgs_B)

                # Condition on B and generate a translated version
                fake_A = self.generator.predict(imgs_B)

                # Train the discriminators (original images = real / generated = Fake)
                d_loss_real = self.discriminator.train_on_batch([imgs_A, imgs_B], valid)
                d_loss_fake = self.discriminator.train_on_batch([fake_A, imgs_B], fake)
                d_loss = np.add(d_loss_real, d_loss_fake)

                # -----------------
                #  Train Generator
                # -----------------

                # Train the generators
                '''
                if self.fmloss:
                    valid_feat = self.discriminator_feat.predict([imgs_A, imgs_B])
                    g_loss = self.combined.train_on_batch([imgs_A, imgs_B], [valid, imgs_A, valid_feat])
                else:
                    g_loss = self.combined.train_on_batch([imgs_A, imgs_B], [valid, imgs_A])
                '''
                g_loss = self.combined.train_on_batch([imgs_A, imgs_B], [valid, imgs_A])

                elapsed_time = datetime.datetime.now() - start_time

                # Plot the progress
                newlog = "[Epoch %d/%d] [Batch %d/%d] [D loss: %f, D fake: %f] [G loss: %f] time: %s" % (
                    epoch+1, epochs, batch_i+1, self.data_loader.n_batches,
                    d_loss[0], d_loss_fake[0], #100*d_loss[1],
                    g_loss[0],
                    elapsed_time
                )

                print(newlog)
                f = open(os.path.join(self.savepath, 'log.txt'), 'a')
                f.write(newlog+'\n')
                f.close()

                # If at save interval => save generated image samples
                if (batch_i+1) % sample_interval == 0:
                    self.sample_images(epoch, batch_i+1)
            # save weights
            if model_interval>0:
                if (epoch+1) % model_interval == 0:
                    self.save_weights('{}'.format(str(epoch+1)))

        # final sample image & save weights
        self.sample_images(epochs, 'final')
        if model_interval>0:
            self.save_weights('final{}'.format(str(epochs)))

    def predict_on_batch(self, imgs_A, imgs_B):
        fake_A = self.generator.predict(imgs_B)
        gen_imgs = np.concatenate([imgs_B, fake_A, imgs_A])
        return gen_imgs

    def sample_images(self, epoch, batch_i):
        r, c = 3, 3

        imgs_A, imgs_B = self.data_loader.load_data(batch_size=3)
        fake_A = np.concatenate([self.generator.predict(np.expand_dims(x,axis=0)) for x in imgs_B], axis=0)
        fake_A = self.invert_resoutput(fake_A, imgs_B)

        gen_imgs = np.concatenate([imgs_B[:,:,:,0,-1], fake_A[:,:,:,0,-1], imgs_A[:,:,:,0,-1]])

        # Rescale images 0 - 1
        gen_imgs = 0.5 * gen_imgs + 0.5

        titles = ['Condition', 'Generated', 'Original']
        plt.style.use('default')
        fig, axs = plt.subplots(r, c)
        cnt = 0
        for i in range(r):
            dr = imgs_B.shape[1]//4+imgs_B.shape[1]//8
            dc = imgs_B.shape[2]//4+imgs_B.shape[2]//8
            pr = np.random.choice(imgs_B.shape[1]//2-dr)+imgs_B.shape[1]//4
            pc = np.random.choice(imgs_B.shape[2]//2-dc)+imgs_B.shape[2]//4

            for j in range(c):
                axs[j,i].imshow(gen_imgs[cnt], cmap='gray', vmin=0, vmax=1)
                axs[j,i].set_title(titles[i])
                axs[i,j].axis([pc,pc+dc,pr+dr,pr])
                cnt += 1
        samplepath = self.make_directory('samples')
        fig.savefig(os.path.join(samplepath, '{}_{}.png'.format(epoch, batch_i)))
        plt.close()

    def make_directory(self, dirname):
        dirpath = os.path.join(self.savepath, dirname)
        if not os.path.isdir(dirpath):
            os.mkdir(dirpath)
        return dirpath

    def load_weights(self, weightfile, summary=True):
        loadweightspath = os.path.join(self.savepath, 'models', '{}.h5'.format(weightfile))
        self.combined.load_weights(loadweightspath)
        if summary:
            self.combined.summary()

    def load_final_weights(self, *args, **kwargs):
        wdir = os.path.join(self.savepath, 'models')
        wlist = [os.path.splitext(x)[0] for x in os.listdir(wdir) if x.lower().endswith('.h5')]
        wlist.sort()

        s = None
        for x in wlist:
            if 'final' in x:
                s = x
        if s is None:
            s = wlist[-1]

        self.load_weights(s,*args, **kwargs)

    def save_weights(self, weightfile):
        modelpath = self.make_directory('models')
        self.combined.save_weights(os.path.join(modelpath, '{}.h5'.format(weightfile)))
        print('saved {}.h5 to {}'.format(weightfile, modelpath))

## not used
def ssim_mae_loss(y_true, y_pred):
    return 0.15*losses.mean_absolute_error(y_true, y_pred) + 0.85*ssim_loss(y_true, y_pred)

def ssim_loss(y_true, y_pred):
    # ssim for 3d data 구현
    ashape = y_true.get_shape().as_list()

    ut = tf.transpose(y_true[:,:,:,:,0], perm=[0,3,1,2])
    #ut = tf.reshape(ut, shape=[-1]+ashape[1:3])
    ut = tf.expand_dims(ut, axis=-1)
    up = tf.transpose(y_pred[:,:,:,:,0], perm=[0,3,1,2])
    #up = tf.reshape(up, shape=[-1]+ashape[1:3])
    up = tf.expand_dims(up, axis=-1)

    return tf.reduce_mean(1-tf.image.ssim_multiscale(ut, up, 2.0))

# modified from "An Intriguing Failing of Convolutional Neural Networks and the CoordConv Solution" from Uber research
# https://arxiv.org/abs/1807.03247

class _CoordinateChannel(Layer):
    def __init__(self, rank,
                 data_format=None,
                 **kwargs):
        super(_CoordinateChannel, self).__init__(**kwargs)

        self.rank = rank
        self.data_format = K.image_data_format() if data_format is None else data_format
        self.axis = 1 if K.image_data_format() == 'channels_first' else -1

        self.input_spec = InputSpec(min_ndim=2)

    def build(self, input_shape):
        assert len(input_shape) >= 2
        input_dim = input_shape[self.axis]

        self.input_spec = InputSpec(min_ndim=self.rank + 2,
                                    axes={self.axis: input_dim})
        self.built = True

    def call(self, inputs):
        input_shape = K.shape(inputs)

        if self.rank == 2:

            input_shape = [input_shape[i] for i in range(4)]
            batch_shape, dim1, dim2, channels = input_shape

            xx_ones = K.ones(K.stack([batch_shape, dim2]), dtype='int32')
            xx_ones = K.expand_dims(xx_ones, axis=-1)

            xx_range = tf.tile(K.expand_dims(K.arange(0, dim1), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=1)
            xx_channels = K.batch_dot(xx_ones, xx_range, axes=[2, 1])
            xx_channels = K.expand_dims(xx_channels, axis=-1)
            xx_channels = K.permute_dimensions(xx_channels, [0, 2, 1, 3])

            yy_ones = K.ones(K.stack([batch_shape, dim1]), dtype='int32')
            yy_ones = K.expand_dims(yy_ones, axis=1)

            yy_range = tf.tile(K.expand_dims(K.arange(0, dim2), axis=0),
                              K.stack([batch_shape, 1]))
            yy_range = K.expand_dims(yy_range, axis=-1)

            yy_channels = K.batch_dot(yy_range, yy_ones, axes=[2, 1])
            yy_channels = K.expand_dims(yy_channels, axis=-1)
            yy_channels = K.permute_dimensions(yy_channels, [0, 2, 1, 3])

            xx_channels = K.cast(xx_channels, K.floatx())
            xx_channels = xx_channels / K.cast(dim1 - 1, K.floatx())
            xx_channels = (xx_channels * 2) - 1.

            yy_channels = K.cast(yy_channels, K.floatx())
            yy_channels = yy_channels / K.cast(dim2 - 1, K.floatx())
            yy_channels = (yy_channels * 2) - 1.

            outputs = K.concatenate([inputs, xx_channels, yy_channels], axis=-1)

        if self.rank == 3:

            input_shape = [input_shape[i] for i in range(5)]
            batch_shape, dim1, dim2, dim3, channels = input_shape

            xx_ones = K.ones(K.stack([batch_shape, dim3]), dtype='int32')
            xx_ones = K.expand_dims(xx_ones, axis=-1)

            xx_range = tf.tile(K.expand_dims(K.arange(0, dim2), axis=0),
                              K.stack([batch_shape, 1]))
            xx_range = K.expand_dims(xx_range, axis=1)

            xx_channels = K.batch_dot(xx_ones, xx_range, axes=[2, 1])
            xx_channels = K.expand_dims(xx_channels, axis=-1)
            xx_channels = K.permute_dimensions(xx_channels, [0, 2, 1, 3])

            xx_channels = K.expand_dims(xx_channels, axis=1)
            xx_channels = tf.tile(xx_channels,
                                 [1, dim1, 1, 1, 1])

            yy_ones = K.ones(K.stack([batch_shape, dim2]), dtype='int32')
            yy_ones = K.expand_dims(yy_ones, axis=1)

            yy_range = tf.tile(K.expand_dims(K.arange(0, dim3), axis=0),
                              K.stack([batch_shape, 1]))
            yy_range = K.expand_dims(yy_range, axis=-1)

            yy_channels = K.batch_dot(yy_range, yy_ones, axes=[2, 1])
            yy_channels = K.expand_dims(yy_channels, axis=-1)
            yy_channels = K.permute_dimensions(yy_channels, [0, 2, 1, 3])

            yy_channels = K.expand_dims(yy_channels, axis=1)
            yy_channels = tf.tile(yy_channels,
                                 [1, dim1, 1, 1, 1])

            zz_range = tf.tile(K.expand_dims(K.arange(0, dim1), axis=0),
                              K.stack([batch_shape, 1]))
            zz_range = K.expand_dims(zz_range, axis=-1)
            zz_range = K.expand_dims(zz_range, axis=-1)

            zz_channels = tf.tile(zz_range,
                                 [1, 1, dim2, dim3])
            zz_channels = K.expand_dims(zz_channels, axis=-1)

            xx_channels = K.cast(xx_channels, K.floatx())
            xx_channels = xx_channels / K.cast(dim2 - 1, K.floatx())
            xx_channels = xx_channels * 2 - 1.

            yy_channels = K.cast(yy_channels, K.floatx())
            yy_channels = yy_channels / K.cast(dim3 - 1, K.floatx())
            yy_channels = yy_channels * 2 - 1.

            zz_channels = K.cast(zz_channels, K.floatx())
            zz_channels = zz_channels / K.cast(dim1 - 1, K.floatx())
            zz_channels = zz_channels * 2 - 1.

            outputs = K.concatenate([inputs, zz_channels, xx_channels, yy_channels],axis=-1)

        return outputs

    def compute_output_shape(self, input_shape):
        assert input_shape and len(input_shape) >= 2
        assert input_shape[self.axis]

        channel_count = self.rank

        output_shape = list(input_shape)
        output_shape[self.axis] = input_shape[self.axis] + channel_count
        return tuple(output_shape)

    def get_config(self):
        config = {
            'rank': self.rank,
            'data_format': self.data_format
        }
        base_config = super(_CoordinateChannel, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class CoordinateChannel3D(_CoordinateChannel):
    def __init__(self,
                 data_format=None,
                 **kwargs):
        super(CoordinateChannel3D, self).__init__(
            rank=3,
            data_format=data_format,
            **kwargs
        )

    def get_config(self):
        config = super(CoordinateChannel3D, self).get_config()
        config.pop('rank')
        return config