#!/usr/bin/env python

##########################################################################################################
# Modulo Bioestadistica - 2015 de la Universidad del Comahue. Centro Regional Bariloche
#http://crubweb.uncoma.edu.ar/
# Dr. <NAME>
# email: <EMAIL>
# licence: MIT. http://opensource.org/licenses/MIT

#Ejemplo del libro Sokal, Introduccion a la bioestadistica, Reverte, 2002. version en ingles  p.79
#Distribucion normal o gaussiana
#Ejecutar:
#python 09_introNormal.py
##########################################################################################################

# Importo librerias
import numpy as np
import pylab as P
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import math
import scipy.stats as stats

#Generador de coef binomial
def normpdf(x, mean, sd):
    var = float(sd)**2
    pi = 3.1415926
    denom = (2*pi*var)**.5
    num = math.exp(-(float(x)-float(mean))**2/(2*var))
    return num/denom

#Generador de coef binomial
def binCoefGen(k):
    f = math.factorial
    coef = []
    print k
    for i in range(k+1):
        #print i
        a = f(k)/f(i)/f(k-i)
        print a
        coef.append(a)
    return coef

#Defino el k
print "Efecto del k -> inf. Probar k = 5, 10, 50, 150. criterio pqk > 3"
k = 150

#calculo los bin
bins = np.array([x for x in range(0, k+1)])

#Defino la probabilidad de ocurrencia
p = 0.5

columna2p = np.array([p**(k-x) for x in range(0, k+1)])
print columna2p

#calculo la probabilidad de no ocurrencia
q = 1-p
columna3q = np.array([q**x for x in range(0, k+1)])
print columna3q
print "approx: ", p*q*k
bincoef = np.array(binCoefGen(k))

columna4relfrec = np.array([columna2p[x]*columna3q[x]*bincoef[x] for x in range(0, k+1)])
print columna4relfrec

totalMuestras = 1
frecAbsTeor = np.array([columna4relfrec[x]*totalMuestras for x in range(0, k+1)])
print frecAbsTeor
mean= k*p
sigma = math.sqrt(k)*p*q
print "mean teorico:", mean
print "sigma teorico:", sigma
#Defino titulos
plt.ylabel("Funcion densidad de probabilidad Normal")
plt.xlabel("Y")
plt.title("Distribuciones normales")

#Defino ancho de barra
width = 1.0

#Defino 2 graficas de barras
normalWLib1stats = stats.norm.pdf(bins,mean,2*sigma)
normalWLib2mlab = mlab.normpdf(bins,mean,2*sigma)

#Check with other library

means = [4,8,8]
stds = [2,2,1]

normVals = []
normVals.append([])
normVals.append([])
normVals.append([])

for k in range(len(means)):
	for i in range(len(bins)):
		normVals[k].append(normpdf(bins[i],means[k],stds[k]))

colors = ["red","blue","green"]
labels = ["Propongan 1","Propongan 2","Propongan 3"]

#labels = ["Norm mu=%d sigma=%d"%(means[0],stds[0]),"Norm mu=%d sigma=%d"%(means[1],stds[1]),"Norm mu=%d sigma=%d"%(means[2],stds[2])]

for k in range(len(means)):
	plt.plot(bins,normVals[k],lw = 2, color = colors[k], label=labels[k])
plt.xlim(0,15)

#Ubico leyenda
plt.legend( loc='upper left', numpoints = 1 )

#Dibujo
plt.show()


import unittest
from unittest import TestCase

from escnn.group import CyclicGroup
from escnn.group import DihedralGroup

import numpy as np
import scipy.sparse as sparse

class TestComputeRegularRepresentations(TestCase):

    def test_dihedral_rr_odd(self):
        dg = DihedralGroup(9)
        self.dihedral_rr_eval(dg, dg.representations['regular'])

    def test_dihedral_rr_even(self):
        dg = DihedralGroup(10)
        self.dihedral_rr_eval(dg, dg.representations['regular'])

    def test_dihedral_rr_large(self):
        dg = DihedralGroup(16)
        self.dihedral_rr_eval(dg, dg.representations['regular'])

    def test_dihedral_rr_small(self):
        dg = DihedralGroup(2)
        self.dihedral_rr_eval(dg, dg.representations['regular'])

    def test_cyclic_rr_odd(self):
        cg = CyclicGroup(11)
        self.cyclic_rr_eval(cg, cg.representations['regular'])

    def test_cyclic_rr_even(self):
        cg = CyclicGroup(10)
        self.cyclic_rr_eval(cg, cg.representations['regular'])

    def test_cyclic_rr_large(self):
        cg = CyclicGroup(20)
        self.cyclic_rr_eval(cg, cg.representations['regular'])

    def test_cyclic_rr_small(self):
        cg = CyclicGroup(2)
        self.cyclic_rr_eval(cg, cg.representations['regular'])

    def cyclic_rr_eval(self, cg, rr):
        # rr = cg.representations['regular']

        # np.set_printoptions(precision=4, suppress=True)
        # print('Change of Basis')
        # print(rr.change_of_basis)
        # print('Change of Basis Inv')
        # print(rr.change_of_basis_inv)
        # print('RR')
        # n = cg.order
        # for i in range(n):
        #     print(rr(i * 2 * np.pi / n))

        D = rr.change_of_basis
        D_inv = rr.change_of_basis_inv
        for i, element in enumerate(cg.elements):

            rho_i = np.zeros([cg.order(), cg.order()])

            for k in range(cg.order()):
                rho_i[(i + k) % cg.order(), k] = 1.0

            # Build the direct sum of the irreps for this element
            blocks = []
            for irrep in rr.irreps:
                repr = cg.irrep(*irrep)(element)
                blocks.append(repr)

            P = sparse.block_diag(blocks, format='csc')
            R = D @ P @ D_inv
            self.assertTrue(np.allclose(R, rho_i), f"{element}:\n{R}\n!=\n{rho_i}\n")
            self.assertTrue(np.allclose(rr(element), rho_i), f"{element}:\n{rr(element)}\n!=\n{rho_i}\n")

    def dihedral_rr_eval(self, dg, rr):

        # rr = dg.representations['regular']

        # np.set_printoptions(precision=2, suppress=True)
        # print('Change of Basis')
        # print(rr.change_of_basis)
        # print('Change of Basis Inv')
        # print(rr.change_of_basis_inv)
        # print('RR')
        # n = dg.rotation_order
        # for i in range(n):
        #     print(rr((0, i * 2 * np.pi / n)))
        # for i in range(n):
        #     print(rr((1, i * 2 * np.pi / n)))

        D = rr.change_of_basis
        D_inv = rr.change_of_basis_inv

        # np.set_printoptions(precision=3, threshold=10*rr.size**2, suppress=True, linewidth=25*rr.size + 5)

        for i, element in enumerate(dg.elements):

            rho_i = np.zeros([dg.order(), dg.order()])

            f = -1 if element.to('int')[0] else 1
            # r = int(np.round(element[1] * dg.rotation_order / (2 * np.pi)))
            r = element.to('int')[1]

            for k in range(dg.rotation_order):
                rho_i[dg.rotation_order * element.to('int')[0] + ((r + k * f) % dg.rotation_order), k] = 1.0
            for k in range(dg.rotation_order):
                rho_i[dg.rotation_order * (1 - element.to('int')[0]) + ((r + k * f) % dg.rotation_order), dg.rotation_order + k] = 1.0

            # Build the direct sum of the irreps for this element
            blocks = []
            for irrep in rr.irreps:
                repr = dg.irrep(*irrep)(element)
                blocks.append(repr)

            P = sparse.block_diag(blocks, format='csc')
            R = D @ P @ D_inv
            self.assertTrue(np.allclose(R, rho_i), f"{element}:\n{R}\n!=\n{rho_i}\n")
            self.assertTrue(np.allclose(rr(element), rho_i), f"{element}:\n{rr(element)}\n!=\n{rho_i}\n")

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


###__________________________ SymPy - Mechanics ____________________________###

##_________ Reference System _________##
# El objeto primordial que vamos a manejar van a ser los sistemas de referencia.
# Podremos definir relaciones geométricas entre ellos y de esta forma las transformaciones
# de vectores entre un sistema y otro serán triviales.

# La manera usual de empezar a trabajar con SymPy es importar la función init_session:
from sympy import init_session, Symbol, symbols, pi, I, E, cos, sin, exp, tan, simplify, expand, factor, collect, apart, cancel, expand_trig, diff, Derivative, Function, integrate, limit, series, Eq, solve, dsolve, Matrix, N
# from sympy import init_session
init_session(use_latex=True)

# Todo lo que necesitamos está en sympy.physics.mechanics, incluyendo la clase ReferenceFrame.
# Nada más crear un sistema de referencia podemos acceder a sus versores unitarios: x, y y z.

from sympy.physics.mechanics import ReferenceFrame

# A = ReferenceFrame("A")
# print(A.x)

# Y para definir vectores solo tenemos que multiplicar cada componente por su versor:
# print(2 * A.x - 1 * A.y)

# De ahora en adelante, para trabajar como si nos enfrentáramos a un problema de
# la escuela, vamos a hacer dos cosas:
    # [-] Definir un sistema inercial 1 del que partir, para así poder referir todos
    # los demás sistemas a él.
    # [-] Que los versores de ese sistema sean 𝑖,𝑗,𝑘

# A = ReferenceFrame("1", latexs=['\mathbf{i}', '\mathbf{j}', '\mathbf{k}'])
# print(A.x + A.y + A.z)

# Y para no tener que hacerlo siempre, un pequeño truco de magia:
# Definimos nuestra propia clase para que los versores sean IJK
#------------------------------------------------------------------------------
class IJKReferenceFrame(ReferenceFrame):
    def __init__(self, name):
        super().__init__(name, latexs=['\mathbf{%s}_{%s}' % (idx, name) for idx in ("i", "j", "k")])
        self.i = self.x
        self.j = self.y
        self.k = self.z
#------------------------------------------------------------------------------
A = IJKReferenceFrame("1")
print(A.x + A.y + A.z)

##_________ Vectorial Algebra _________##
# Nuestros vectores funcionan también con símbolos, y podemos realizar las operaciones
# de producto escalar y producto vectorial con ellos.
R, V = symbols('R, V', positive=True)
r1 = R * (A.x + A.y + A.z)
v1 = V * (A.x - 2 * A.z)

from sympy.physics.mechanics import dot, cross

print(r1.dot(v1))
print(dot(r1, v1))
print(r1 & v1)

print(r1.cross(v1))
print(cross(r1, v1))
print(r1 ^ v1)

# Podemos hallar también la norma de los vectores con su método magnitude e incluso
# normalizarlos con normalize:
print((r1 ^ v1).magnitude())
print((r1 ^ v1).normalize())

##_________ Exercise _________##
# Usando directamente la fórmula para la derivada en ejes móviles:
    # $$\left(\frac{\operatorname{d}\!\mathbf{a}}{\operatorname{d}\!t}\right)_1
    # = \left(\frac{\operatorname{d}\!\mathbf{a}}{\operatorname{d}\!t}\right)_0
    # + \mathbf{\omega}_{01}\! \times \mathbf{a}$$
# Calcula la derivada del vector de posición $R \mathbf{i}_0$, siendo $A_0$ un
# sistema de referencia que gira respecto al inercial con velocidad angular
# $\mathbf{\omega}_{01}=\Omega \mathbf{k}_0$. **¿Cuál es el módulo de la derivada?**
R, Omega = symbols('R, Omega', positive=True)
A0 = IJKReferenceFrame('0')
a = R * A0.i
omega01 = Omega * A0.k
da = omega01 ^ a # cross producut
print(da.magnitude())
# Si no especificaste `positive=True` vas a ver algo como sqrt(omega^2*R^2).
# Debería haber una forma de simplificar esta expresión _a posteriori_, pero de
# momento no funciona del todo bien.

##_________ Relative Movement _________##
# ¿A quién no le gusta multiplicar matrices de rotación? Para esa minoría que lo
# detesta, existe SymPy. Para ello debemos especificar la orientación de nuestros
# sistemas de referencia usando el *método orient*, y recuperaremos la matriz de
# cosenos directores usando el *método dcm*.
A1 = IJKReferenceFrame("1")
A0 = IJKReferenceFrame("0")
phi = symbols('phi')
A0.orient(A1, 'Axis', [phi, A1.z])  # Rotación phi alrededor del eje A1.z
print(A0.dcm(A1))                          # "Direct Cosine Matrix"
# Usando el argumento `Axis` hemos especificado que rotamos el sistema un ángulo
# especificado alrededor de un eje. Otros métodos son:
    # * `Body`: se especifican los tres ángulos de Euler.
    # * `Space`: igual que `Body`, pero las rotaciones se aplican en orden inverso.
    # * `Quaternion`: utilizando cuaternios, rotación alrededor de un vector unitario
    # $\lambda$ una cantidad $\theta$.

##_________ Different Reference System _________##
# Para expresar un vector en otro sistema de referencia, no hay más que usar los
# métodos express o to_matrix:
print(A0.x.express(A1))
print(A0.x.to_matrix(A1))

##_________ Dynamic Symbols (time dependent) _________##
# Si queremos especificar que un símbolo puede variar con el tiempo, hay que usar
# la función dynamicsymbols:
from sympy.physics.mechanics import dynamicsymbols

alpha = dynamicsymbols('alpha')
# Y pedir su derivada con el método diff:
print(alpha.diff())

##_________ Exercise1 _________##
print('_______________Exercise1___________________')
# from notebook completos/041-SymPy
## Obtener la matriz de rotación de la pala B respecto a los ejes A1. ##
print('______Rot Matrix_________')
A = IJKReferenceFrame("A")
A1 = IJKReferenceFrame("A1")
psi = dynamicsymbols('psi')
A1.orient(A, 'Axis', [psi, A.z])
print('A1 DCM A --> =', A1.dcm(A))  # T_{A1A}

A2 = IJKReferenceFrame("A2")
beta = dynamicsymbols('beta')
A2.orient(A1, 'Axis', [beta, -A1.y])
print('A2 DCM A1 --> =', A2.dcm(A1))  # T_{A2A1}

A3 = IJKReferenceFrame("A3")
zeta = dynamicsymbols('zeta')
A3.orient(A2, 'Axis', [zeta, A2.z])
print('A3 DCM A1 --> =',A3.dcm(A1))  # T_{A3A1}

B = IJKReferenceFrame("B")
theta = dynamicsymbols('theta')
B.orient(A3, 'Axis', [theta, A3.x])
print('B DCM A3 --> =', B.dcm(A3))  # T_{BA3}

print('B DCM A2 --> =',B.dcm(A2))

print('B DCM A1 --> =',B.dcm(A1))

## Angular Velocity ##
print('______Ang Vel______')
# También podemos hallar la velocidad angular de un sistema respecto a otro
# usando el método ang_vel_in:
print(B.ang_vel_in(A2))
print(B.ang_vel_in(A))
print(B.ang_vel_in(A).express(A))

## Derivative in moving axis ##
print('______Derivative in moving axis______')
# Hacer una derivada con la fórmula lo hace cualquiera, pero SymPy puede
# encargarse automáticamente.
v1 = A1.x
dv1 = v1.diff(symbols('t'), A)
print(dv1.to_matrix(A1))
print((dv1 & A1.j).simplify())
print('_______________End Exercise1___________________')

##_________ Puntos, velocidades y la rueda que no desliza _________##
# El último paso que nos queda para completar la cinemática es la posibilidad de
# definir puntos en sólidos y aplicar su campo de velocidades. SymPy también
# permite esto, y para ello no tenemos más que importar la clase Point.
from sympy.physics.mechanics import Point
O = Point("O")
# Para trabajar como lo haríamos en la escuela, vamos a especificar que O es el
# origen de A, y para eso vamos a imponer que su velocidad es cero con el método set_vel:
O.set_vel(A, 0)
# Para definir nuevos puntos, podemos utilizar el método locate_new:
e_b = symbols('e_b')
E_b = O.locatenew('E_b', e_b * A1.x)
# Y para obtener vectores de un punto a otro, el método pos_from:
print(E_b.pos_from(O))
# La notación de este paquete está influenciada por el libro
#  <NAME>. & <NAME>. "Dynamics, Theory and Applications"

# Por último, el **campo de velocidades de un sólido rígido** se formula usando
# el método `v2pt_theory`.
    # $$v^P_A = v^O_A + \omega_{A_1 A} \times \mathbf{OP}$$
# Este método pertenece *al punto del cual queremos conocer la velocidad* y
# recibe tres parámetros:
    # * `O`, punto de velocidad conocida respecto a A
    # * `A`, sistema de referencia donde queremos calcular la velocidad
    # * `A1`, sistema de referencia donde están fijos ambos puntos (_sistema de arrastre_)
print(E_b.v2pt_theory(O, A, A1))

##_________ Exercise2 _________##
print('_______________Exercise2___________________')
# from notebook completos/041-SymPy
# ¡Halla la velocidad y la aceleración de P!

# Creamos nuestros sistemas de referencia
A1 = IJKReferenceFrame('1')
A0 = IJKReferenceFrame('0')
A2 = IJKReferenceFrame('2')

# Creamos los símbolos dinámicos necesarios
xi, theta = dynamicsymbols('xi, theta')

# Orientamos los sistemas de referencia
A0.orient(A1, 'Axis', [0, A1.k])  # A0 no gira respecto a A1
A2.orient(A0, 'Axis', [theta, A0.k])
print('A2 DCM A1 --> =', A2.dcm(A1))

# Creamos el punto C, centro del disco, y especificamos su velocidad
# respecto a A1
C = Point('C')
C.set_vel(A1, xi.diff() * A1.x)

# Localizamos el punto P, punto fijo del disco, respecto a C, en
# el sistema A2 (que gira solidariamente con el disco)
R = symbols('R')
P = C.locatenew('P', -R * A2.j)
print(P.pos_from(C))

# Hallamos la velocidad de P en A1, expresada en A0
# ¡Con esta llamada ya estamos diciendo que C y P son fijos en A2!
print(P.v2pt_theory(C, A1, A2).express(A0))

#______________________________________________________________________________
#______________________________________________________________________________
#
# Estabilidad y control dinámicos longitudinales en cadena abierta
# Análisis de la estabilidad longitudinal de un B747-100
# https://nbviewer.jupyter.org/github/AlexS12/Mecanica_Vuelo/blob/master/MVII_MatrizSistema.ipynb


'''
Created on Jun 7, 2012

@author: vinnie
'''

import os
import sys
import cv2 as cv
from scipy.ndimage import filters
#import scipy.signal as signal
import numpy as np
import pylab as plb
import matplotlib.cm as cm
from itertools import product

_ALPHA = 100
_BETA = 200
_W_LINE = 250
_W_EDGE = 30
_MIN_DISTANCE = 10
_INITIAL_SMOOTH = 15
_INITIAL_ITERATIONS = 30
_ITERATIONS_DELTA = 5
_SMOOTH_FACTOR_DELTA = 4

_NUM_NEIGHBORS = 9
_MAX_SNAXELS = 10000
_INITIAL_DISTANCE_BETWEEN_SNAXELS = 50

def _display(image, snaxels=None):
    """
    Display a grayscale image with pylab, and draw the contour if there is any.
    """
    plb.clf()
    if snaxels is not None:
        for s in snaxels:
            plb.plot(s[0],s[1],'g.',markersize=10.0)

    plb.imshow(image, cmap=cm.Greys_r)
    plb.draw()

    return

def _gradientImage(image):
    """
    Obtain a gradient image (in both x and y directions)
    """
    gradient = np.sqrt(filters.sobel(image, 0)**2 + filters.sobel(image, 1)**2)
    gradient -= gradient.min()

    return gradient

def _inBounds(image, point):
    """
    Is the point within the bounds of the image?
    """
    return np.all(point < np.shape(image)) and np.all(point > 0)

def _externalEnergy(image, smooth_image, point):
    """
    The external energy of the point, a combination of line and edge
    """
    pixel = 255 * image[point[1]][point[0]]
    smooth_pixel = 255 * smooth_image[point[1]][point[0]]
    external_energy = (_W_LINE * pixel) - (_W_EDGE * (smooth_pixel**2))
    return external_energy

def _energy(image, smooth_image, current_point, next_point, previous_point=None):
    """
    Total energy (internal and external).
    Internal energy measures the shape of the contour
    """
    d_squared = np.linalg.norm(next_point -current_point)**2

    if previous_point is None:
        e =  _ALPHA * d_squared + _externalEnergy(image, smooth_image, current_point)
        return e
    else:
        deriv = np.sum((next_point - 2 * current_point + previous_point)**2)
        e = 0.5 * (_ALPHA * d_squared + _BETA * deriv + _externalEnergy(image, smooth_image, current_point))
        return e

def _iterateContour(image, smooth_image, snaxels, energy_matrix, position_matrix, neighbors):
    """
    Compute the minimum energy locations for all the snaxels in the contour
    """
    snaxels_added = len(snaxels)
    for curr_idx in range(snaxels_added - 1, 0, -1):
        energy_matrix[curr_idx][:][:] = float("inf")
        prev_idx = (curr_idx - 1) % snaxels_added
        next_idx = (curr_idx + 1) % snaxels_added

        for j, next_neighbor in enumerate(neighbors):
            next_node = snaxels[next_idx] + next_neighbor

            if not _inBounds(image, next_node):
                continue

            min_energy = float("inf")
            for k, curr_neighbor in enumerate(neighbors):
                curr_node = snaxels[curr_idx] + curr_neighbor
                distance = np.linalg.norm(next_node - curr_node)

                if not _inBounds(image, curr_node) or (distance < _MIN_DISTANCE):
                    continue

                min_energy = float("inf")
                for l, prev_neighbor in enumerate(neighbors):
                    prev_node = snaxels[prev_idx] + prev_neighbor

                    if not _inBounds(image, prev_node):
                        continue

                    energy = energy_matrix[prev_idx][k][l] + _energy(image, smooth_image, curr_node, next_node, prev_node)

                    if energy < min_energy:
                        min_energy = energy
                        min_position_k = k
                        min_position_l = l

                energy_matrix[curr_idx][j][k] = min_energy
                position_matrix[curr_idx][j][k][0] = min_position_k
                position_matrix[curr_idx][j][k][1] = min_position_l

    min_final_energy = float("inf")
    min_final_position_j = 0
    min_final_position_k = 0

    for j in range(_NUM_NEIGHBORS):
        for k in range(_NUM_NEIGHBORS):
            if energy_matrix[snaxels_added - 2][j][k] < min_final_energy:
                min_final_energy = energy_matrix[snaxels_added - 2][j][k]
                min_final_position_j = j
                min_final_position_k = k

    pos_j = min_final_position_j
    pos_k = min_final_position_k

    for i in range(snaxels_added - 1, -1, -1):
        snaxels[i] = snaxels[i] + neighbors[pos_j]
        if i > 0:
            pos_j = position_matrix[i - 1][pos_j][pos_k][0]
            pos_k = position_matrix[i - 1][pos_j][pos_k][1]

    return min_final_energy

def activeContour(image, snaxels):
    """
    Iterate the contour until the energy reaches an equilibrium
    """
    energy_matrix = np.zeros( (_MAX_SNAXELS - 1, _NUM_NEIGHBORS, _NUM_NEIGHBORS), dtype=np.float32)
    position_matrix = np.zeros( (_MAX_SNAXELS - 1, _NUM_NEIGHBORS, _NUM_NEIGHBORS, 2), dtype=np.int32 )
    neighbors = np.array([[i, j] for i in range(-1, 2) for j in range(-1, 2)])
    min_final_energy_prev = float("inf")

    counter = 0
    smooth_factor = _INITIAL_SMOOTH
    iterations = _INITIAL_ITERATIONS
    gradient_image = _gradientImage(image)
    smooth_image = cv.blur(gradient_image, (smooth_factor, smooth_factor))

    while True:
        counter += 1
        if not (counter % iterations):
            iterations += _ITERATIONS_DELTA
            if smooth_factor > _SMOOTH_FACTOR_DELTA:
                smooth_factor -= _SMOOTH_FACTOR_DELTA
            smooth_image = cv.blur(gradient_image, (smooth_factor, smooth_factor))
            print "Deblur step, smooth factor now: ", smooth_factor

        _display(smooth_image, snaxels)
        min_final_energy = _iterateContour(image, smooth_image, snaxels, energy_matrix, position_matrix, neighbors)

        if (min_final_energy == min_final_energy_prev) or smooth_factor < _SMOOTH_FACTOR_DELTA:
            print "Min energy reached at ", min_final_energy
            print "Final smooth factor ", smooth_factor
            break
        else:
            min_final_energy_prev = min_final_energy

def _pointsOnCircle(center, radius, num_points=12):
    points = np.zeros((num_points, 2), dtype=np.int32)
    for i in range(num_points):
        theta = float(i)/num_points * (2 * np.pi)
        x = center[0] + radius * np.cos(theta)
        y = center[1] + radius * np.sin(theta)
        p = [x, y]
        points[i] = p

    return points

def activeContourFromCircle(image_file, center, radius):
    image = plb.imread(image_file)
    if image.ndim > 2:
        image = np.mean(image, axis=2)
    print "Image size: ", image.shape

    plb.ion()
    plb.figure(figsize=np.array(np.shape(image))/50.)

    _display(image)
#    num_points = int((2 * np.pi * radius)/_INITIAL_DISTANCE_BETWEEN_SNAXELS)
    snaxels = _pointsOnCircle(center, radius, 30)
    _display(image, snaxels)
    activeContour(image, snaxels)

    # show and save the result
    plb.ioff()
    _display(image, snaxels)
    plb.savefig(os.path.splitext(image_file)[0] + "-contour-result.png")
    plb.show()
    return

def _test():
    """
    Run the active contour on an image file
    """
    activeContourFromCircle("mri.png", (290, 440), 125)

    return

if __name__ == '__main__':
    _test()


from __future__ import print_function
import cv2
import PIL.Image
import numpy as np
import scipy.stats
import sys
import itertools
from line_intersection import *
np.set_printoptions(suppress=True, precision=2)

def scaleImageIfNeeded(img, max_width=1024, max_height=1024):
  """Scale image down to max_width / max_height keeping aspect ratio if needed. Do nothing otherwise."""
  # Input and Output is a numpy array
  img = PIL.Image.fromarray(img)
  img_width, img_height = img.size
  # print("Image size %dx%d" % (img_width, img_height))
  aspect_ratio = min(float(max_width)/img_width, float(max_height)/img_height)
  if aspect_ratio < 1.0:
    new_width, new_height = ((np.array(img.size) * aspect_ratio)).astype(int)
    # print(" Resizing to %dx%d" % (new_width, new_height))
    return np.array(img.resize((new_width,new_height)))
  return np.array(img)

def getAngle(a,b,c):
  # Get angle given 3 side lengths, in degrees
  k = (a*a+b*b-c*c) / (2*a*b)
  # Handle floating point errors
  if (k < -1):
    k=-1
  elif k > 1:
    k=1
  return np.arccos(k) * 180.0 / np.pi

def angleCloseDeg(a, b, angle_threshold=10):
  d = np.abs(a - b)
  # Handle angles that are ~180 degrees apart
  return d <= angle_threshold or np.abs(d-180) <= angle_threshold

def getSegmentThetaRho(line):
  x1,y1,x2,y2 = line
  theta = np.math.atan2(y2-y1, x2-x1)
  m = np.tan(theta)
  # rho = np.abs(y1 + m*x1) / np.sqrt(m*m+1)
  rho = x1*np.cos(theta) + y1*np.sin(theta)
  return theta, rho

def getTwoLineSegmentIntersection(p,pr,q,qs):
  # Uses http://stackoverflow.com/a/565282/2574639
  # Given two line segments defined by sets of points
  # (p -> pr) and (q -> qs).
  # Return the intersection point between them
  # *assumes it always exists for our particular use-case*

  # Convert to floats
  p = p.astype(np.float32)
  pr = pr.astype(np.float32)
  q = q.astype(np.float32)
  qs = qs.astype(np.float32)
  r = pr-p
  s = qs-q
  # print(p, pr, r)
  # print(q, qs, s)
  rxs = np.cross(r, s)
  if rxs == 0:
    return [] # parallel
  t = np.cross((q - p), s) / rxs
  return p + t*r # intersect

def chooseRandomGoodQuad(lines_a, lines_b, median_contour):
  # Get random set of points
  # Redo until min side distance of random corners greater than a multiple
  # of the median tile found from initial estimator.
  sides_tile = getSquareSides(median_contour)
  for i in range(50):
    corners = chooseRandomQuad(lines_a, lines_b)
    sides_quad = getSquareSides(corners)
    if (i < 5):
      tile_size_mult = 5
    elif (i < 10):
      tile_size_mult = 4
    elif (i < 20):
      tile_size_mult = 3
    elif (i < 30):
      tile_size_mult = 2
    else:
      tile_size_mult = 1

    if min(sides_quad) > min(sides_tile*tile_size_mult):
      return corners

  print('chooseRandomGoodQuad hit max iter: %d' % i)
  return corners

def chooseRandomQuad(lines_a, lines_b):
  # Return 1 random quad (4 points) by choosing
  # 2 lines from lines_a and 2 lines from lines_b
  # and returning their intersections
  a = np.random.choice(range(len(lines_a)),2, replace=False)
  b = np.random.choice(range(len(lines_b)),2, replace=False)

  pairs = np.array([
    [a[0], b[0]],
    [a[0], b[1]],
    [a[1], b[1]],
    [a[1], b[0]],
    ])

  corners = np.zeros([4,2])
  for i in range(4):
    k1 = lines_a[pairs[i,0]]
    k2 = lines_b[pairs[i,1]]
    corners[i,:] = getTwoLineSegmentIntersection(k1[:2], k1[2:], k2[:2], k2[2:])
  return corners

def getSegmentTheta(line):
  x1,y1,x2,y2 = line
  theta = np.math.atan2(y2-y1, x2-x1)
  return theta

def is_square(cnt, eps=3.0, xratio_thresh = 0.5):
  # 4x2 array, rows are each point, columns are x and y
  center = cnt.sum(axis=0)/4

  # Side lengths of rectangular contour
  dd0 = np.sqrt(((cnt[0,:] - cnt[1,:])**2).sum())
  dd1 = np.sqrt(((cnt[1,:] - cnt[2,:])**2).sum())
  dd2 = np.sqrt(((cnt[2,:] - cnt[3,:])**2).sum())
  dd3 = np.sqrt(((cnt[3,:] - cnt[0,:])**2).sum())

  # diagonal ratio
  xa = np.sqrt(((cnt[0,:] - cnt[2,:])**2).sum())
  xb = np.sqrt(((cnt[1,:] - cnt[3,:])**2).sum())
  xratio = xa/xb if xa < xb else xb/xa

  # Check whether all points part of convex hull
  # ie. not this http://i.stack.imgur.com/I6yJY.png
  # all corner angles, angles are less than 180 deg, so not necessarily internal angles
  ta = getAngle(dd3, dd0, xb)
  tb = getAngle(dd0, dd1, xa)
  tc = getAngle(dd1, dd2, xb)
  td = getAngle(dd2, dd3, xa)
  angle_sum = np.round(ta+tb+tc+td)

  is_convex = angle_sum == 360

  angles = np.array([ta,tb,tc,td])
  good_angles = np.all((angles > 40) & (angles < 140))

  # side ratios
  dda = dd0 / dd1
  if dda < 1:
    dda = 1. / dda
  ddb = dd1 / dd2
  if ddb < 1:
    ddb = 1. / ddb
  ddc = dd2 / dd3
  if ddc < 1:
    ddc = 1. / ddc
  ddd = dd3 / dd0
  if ddd < 1:
    ddd = 1. / ddd
  side_ratios = np.array([dda,ddb,ddc,ddd])
  good_side_ratios = np.all(side_ratios < eps)

  # Return whether side ratios within certain ratio < epsilon
  return (
    # abs(1.0 - dda) < eps and
    # abs(1.0 - ddb) < eps and
    # xratio > xratio_thresh and
    # good_side_ratios and
    # is_convex and
    good_angles)

def minimum_distance2(v, w, p):
  # Return squared min distance between point p and line segment vw
  # Via http://stackoverflow.com/a/1501725
  # Return minimum distance between line segment vw and point p
  l2 = np.sum((v - w)**2)  # i.e. |w-v|^2 -  avoid a sqrt
  if (l2 == 0.0):
    return np.sum((p - v)**2)   # v == w case
  # Consider the line extending the segment, parameterized as v + t (w - v).
  # We find projection of point p onto the line.
  # It falls where t = [(p-v) . (w-v)] / |w-v|^2
  # We clamp t from [0,1] to handle points outside the segment vw.
  t = max(0, min(1, np.dot(p - v, w - v) / l2))
  projection = v + t * (w - v)  # Projection falls on the segment
  return np.sum((p - projection)**2)

def getMinLineAngleDistance(a0, a1):
  # Compare line angles (which can be 180 off from one another, or +- 180)
  v0 = abs(a1-a0)
  v1 = abs((a1+np.pi) - a0)
  v2 = abs(a1 - (a0+np.pi))
  return min([v0,v1,v2])

def getBestCorners(tile_corners, hough_lines, angle_threshold = 10*np.pi/180):
  # Given 4x2 imperfect tile corners and Nx4 line segments
  # Expects line segments and corner points to be in same cartesian space
  #
  # Find 4 best line segments that are best match to the tile corners
  # and return the corners based off of those line segments, and those line segments
  best_lines = np.zeros([4,4])
  for i in range(4):
    corner_theta = getSegmentTheta(tile_corners[[i,i,((i+1)%4),((i+1)%4)], [0,1,0,1]])
    corner_ctr_pt = (tile_corners[i,:] + tile_corners[((i+1)%4),:]) / 2

    best_d = 1e6
    for line in hough_lines:
      theta = getSegmentTheta(line)
      # If angle within 10 degrees
      # if abs(corner_theta - theta) < angle_threshold:
      if getMinLineAngleDistance(corner_theta, theta) < angle_threshold:
        d = minimum_distance2(line[:2], line[2:], corner_ctr_pt)
        if d < best_d:
          best_d = d
          best_lines[i,:] = line

  new_corners = np.zeros([4,2], dtype=np.float32)
  for i in range(4):
    x = getTwoLineSegmentIntersection(
      best_lines[i,:2], best_lines[i,2:],
      best_lines[(i+1)%4,:2], best_lines[(i+1)%4,2:])
    # print(best_lines, x)
    # print(best_lines[i,:2], best_lines[i,2:], best_lines[(i+1)%4,:2], best_lines[(i+1)%4,2:])
    new_corners[i,:] = x

  return new_corners, best_lines

def simplifyContours(contours):
  for i in range(len(contours)):
    # Approximate contour and update in place
    contours[i] = cv2.approxPolyDP(contours[i],0.04*cv2.arcLength(contours[i],True),True)

def pruneContours(contours):
  new_contours = []
  for i in range(len(contours)):
    cnt = contours[i]
    # Only contours that fill an area of at least 8x8 pixels
    if cv2.contourArea(cnt) < 8*8:
      continue
    # Only rectangular contours allowed
    if len(cnt) != 4:
      continue

    if not is_square(cnt):
      continue

    new_contours.append(cnt)
  new_contours = np.array(new_contours)
  if len(new_contours) == 0:
    return new_contours, None

  norm_contours = new_contours[:,:,0,:] - new_contours[:,[0],0,:]
  median_contour = np.median(norm_contours, axis=0).astype(int)
  diff = np.sqrt(np.sum((norm_contours - median_contour)**2,axis=2))

  new_contours = new_contours[np.all(diff < 60, axis=1)]

  return np.array(new_contours), median_contour

def getSquareSides(cnt):
  # 4x2 array, rows are each point, columns are x and y
  center = cnt.sum(axis=0)/4

  # Side lengths of rectangular contour
  dd0 = np.sqrt(((cnt[0,:] - cnt[1,:])**2).sum())
  dd1 = np.sqrt(((cnt[1,:] - cnt[2,:])**2).sum())
  dd2 = np.sqrt(((cnt[2,:] - cnt[3,:])**2).sum())
  dd3 = np.sqrt(((cnt[3,:] - cnt[0,:])**2).sum())
  return np.array([dd0, dd1, dd2, dd3])

from time import time
def calculateMask(mask_shape, contours, iters=10):
  a = time()
  sum_mask = np.zeros(mask_shape, dtype=int)
  tmp_mask = np.zeros(mask_shape, dtype=int)
  for i, cnt in enumerate(contours):
    for i in np.linspace(5,23,iters):
    # for i in [3,5,7,9,11,13,15]:
      # Calculate oversized tile mask and add to sum
      # big_cnt = (cnt.mean(axis=0) + (cnt-cnt.mean(axis=0))*i).astype(int)
      cnt_center = cnt.mean(axis=0)
      big_cnt = (cnt*i + cnt_center*(1-i)).astype(int)
      tmp_mask[:] = 0 # Reset
      cv2.drawContours(tmp_mask,[big_cnt],0,1,-1) # Fill mask with 1's inside contour
      sum_mask += tmp_mask

  # Normalize mask to 0-1 range
  sum_mask = sum_mask.astype(float) / sum_mask.max()
  # print("Mask calc took %.4f seconds." % (time() - a))
  return sum_mask

def getContourThetas(contours):
  thetas = []
  for cnt in contours:
    cnt = cnt[:,0,:]
    if cnt[0,0] < cnt[1,0]:
      side0 = np.hstack([cnt[1,:],cnt[0,:]])
    else:
      side0 = np.hstack([cnt[0,:],cnt[1,:]])
    if cnt[1,0] < cnt[2,0]:
      side1 = np.hstack([cnt[2,:],cnt[1,:]])
    else:
      side1 = np.hstack([cnt[1,:],cnt[2,:]])
    if cnt[2,0] < cnt[3,0]:
      side2 = np.hstack([cnt[3,:],cnt[2,:]])
    else:
      side2 = np.hstack([cnt[2,:],cnt[3,:]])
    if cnt[3,0] < cnt[0,0]:
      side3 = np.hstack([cnt[0,:],cnt[3,:]])
    else:
      side3 = np.hstack([cnt[3,:],cnt[0,:]])
    theta0 = getSegmentTheta(side0)
    theta1 = getSegmentTheta(side1)
    theta2 = getSegmentTheta(side2)
    theta3 = getSegmentTheta(side3)
    thetas.extend([theta0,theta1,theta2,theta3])
  return np.array(thetas)

def getEstimatedChessboardMask(img, edges, iters=10):
  # Morphological Gradient to get internal squares of canny edges.
  kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(3,3))
  edges_gradient = cv2.morphologyEx(edges, cv2.MORPH_GRADIENT, kernel)

  _, contours, hierarchy = cv2.findContours(edges_gradient, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
  # Approximate polygons of contours
  simplifyContours(contours)

  if len(contours) == 0:
    return np.ones(img.shape[:2], dtype=float), None, None, None

  # Prune contours to rectangular ones
  contours, median_contour = pruneContours(contours)

  if len(contours) == 0 or median_contour is None:
    return np.ones(img.shape[:2], dtype=float), None, None, None

  thetas = getContourThetas(contours)

  top_two_angles = calculateKDE(thetas)

  mask = calculateMask(edges_gradient.shape, contours, iters)

  min_area_rect = getMinAreaRect(mask)

  return mask, top_two_angles, min_area_rect, median_contour

def calculateKDE(thetas):
  thetas *= 180/np.pi
  thetas[thetas<0] += 180

  kde_func = scipy.stats.gaussian_kde(thetas)
  positions = np.linspace(-40,180+40,360)
  kde_res = kde_func(positions)

  left_half = np.diff(kde_res)
  right_half = np.diff(kde_res[::-1])

  f = kde_res.copy()
  f[1:][left_half<0] = 0
  f[:-1][right_half[::-1]<0] = 0

  peak_indices = np.argwhere(f).flatten()
  peak_angles = positions[peak_indices]

  order = np.argsort(kde_res[peak_indices])[::-1] # strongest to weakest peaks

  return peak_angles[order[:2]] # top two strongest angles in degrees

# def plotKDE(thetas):
#   thetas *= 180/np.pi
#   thetas[thetas<0] += 180

#   ax1 = plt.subplot(211)
#   plt.plot(thetas,np.zeros(thetas.shape),'.')
#   plt.hist(thetas,20)

#   plt.subplot(212, sharex=ax1)
#   kde_func = scipy.stats.gaussian_kde(thetas)
#   positions = np.linspace(-40,180+40,360)
#   kde_res = kde_func(positions)
#   plt.plot(positions, kde_res)

#   c = kde_res.copy()
#   left_half = np.diff(kde_res)
#   right_half = np.diff(kde_res[::-1])

#   f = c.copy()
#   f[1:][left_half<0] = 0
#   f[:-1][right_half[::-1]<0] = 0
#   peak_indices = np.argwhere(f).flatten()
#   print(peak_indices, positions[peak_indices])
#   peak_angles = positions[peak_indices]

#   plt.plot(peak_angles, kde_res[peak_indices],'go')
#   order = np.argsort(kde_res[peak_indices][::-1]) # strongest to weakest peaks
#   return peak_angles[order[:2]], peak_angles, kde_res[peak_indices] # two angles in degrees

def getHoughLines(edges, min_line_size=100):
  # Expects chessboard to take up over 50% of edge map
  # min_line_size = int(min(edges.shape)/8)
  lines = cv2.HoughLinesP(edges,1,np.pi/360.0, int(min_line_size),
    minLineLength = min_line_size, maxLineGap = min_line_size/2)

  if lines is None:
    return np.array([])

  return lines[:,0,:]

def getSegmentTheta(line):
  x1,y1,x2,y2 = line
  theta = np.math.atan2(y2-y1, x2-x1)
  return theta

def parseHoughLines(lines, top_two_angles, angle_threshold_deg=20):
  is_good = np.zeros(len(lines)) # 0 = bad, 1 = close to 1st angle, 2 = close to 2nd angle
  for i, line in enumerate(lines):
    theta = getSegmentTheta(line) * 180/np.pi # to degrees
    d1 = getMinLineAngleDistanceDeg(theta, top_two_angles[0])
    d2 = getMinLineAngleDistanceDeg(theta, top_two_angles[1])
    if (d1 < angle_threshold_deg):
      is_good[i] = 1
    elif (d2 < angle_threshold_deg):
      is_good[i] = 2
  lines_a = lines[is_good==1]
  lines_b = lines[is_good==2]
  return lines_a, lines_b

def getMinLineAngleDistance(a0, a1):
  # In radians
  # Compare line angles (which can be 180 off from one another, or +- 180)
  v0 = abs(a1-a0)
  v1 = abs((a1+np.pi) - a0)
  v2 = abs(a1 - (a0+np.pi))
  return min([v0,v1,v2])

def getMinLineAngleDistanceDeg(a0, a1):
  # In degrees
  # Compare line angles (which can be 180 off from one another, or +- 180)
  v0 = abs(a1-a0)
  v1 = abs((a1+180) - a0)
  v2 = abs(a1 - (a0+180))
  return min([v0,v1,v2])

def plotHoughLines(img, lines, color=(255,255,255), line_thickness=2):
  # colors = np.random.random([lines.shape[0],3])*255
  # colors = np.array([
  #     [20,20,20],
  #     [255,0,0],
  #     [0,255,0],
  #     [255,255,0],
  #     [0,0,255],
  #     [255,0,255],
  #     [0,255,255],
  #     [200,200,200],
  #     ], dtype=np.uint8)
  # Plot lines
  for i, line in enumerate(lines):
    # color = list(map(int,colors[i%len(colors)]))
    cv2.line(img,
      tuple(line[:2].astype(np.int)),
      tuple(line[2:].astype(np.int)), color, thickness=line_thickness)

def getMinAreaRect(mask):
  a = np.argwhere(mask.T>0.5)
  # rect = cv2.boundingRect(a)
  rect = cv2.minAreaRect(a)
  return rect

def drawMinAreaRect(img, rect, color=(0,255,255)):
  ctr = tuple(map(int,rect[0]))

  box = cv2.boxPoints(rect)
  box = np.int0(box)

  cv2.drawContours(img,[box],0,color,2)
  cv2.circle(img, ctr, 3, (255,0,0),-1)

def skeletonize_1d(data):
  c = data.copy()
  left_half = np.diff(data)
  right_half = np.diff(data[::-1])

  f = data.copy()
  f[1:][left_half<0] = 0
  f[:-1][right_half[::-1]<0] = 0
  return f

def getWarpCheckerLines(img):
  """Given a warped axis-aligned image of a chessboard, return internal line crossings"""
  # TODO: Fix awkward conversion
  # Convert RGB numpy array to image, then to grayscale image, then back to numpy array
  img_gray = np.array(PIL.Image.fromarray(img).convert('L'))
  img_gray = cv2.bilateralFilter(img_gray,15,75,75)

  # Find gradients
  sobelx = cv2.Sobel(img_gray,cv2.CV_64F,1,0,ksize=5)
  sobely = cv2.Sobel(img_gray,cv2.CV_64F,0,1,ksize=5)

  sobelx_pos = sobelx.copy()
  sobelx_pos[sobelx <= 0] = 0
  sobelx_neg = sobelx.copy()
  sobelx_neg[sobelx > 0] = 0

  sobely_pos = sobely.copy()
  sobely_pos[sobely <= 0] = 0
  sobely_neg = sobely.copy()
  sobely_neg[sobely > 0] = 0

  kernel = np.array([ 0.  ,  0.  ,  0.04,  0.32,  0.88,  0.88,  0.32,  0.04,  0.  ,  0.  ])

  checker_x = np.sum(sobelx_pos, axis=0) * np.sum(-sobelx_neg, axis=0)
  checker_x = np.convolve(checker_x, kernel, 'same')
  checker_x = checker_x / checker_x.max()
  checker_x[checker_x<0.1] = 0
  checker_x = skeletonize_1d(checker_x)

  checker_y = np.sum(sobely_pos, axis=1) * np.sum(-sobely_neg, axis=1)
  checker_y = np.convolve(checker_y, kernel, 'same')
  checker_y = checker_y / checker_y.max()
  checker_y[checker_y<0.1] = 0
  checker_y = skeletonize_1d(checker_y)

  x_lines = np.argwhere(checker_x).flatten()
  y_lines = np.argwhere(checker_y).flatten()

  #######
  ## Semi-brute force approach, merge all combinations of 3 points
  # with equal spacing under one standard deviation
  x_lines = getBestEqualSpacing(x_lines)
  y_lines = getBestEqualSpacing(y_lines)

  ###########

  if len(x_lines) < 7 or len(y_lines) < 7:
    return [], [], [], []

  # Select set of 7 adjacent lines with max sum score
  x_scores = np.zeros(x_lines.shape[0]-7+1)
  for i in range(0,x_lines.shape[0]-7+1):
    x_scores[i] = np.sum(checker_x[x_lines[i:i+7]])
  x_start = np.argmax(x_scores)
  strongest_x_lines = range(x_start,x_start+7)

  y_scores = np.zeros(y_lines.shape[0]-7+1)
  for i in range(0,y_lines.shape[0]-7+1):
    y_scores[i] = np.sum(checker_y[y_lines[i:i+7]])
  y_start = np.argmax(y_scores)
  strongest_y_lines = range(y_start,y_start+7)

  # TODO: Sanity check areas between lines for consistent color when choosing?

  # Choose best internal 7 chessboard lines
  lines_x = x_lines[strongest_x_lines]
  lines_y = y_lines[strongest_y_lines]

  # Add outer chessboard edges assuming consistent step size
  step_x = np.median(np.diff(lines_x))
  step_y = np.median(np.diff(lines_y))

  lines_x = np.hstack([lines_x[0]-step_x, lines_x, lines_x[-1]+step_x])
  lines_y = np.hstack([lines_y[0]-step_y, lines_y, lines_y[-1]+step_y])

  return lines_x, lines_y, step_x, step_y

  # x_lines = np.argwhere(checker_x).flatten()
  # y_lines = np.argwhere(checker_y).flatten()

  # x_diff = np.diff(x_lines)
  # y_diff = np.diff(y_lines)

  # step_x_pred = np.median(x_diff)
  # step_y_pred = np.median(y_diff)

def pruneGradLines(a, b, eta=10):
  # Remove values from vector 'a' that aren't close to values in vector b
  is_good = np.zeros(len(a),dtype=bool)
  for i,v in enumerate(a):
    if min(b-v) < eta:
      is_good[i] = True
  return a[is_good]

def main(filenames):
  for filename in filenames:
    img = cv2.imread(filename)
    img = scaleImageIfNeeded(img, 600, 480)

    # Edges
    edges = cv2.Canny(img, 100, 550)
    mask, _, _, _ = getEstimatedChessboardMask(img, edges, iters=10)

    img_masked_full = cv2.bitwise_and(img,img,mask = (mask > 0.5).astype(np.uint8))
    img_masked = cv2.addWeighted(img,0.2,img_masked_full,0.8,0)
    edges_masked = cv2.bitwise_and(edges,edges,mask = (mask > 0.5).astype(np.uint8))

    cv2.imshow('img %s' % filename,img_masked)
    cv2.imshow('edges %s' % filename, edges_masked)
    cv2.imshow('mask %s' % filename, mask)

  cv2.waitKey(0)
  cv2.destroyAllWindows()

if __name__ == '__main__':
  if len(sys.argv) > 1:
    filenames = sys.argv[1:]
  else:
    # filenames = ['input/1.jpg']
    filenames = ['input2/18.jpg']
  print("Loading", filenames)
  main(filenames)


<reponame>bartolsthoorn/PH_XXZ
import h5py
import numpy as np
from ripser import ripser
import gudhi as gd
from persim import plot_diagrams
from scipy.spatial.distance import pdist
from scipy.spatial.distance import squareform
from tqdm import tqdm
from multiprocessing import Pool
import gudhi as gd
import pickle
import persim

filename = 'snapshots_4_100000_32.hdf5'
#filename = 'snapshots_6_100000_32.hdf5'
f = h5py.File(filename, 'r')
L = 4
N = L*L*L*16

spins_type = 'global'

coordinates = f['coordinates'][:].transpose()
coordinates = coordinates.reshape(N,3)
coordinates_D = squareform(pdist(coordinates, metric='euclidean'))
print(np.min(coordinates_D[coordinates_D > 0]))
spins = f['spins_' + spins_type][:].transpose()
print(spins.shape)

J_grid = f['J'][:]
T_grid = f['T'][:]
print(J_grid)

N_T = len(T_grid)
N_J = len(J_grid)

f.close()

params = []
for i, J in enumerate(J_grid):
    for j, T in enumerate(T_grid):
        params.append((J,T))

n_samples = 32
def calculate_barcodes(samples):
    points_list = [coordinates + (spins * 0.3535533905932738 * 0.25) for spins in samples]
    points_list = points_list[0:n_samples]
    dgms = [[] for i in range(3)]
    for points in points_list:
        #rips = ripser(points, thresh=max_death, maxdim=maxdim)
        alpha_complex = gd.AlphaComplex(points=points)
        st_alpha = alpha_complex.create_simplex_tree()
        barcodes = st_alpha.persistence()
        #for i, d in enumerate(rips['dgms']):
        #    dgms[i] += list(d)
        for d in np.arange(3):
            dgms[d] += [bar for dim, bar in barcodes if dim == d]
    return dgms

barcodes = []
for i, J in tqdm(enumerate(J_grid), total=len(J_grid)):
    for j, T in tqdm(enumerate(T_grid), total=len(T_grid)):
        barcodes.append(calculate_barcodes(spins[i, j, :]))

pickle.dump(barcodes, open(filename.split('.')[0] + '_' + spins_type + ('_barcodes_%d.p' % n_samples), 'wb' ))


<reponame>HuguesMoreau/TMD_fusion_benchmark<filename>preprocess/transforms.py
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
This file contains diverse preprocessing functions (mostly norms ans spectrograms),
and basic tests and visualizations.
If you are to work with any IPython console (ex: with Jupyter or spyder), is is advised
to launch a '%matplotlib qt' ,to get clean widow
"""

if __name__ == '__main__': # this is used to launch the file from anywhere
    import sys
    sys.path.append("..")

import numpy as np
import torch
import scipy.signal, scipy.interpolate, scipy.ndimage

from param import classes_names, fs, duration_window, duration_overlap, duration_segment, spectro_batch_size
from preprocess import Datasets

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    n_classes = len(classes_names)
    # We will need this for the tests
    DS = Datasets.SignalsDataSet(mode='train', split='balanced', comp_preprocess_first=False)

#%% transform functions

"""In all following functions, the input parameter (data) is, by default,
 a dict of numpy arrays, containing signal names (eg. "Gyr_z") as keys, and 1-dimensional
 arrays as values

Most of this part contains basic visualizations to make sure the preprocessing is correct"""

class TemporalTransform():
    """  create the base transform to use to each element of the data
    Also generates data for thr 'Random' and 'Zero' cases

    Parameters
    ----------
    signal_name_list: a list of string signals (ex: 'Gyr_y', 'Ori_x')
        If a string ends by "_norm" (ex: "Mag_norm"), the output will
        be the norm of the three (or four) axis of the signal.
        The signals can also be 'Zero', in which case the segments are only
        zeros, or 'Random', in which case each data oint is sampled with normal
        distribution (zero mean, unit variance)

    Returns
    -------
    a function with input:  a dict of (_, 6000) arrays (key example: 'Gyr_y')
                and output: a dictionnary of arrays.
    """
    def __init__(self, signal_name_list):
        super(TemporalTransform, self).__init__()
        self.signal_name_list = signal_name_list

    def __call__(self, data):
        """
        Parameters
        ----------
        a dict of (B, 6000) arrays (key example: 'Gyr_y')

        Returns
        -------
        a dictionnary of arrays. This time, the keys are from signal_name_list,
            and the values are either raw signals (if the key ends with '_x',
            '_y', '_z', or '_w'); a norm of several signals (if the key ends
            with '_norm'); or a specific signal (if the key is 'Random' or
            'Zero'). The shape of each array is (B, 6000), where B (batch size)
            depends on the input shape.
        """

        outputs = {}
        for signal_name in self.signal_name_list:

            if signal_name[-2:] in ['_x', '_y', '_z', '_w'] or signal_name == "Pressure":
                processed_signal = data[signal_name]

            elif signal_name == 'Random':
                data_shape = data["Acc_x"].shape
                processed_signal = np.random.randn(data_shape[0], data_shape[1]).astype(np.float32)

            elif signal_name == 'Zero':
                data_shape = data["Acc_x"].shape
                processed_signal = np.zeros(data_shape).astype(np.float32)

            elif signal_name[-5:] == '_norm':
                suffix_location = signal_name.index("_") # 4 if signal_name == "LAcc", 3 otherwise
                sensor = signal_name[:suffix_location]    # ex: 'Acc', 'LAcc'

                if sensor == "Ori":
                    # in that case, data[sensor+"_x"]**2 + data[sensor+"_y"]**2 + data[sensor+"_z"]**2 should be 1.0
                    processed_signal = np.sqrt(data[sensor+"_x"]**2 + data[sensor+"_y"]**2 + data[sensor+"_z"]**2 \
                                             + data[sensor+"_w"]**2)
                else :
                    processed_signal = np.sqrt(data[sensor+"_x"]**2 + data[sensor+"_y"]**2 + data[sensor+"_z"]**2)
            else :
                raise ValueError("unknown signal name: '{}'. Signal names should end with either '_x', '_y', '_z', '_w', or '_norm'".format(signal_name))

            outputs[signal_name] = processed_signal

        return outputs

    def __str__(self):
        """purely for visual purposes, so that we can print() the function"""
        str_to_return = "Temporal_transform"
        str_to_return += "\n\t Signals: {}".format(self.signal_name_list)
        return str_to_return

if __name__ == "__main__":

    # plot one figure per sensor
    # on each figure, one subplot per class,
    # to find one instance per each class, we start looking at index = index0
    index0 = 0

    for tested_signal_name in ["Acc_norm", "Ori_norm", "Mag_norm", "LAcc_x"]:
        # plot 1 segment from each class.
        plt.figure()

        if tested_signal_name != 'Pressure':
            suffix_location = tested_signal_name.index("_")
            tested_sensor = tested_signal_name[:suffix_location]    # ex: 'Acc', 'LAcc'
        else:
            tested_sensor = 'Pressure'

        sensor_axis = [tested_sensor + axis for axis in ["_x", "_y", "_z"]] if tested_sensor != 'Pressure' else ['Pressure']
        if tested_sensor == "Ori" :  sensor_axis.append(tested_sensor+"_w")

        temporal_transform = TemporalTransform([tested_signal_name])

        remaining_classes = classes_names.copy()
        index = index0

        while len(remaining_classes)>0:
            data_tensor, class_tensor = DS[index] # data is a dict of 2D tensors (1,nb)
            data_cpu = {signal:data_tensor[signal].to(torch.device('cpu')).detach().numpy() for signal in data_tensor.keys()}
            class_index = int(class_tensor)
            class_name = classes_names[class_index-1]

            if class_name in remaining_classes:

                remaining_classes.remove(class_name)
                plt.subplot(2, 4, n_classes - len(remaining_classes))

                for k,signal in enumerate(sensor_axis):

                    if k==0:  # compute the temporal axis once
                        nb = data_cpu[signal].shape[1]
                        x_t = np.linspace(0, nb/fs, nb)

                    plt.plot(x_t, data_cpu[signal][0,:])

                selected_signal = temporal_transform(data_cpu)[tested_signal_name]

                error_message_dtype = "One of the signals does not have the correct type: {}, {} \n dtype should be float32, is actually {}".format(tested_signal_name, str(temporal_transform), selected_signal.dtype)
                assert (selected_signal.dtype == 'float32'), error_message_dtype

                plt.plot(x_t, selected_signal[0,:], '--')
                plt.xlabel("t (s)")
                legend = sensor_axis + [tested_signal_name+' (selected)']
                plt.legend(legend)
                plt.title("{} ({}, index={})".format(tested_sensor, classes_names[class_index-1], index))

            index +=1

        plt.show()

#%% FFT

class FFTTransform():
    """  create a transform to use to return the power of the spectrum
    (computed through a Fourier transform) of each element of the data

    Parameters
    ----------
    signal_name_list: a list of string signals (ex: 'Gyr_y', 'Ori_x')
        If a string ends by "_norm" (ex: "Mag_norm"), the output will
        be the norm of the three (or four) axis of the signal.

    Returns
    -------
    a function with input:  a dict of (B, 6000) arrays (key example: 'Gyr_y')
                and output: a dictionnary of (B, 6000) arrays.
    """
    def __init__(self, signal_name_list):
        super(FFTTransform, self).__init__()
        self.signal_name_list = signal_name_list
        self.temporal_transform = TemporalTransform(signal_name_list)

    def __call__(self, data):
        """
        Parameters
        ----------
        a dict of (B, 6000) arrays (key example: 'Mag_x')

        Returns
        -------
        a dictionnary of arrays. The keys are from signal_name_list, and the
            values are the power spectra of each signal. The shape of each array is
            (B, 6000), where B (batch size) depends on the input shape
        """
        temporal_signals = self.temporal_transform(data)
        del data  # free some memory

        outputs = {}
        for signal_name in self.signal_name_list:
            complex_fft = np.fft.fft(temporal_signals[signal_name], axis=1)
            power_fft = np.abs(complex_fft)
            power_fft[:,0] = 0. # remove the DC component (to avoid this component
              #  outscales te others)

            centered_power_fft = np.fft.fftshift(power_fft, axes=1) # so 0 Hz is in the middle

            outputs[signal_name] = centered_power_fft.astype('float32')
            del temporal_signals[signal_name] # release the memory

            # a faire, calculer les f et les sauver
            # self.f=f

        return outputs

    def __str__(self):
        """purely for visual purposes, so that we can print() the function"""
        str_to_return = "FFT_transform"
        str_to_return += "\n\t Signals: {}".format(self.signal_name_list)
        return str_to_return

#%%

if __name__ == "__main__":

    # classes to plot,
    sel_classes = ["Still","Walk","Run","Train"]
    nsel = len(sel_classes)

    for tested_signal_name in ["Acc_norm", "Gyr_y", "Mag_norm", "Pressure"]:
        # plot 1 segment from each class.
        plt.figure()
        tested_sensor = tested_signal_name[:3]
        if "_" in tested_sensor:
            sensor_axis = [tested_sensor + axis for axis in ["_x", "_y", "_z"]]
        else: # Pressure
            sensor_axis = [tested_sensor]
        if tested_sensor == "Ori" :  sensor_axis.append(tested_sensor+"_w")

        fft_transform =      FFTTransform([tested_signal_name])
        temporal_transform = TemporalTransform([tested_signal_name])

        remaining_classes = sel_classes.copy()

        index = 0
        isub = 1

        while len(remaining_classes)>0:
            data_tensor, class_tensor = DS[index]
            data_cpu = {signal:data_tensor[signal].to(torch.device('cpu')).detach().numpy() for signal in data_tensor.keys()}
            class_index = int(class_tensor)

            class_name = classes_names[class_index-1]

            if class_name in remaining_classes:
                remaining_classes.remove(class_name)

                plt.subplot(2, nsel, isub)
                selected_signal = temporal_transform(data_cpu)[tested_signal_name]

                 # plot the temporal signal
                nb = selected_signal.shape[1]
                x_t = np.linspace(0, nb/fs, nb)

                x_f = np.linspace(-fs/2,fs/2, nb)

                plt.plot(x_t, selected_signal[0,:])
                plt.xlabel("t (s)")
                plt.title("{} ({}, index={})".format(tested_signal_name, classes_names[class_index-1], index))

                 # plot the fft
                plt.subplot(2, nsel, isub + 4)
                isub += 1
                selected_power = fft_transform(data_cpu)[tested_signal_name]

                error_message_dtype = "One of the signals does not have the correct type: {}, {} \n dtype should be float32, is actually {}".format(tested_signal_name, str(fft_transform), selected_power.dtype)
                assert (selected_power.dtype == 'float32'), error_message_dtype

                plt.plot(x_f, selected_power[0,:])
                plt.xlabel("f (Hz)")
                plt.title("FFT of {} ({}, index={})".format(tested_signal_name, classes_names[class_index-1], index))

            index +=1
        plt.show()

#%%

#  ----------------  Spectral transforms  ---------------------

# Interpolation functions
def interpol_log(f, t, spectrogram, out_size):
    """interpolates the spectrogram in input using a linear axis for the timestamps and a LOG axis for the frequencies

    Parameters
    ----------
    f : numpy array, shape: (F_in,), frequencies of the spectrogram
    t : numpy array, shape: (T_in,), timestamps of the spectrogram
    spectrogram : (B, F_in, T_in), B is batch size; 3D numpy array

    out_size : couple of ints (F_out, T_out)

    Returns
    -------
    f_interpolated : numpy array, shape: (F_out,), frequencies of the spectrogram AFTER interpolation
    t_interpolated : numpy array, shape: (T_out,), timestamps of the spectrogram AFTER interpolation
    a spectrogram, where the f axis (second dimension) has been re-interpolated
    using a log axis

    """
    B = spectrogram.shape[0]
    out_f, out_t = out_size

    log_f = np.log(f+f[1]) #  log between 0.2 Hz and 50.2 Hz

    log_f_normalized    = (log_f-log_f[0])/(log_f[-1]-log_f[0]) # between 0.0 and 1.0
    t_normalized        = (t-t[0])/(t[-1]-t[0])

    rescaled_f = out_f*log_f_normalized # 0 and 48
    # rescaled_f = (out_f-1)*log_f_normalized ??
    rescaled_t = out_t*t_normalized

    spectrogram_interpolated = np.zeros( (B, out_f, out_t), dtype='float32')
    index_f, index_t = np.arange(out_f), np.arange(out_t) # between 0 and 47

    for i in range(B):
        spectrogram_fn = scipy.interpolate.interp2d(rescaled_t, rescaled_f, spectrogram[i,:,:], copy=False)
        # interp2d returns a 2D function
        spectrogram_interpolated[i,:,:] = spectrogram_fn(index_t, index_f)  # care to the order

    f_fn = scipy.interpolate.interp1d(rescaled_f, f, copy=False)
    f_interpolated = f_fn(index_f)

    t_fn = scipy.interpolate.interp1d(rescaled_t, t, copy=False)
    t_interpolated = t_fn(index_t)

    return f_interpolated, t_interpolated, spectrogram_interpolated

def interpol_lin(f, t, spectrogram, out_size):
    """interpolates the spectrogram in input using a linear axis for the timestamps AND the frequencies

    Parameters
    ----------
    f : numpy array, shape: (F_in,), frequencies of the spectrogram
    t : numpy array, shape: (T_in,), timestamps of the spectrogram
    spectrogram : (B, F_in, T_in) numpy array
    out_size : couple of ints (F_out, T_out)
     (does not need f or t)

    Returns
    -------
    f_interpolated : numpy array, shape: (F_out,), frequencies of the spectrogram AFTER interpolation
    t_interpolated : numpy array, shape: (T_out,), timestamps of the spectrogram AFTER interpolation
    a spectrogram: 3D numpy array, where the f axis (second dimension) has been re-interpolated
          using a linear axis
    """
    B, F_in, T_in = spectrogram.shape
    out_f, out_t = out_size
    output_shape = (B, out_f, out_t ) # result is (B, out_f, out_t )

    rescale_factor_d = 1.   # for depth
    rescale_factor_f = F_in/out_f # typically 550/48
    rescale_factor_t = T_in/out_t

    matrix_transform = np.diag( np.array([rescale_factor_d, rescale_factor_f, rescale_factor_t]) ) # (3,3) matrix

    # spectrogram = matrix_transform * spectrogram_interpolated
    spectrogram_interpolated = scipy.ndimage.affine_transform(spectrogram, matrix_transform, offset=0, order=1, output_shape=output_shape)
        # we only use linear interpolation because we almost always downsample, and because 2nd order methods and above
        # have a nasty tendency to create small negative local minimas between two strictly positive values
        # we do not want this when we apply a log to the values of the spectrogram

    f_interpolated = scipy.ndimage.affine_transform(f, np.array( [rescale_factor_f] ) , offset=0, order=1, output_shape = (out_f,) )
    t_interpolated = scipy.ndimage.affine_transform(t, np.array( [rescale_factor_t] ) , offset=0, order=1, output_shape = (out_t,) )

    return f_interpolated, t_interpolated, spectrogram_interpolated

def no_interpolation(f, t, spectrogram, out_size):
    """ This function is just a placeholder that mimics the arguments
    of the two previous interpolation functions    """
    return f, t, spectrogram

#%%
#  ---------------- The spectrogram class --------------
class SpectrogramTransform():
    """ create the transform to work with spectrograms. This class behaves
    essentially the same as TempralTransform, except the created transform
    returns a dict of 3d array instead of 2d

    Parameters
    ----------
    signal_name_list: a list of string signals (ex: 'Gyr_y', 'Ori_x')
        If a string ends by "_norm" (ex: "Mag_norm"), the output will
        be the norm of the three (or four) axis of the signal.
    fs: sampling frequency
    duration_window, duration_overlap: duration in sec of spectrogram window and overlap
    spectro_batch_size:
        turning 13,000 temporal signals into (550, 500) array
        spectrograms at once is too much: a single (13000, 550, 500) array,
        with simple precision requires 7.15 Go !
        This is why we work with batches of 1000 instead. For each batch,
        we compute the complete sectrogram (1000 x 550 x 500), then
        interpolate it to smaller sizes, before working wit the following batch.

    interpolation :  string ("log", "linear", "none")
    log_power : bool. If True, the values of the  power spectrum are replaced
        by their log
    out_size : tuple of integer (nb_interp_f, nb_interp_t) = size of spectrogram AFTER interpolation        Is ignored if no interpolation occurs. Default: None
        the spectrogram is computed for 2 1D-arrays: f and t

    flag_debug: flag for print debugging info

    Returns
    -------
    a function with input: data : a dict of (_, 6000) arrays  (key example: 'Gyr_y')
                and output: a dictionnary of 2d arrays.

    """
    def __init__(self, signal_name_list, fs, duration_window, duration_overlap, spectro_batch_size, interpolation,
                 log_power, out_size=None, flag_debug=False):
        super(SpectrogramTransform, self).__init__()

        self.temporal_transform = TemporalTransform(signal_name_list)
        self.fs = fs
        self.duration_window = duration_window
        self.duration_overlap = duration_overlap
        self.spectro_batch_size = spectro_batch_size

        self.signal_name_list = signal_name_list
        self.log_power = log_power
        self.interpolation_name = interpolation

        if interpolation == "linear":
            self.interpolation_fn = interpol_lin
            self.out_size = out_size

        elif interpolation == "log":
            self.interpolation_fn = interpol_log
            self.out_size = out_size

        elif interpolation == "none":
            self.interpolation_fn = no_interpolation
            self.out_size = None

        else :
            raise ValueError("Unknown interpolation: '{}'. Use one of 'log', 'linear', 'none'".format(interpolation))

        # if interpolation == "none" and out_size != None :
        #     warnings.warn("No interpolation is to take place, but an target output size was provided. the output_size argument will be ignored",  Warning)

        self.flag_debug = flag_debug

    def __call__(self, data):
        """
        Parameters
        ----------
        data : a dict of (B, 6000) arrays  (key example: 'Gyr_y')

        Returns
        -------
        a dictionnary of 2d arrays. The keys are from signal_name_list,
            and the values are either spectrograms of raw signals (if the key
            ends with '_x', '_y', '_z', or '_w'); or a spectogram of a norm of
            signals (if the key ends with '_norm'). The shape of the spectrogram
            is (B, F, T), where B (batch size) depends on the input shape, and
            F and T are given by self.out_size
        """

        temporal_signals = self.temporal_transform(data)
        del data  # free some memory

        fs = self.fs

        nperseg     = int(self.duration_window * fs)
        noverlap    = int(self.duration_overlap * fs)

        spectro_batch_size = self.spectro_batch_size
        # turning 13,000 temporal signals into (550, 500) array
            # spectrograms at once is too much: a single (13000, 550, 500) array,
            # with simple precision requires 7.15 Go !
            # This is why we work with batches of 1000 instead. For each batch,
            # we compute the complete sectrogram (1000 x 550 x 500), then
            # interpolate it to smaller sizes, before working wit the following batch.

        out_size = self.out_size

        flag_debug = self.flag_debug

        outputs = {}

        for signal_name in self.signal_name_list:
            current_spectro_batch_size = temporal_signals[signal_name].shape[0]

            if current_spectro_batch_size < spectro_batch_size :
                f, t, spectrogram = scipy.signal.spectrogram(temporal_signals[signal_name], fs=fs, nperseg=nperseg, noverlap=noverlap)

                f_interpolated, t_interpolated, interpolated_spectrogram = self.interpolation_fn(f, t, spectrogram, out_size)
                        # f, t, and possibly out_size will be ignored when the function does not need them

            else :
                n_batches = (current_spectro_batch_size-1)//spectro_batch_size +1

                if out_size is not None:  # we actually compute the interpolation
                    nb_interp_f, nb_interp_t = out_size

                else:    # we only recompute the shapes of the raw spectrogram
                    nb_interp_f = int(duration_window*fs/2) +1
                    nb_interp_t = int((duration_segment-duration_window)/(duration_window-duration_overlap)) +1

                interpolated_spectrogram = np.zeros((current_spectro_batch_size, nb_interp_f, nb_interp_t), dtype='float32')
                for i in range(n_batches):
                    i_min =   i   * spectro_batch_size
                    i_max = (i+1) * spectro_batch_size  # does not matter if it goes beyond current_spectro_batch_size
                    this_temporal_signal = temporal_signals[signal_name][i_min:i_max,:]

                    f, t, spectrogram = scipy.signal.spectrogram(this_temporal_signal, fs=fs, nperseg=nperseg, noverlap=noverlap)

                    if out_size is not None:
                        f_interpolated, t_interpolated, interpolated_spectrogram[i_min:i_max,:,:] = self.interpolation_fn(f, t, spectrogram, out_size) # erase the spectrogram by its interpolation
                    else:
                        f_interpolated, t_interpolated, interpolated_spectrogram[i_min:i_max,:,:] = f, t, spectrogram

            if flag_debug:
                print('f_interpolated:')
                print(f_interpolated[:5])
                print(f_interpolated[-5:])

            del temporal_signals[signal_name]

            if self.log_power :
                np.log(interpolated_spectrogram + 1e-10, dtype='float32', out=interpolated_spectrogram)

            outputs[signal_name] = interpolated_spectrogram

            self.f_interpolated = f_interpolated
            self.t_interpolated = t_interpolated

            # for future debug
            # self.f = f
            # self.t = t
            # self.spectrogram = spectrogram

        return outputs

    def __str__(self):
        """purely for visual purposes, so that we can print() the function"""
        str_to_return = "Spectrogram transform"
        str_to_return += "\n\t Signals: {}".format(self.signal_name_list)
        if self.interpolation_fn != no_interpolation:
            if self.out_size != None :
                str_to_return += "\n\t Output size: {}".format(self.out_size)
            else :
                str_to_return += "\n\t Output size: unchanged"

        str_to_return += "\n\t Interpolation: {}".format(self.interpolation_name)
        str_to_return += "\n\t Log power: {}".format(self.log_power)

        return str_to_return

# end of class SpectrogramTransform():

#%%
if __name__ == "__main__":

    flag_debug = 0
    flag_index = False # if true, plot the spectrogram wrt to index; if False, wrt to Timestamps and Frequencies
    fontdict = {'fontsize':10}
    vmin, vmax = 0, 60
    vmin, vmax = None, None
    n_ticks = 10

    # we plot the raw spectrogram and two interpolated spectrograms for the following classes
    sel_classes = ["Run"]
    nsel = len(sel_classes)
    functions = {"raw spectrogram":           "none",
                 "linear interpolation":      "linear",
                 "logarithmic interpolation": "log"}

    remaining_classes = sel_classes.copy()
    index = 3204  # where to tart the search

    signal_name = "Acc_norm"
    temporal_transform = TemporalTransform([signal_name])  # we will plot the result

    while len(remaining_classes)>0:

        flag_debug = len(remaining_classes) == (nsel-1)

        data_tensor, class_tensor = DS[index]
        data_cpu = {signal:data_tensor[signal].to(torch.device('cpu')).detach().numpy() for signal in data_tensor.keys()}
        class_index = int(class_tensor)

        class_name = classes_names[class_index-1]

        prefix_title = '%s (index=%d)'% (class_name,index)

        if class_name in remaining_classes:
            remaining_classes.remove(class_name)

            temporal_signal = temporal_transform(data_cpu)[signal_name]
            nb = temporal_signal.shape[1]
            x_t = np.linspace(0, nb/fs, nb)

            plt.figure(figsize=(30,10))

            plt.subplot(2,4,1)

            plt.plot(x_t, temporal_signal[0,:])
            plt.title(prefix_title + "\nraw signal : {}".format(signal_name), fontdict)
            plt.xlabel("t (sec)")
            plt.ylabel("Acc (m/s²)")

            index_figure = 2

            # for log_power in [False]:
            for log_power in [False, True]:

                if flag_debug:
                    print('\n log_power = %s' % log_power)

                for f_name in functions :

                    if flag_debug:
                        print('\n f_name = %s' % f_name)

                    function_interpol = functions[f_name]

                    data_tensor, _ = DS[index]  # we need to recreate data because the variable is deleted
                    data_cpu = {signal:data_tensor[signal].to(torch.device('cpu')).detach().numpy() for signal in data_tensor.keys()}

                    spectrogram_transform = SpectrogramTransform([signal_name], fs, duration_window, duration_overlap, spectro_batch_size,
                                                                 function_interpol, log_power, out_size=(48, 48), flag_debug = flag_debug)

                    spectrogram_interpolated = spectrogram_transform(data_cpu)[signal_name]

                    f_interpolated = spectrogram_transform.f_interpolated
                    t_interpolated = spectrogram_transform.t_interpolated

                    error_message_dtype = "One of the spectrograms does not have the correct type: {}, log_power={}, {}. \n dtype should be float32, is actually {}".format(signal_name, str(log_power), f_name, spectrogram_interpolated.dtype)
                    assert (spectrogram_interpolated.dtype == 'float32'), error_message_dtype

                    plt.subplot(2,4,index_figure)

                    if flag_index:
                        ylabel = "f (index)"
                        xlabel = "t (index)"

                        plt.imshow(spectrogram_interpolated[0,:,:])

                    else:
                        ylabel = "f (Hz) "
                        xlabel = "t (s)"

                        t_interpolated = spectrogram_transform.t_interpolated
                        f_interpolated = spectrogram_transform.f_interpolated
                        matrix_shape = spectrogram_interpolated.shape
                        time_list = [f'{t_interpolated[i]:.0f}' for i in np.round(np.linspace(0, matrix_shape[2]-1,n_ticks)).astype(int)]
                        freq_list = [f'{f_interpolated[i]:.1f}' for i in np.round(np.linspace(0, matrix_shape[1]-1,n_ticks)).astype(int)]

                        plt.xticks(np.linspace(0, matrix_shape[2]-1, n_ticks), time_list)
                        plt.yticks(np.linspace(0, matrix_shape[1]-1, n_ticks), freq_list)

                        plt.imshow(spectrogram_interpolated[0,:,:])

                    plt.ylabel(ylabel)
                    plt.xlabel(xlabel)
                    plt.colorbar()

                    index_figure += 1

                    log_power_text = 'log power' if log_power==True else 'power'

                    plt.title("{} of {}".format( log_power_text, f_name), fontdict = {'fontsize':10})

                index_figure += 1 # for the vertical alignment
        index +=1
    plt.show()


import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import scipy.stats as sts
from finn.util.gdrive import *
import sklearn
from sklearn.linear_model import LinearRegression
from sklearn.svm import SVR
from sklearn import model_selection
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import cross_val_score
import os, csv

#define the worksheet name from finn-resource-dashboard
worksheet_name = "FCLayer_resources"
#define the directory name where to save the graphs
directory_name = "FCLayer"

##create the directory
new_dir_path = "../graphs/%s" % directory_name
try:
    os.mkdir(new_dir_path)
except OSError:
    print ("Creation of the directory %s failed" % new_dir_path)
else:
    print ("Successfully created the directory %s " % new_dir_path)

filename  = "db_mem_const.csv"

#get all records from the selected worksheet
list_of_dicts = get_records_from_resource_dashboard(worksheet_name)

# convert list of dicts to dataframe
df = pd.DataFrame(list_of_dicts)
print(df)

fpga = df['FPGA'].iloc[0]

#get synth data
df = df[df.apply(lambda r: r.str.contains('synthesis', case=False).any(), axis=1)]

#get records where act=None
df = df[df['act'].astype(str) == 'None']

#get records where mem_mode=external
df_external = df[df['mem_mode'].astype(str) == 'external']

#get records where mem_mode=const
df_const = df[df['mem_mode'].astype(str) == 'const']

def models(res_class):

    #encode wdt, idt
    labelencoder = LabelEncoder()
    df_training['wdt_encoded'] = labelencoder.fit_transform(df_training['wdt'])
    df_training['idt_encoded'] = labelencoder.fit_transform(df_training['idt'])

    features = ['mh', 'mw', 'pe', 'simd', 'wdt_encoded', 'idt_encoded']
    #features = ['mh', 'mw', 'pe', 'simd']
    #extract features
    X = df_training.loc[:, features].values
    #extract target
    Y = df_training.loc[:, [res_class]].values

    #split the data into train/test data sets 30% testing, 70% training
    X_train, X_test, Y_train, Y_test = model_selection.train_test_split(X, Y, test_size = 0.3, random_state=0)

    #linear regression
    linear_reg_model = LinearRegression()
    linear_reg_model = linear_reg_model.fit(X_train, Y_train)
    Y_predict_linear = linear_reg_model.predict(X_test)
    score_linear = linear_reg_model.score(X_test, Y_test)

    #search for the best SVR hyperparameters
    gscv = GridSearchCV(
        estimator=SVR(kernel='poly', max_iter=20000000),
        param_grid={
            #'C': [0.1, 1, 10, 100, 1000],
            #'epsilon': [0.0001, 0.0005, 0.001, 0.005, 0.01, 0.05, 0.1, 0.5, 1, 5, 10],
            #'gamma': [0.0001, 0.001, 0.005, 0.1, 1, 3, 5]
            'C': [1],
            'epsilon': [0.1],
            'gamma': [0.001]
        },
        cv=5, n_jobs = -1, verbose = 2)

    grid_result = gscv.fit(X_train, Y_train.ravel())

    print("Best parameters set found on development set:")
    print()
    print(grid_result.best_params_)

    #get best hyperparameters and define the model
    best_params = grid_result.best_params_
    best_svr = SVR(kernel='poly', C=best_params["C"], epsilon=best_params["epsilon"], gamma=best_params["gamma"])

    #train the SVR model
    best_svr = best_svr.fit(X_train, Y_train.ravel())

    Y_predict_svr = best_svr.predict(X_test)

    print(X_test)

    print(best_svr.score(X_test, Y_test))

    #cross-validation
    scores = cross_val_score(best_svr, X, Y.ravel(), cv=5)
    print(scores)

    return X_test, Y_test, Y_predict_svr, best_svr.score(X_test, Y_test), Y_predict_linear, score_linear

def generate_graph_test_set(parameter, res_class):

    X_test, Y_test, Y_predict_svr, score_svr, Y_predict_linear, score_linear = models(res_class)

    fig = plt.figure(figsize=(20, 11))
    ax = fig.gca()

    X_test = [col[2] for col in X_test]

    ax.scatter(Y_predict_linear, X_test, marker="o", s=200, facecolors='none', edgecolors='g', label='predicted_linear')
    ax.scatter(Y_predict_svr, X_test, marker="^", s=500, facecolors='none', edgecolors='m', label='predicted_svr')
    ax.scatter(Y_test, X_test, marker="x", s=50, color='r', label='synth')

    ax.set_xlabel("%s" % 'LUT')
    ax.set_ylabel("%s" % 'pe')

    if "FCLayer" in worksheet_name:
        ax.set_title("%s vs %s (SVR_score = %s, linear_score = %s)" % ('pe', 'LUT', score_svr, score_linear))
    else:
        ax.set_title("%s vs %s (SVR_score = %s, linear_score = %s)" % ('pe', 'LUT', score_svr, score_linear))

    leg = ax.legend()

    fig.savefig('../graphs/%s/plot_%s_vs_%s_weight_mem_const_test.png' % (directory_name, 'pe', 'LUT'), bbox_inches='tight')

#get dataframe headers
headers = list(df)

#get all parameters and resource classes from the csv file
parameters = []
res_classes = []
separator = 0
for s in headers:
    if s == 'Resources from:':
        separator = 1
    elif separator:
        res_classes.append(s)
    else:
        parameters.append(s)

#remove tools details: FPGA, finn_commit, vivado_version, vivado_build_no
#remove timing
#remove act, mem_mode, Res from
columns_to_remove = ['act', 'mem_mode', 'Resources from:', 'FPGA', 'finn_commit', 'vivado_version', 'vivado_build_no', 'TargetClockPeriod', 'EstimatedClockPeriod', 'Delay', 'TargetClockFrequency [MHz]', 'EstimatedClockFrequency [MHz]']
res_classes = [element for element in res_classes if element not in columns_to_remove]
parameters = [element for element in parameters if element not in columns_to_remove]

df_const = df_const.drop(columns_to_remove, axis=1)
df_external = df_external.drop(columns_to_remove, axis=1)

print(parameters)
print(res_classes)
print(len(df_external))
print(len(df_const))
print(df_external)
print(df_const)

pd.set_option('display.max_columns', 500)

#isolate contribution of weights to overall resource utilization by subtracting resources of equivalent (mem_mode = external) configuration
df_training = pd.DataFrame(columns=list(df_const))

with open(filename, 'w') as csvfile:
    csvwriter = csv.writer(csvfile)
    headers = list(df_training)
    #headers.insert(0, 'set')
    csvwriter.writerow(headers)

found_row = False
for index, row1 in df_external.iterrows():
    for index, row2 in df_const.iterrows():
        for s in parameters:
            if row1[s] == row2[s]:
                found_row = True
            else:
                found_row = False
                break
        if found_row == True:
            print(row1)
            print(row2)
            """
            with open(filename, 'a+') as csvfile:
                csvwriter = csv.writer(csvfile)
                row = pd.concat([pd.Series(['external']), row1])
                csvwriter.writerow(row)
                row = pd.concat([pd.Series(['const']), row2])
                csvwriter.writerow(row)
            """
            for res in res_classes:
                row2[res] = row2[res] - row1[res]
            df_training = df_training.append(row2, ignore_index=True)

            with open(filename, 'a+') as csvfile:
                csvwriter = csv.writer(csvfile)
                #row2 = pd.concat([pd.Series(['result']), row2])
                csvwriter.writerow(row2)

print(df_training)

#for testing
parameters = ['pe']
res_classes = ['LUT']

for parameter in parameters:
    for res_class in res_classes:
        generate_graph_test_set(parameter, res_class)

"""
####test####
parameters = ['mh', 'mw', 'pe','simd']
res_classes = ['LUT', 'BRAM']

list_columns = ['mh', 'mw', 'pe','simd', 'LUT', 'BRAM']
list_data = [ [16, 16, 8, 8, 345, 34],[16, 16, 2, 2, 456, 45], [32, 32, 16, 16, 786, 77]]
df_ext= pd.DataFrame(columns=list_columns, data=list_data)

list_columns = ['mh', 'mw', 'pe','simd', 'LUT', 'BRAM']
list_data = [ [16, 16, 8, 8, 3345, 234],[32, 32, 16, 16, 5786, 477], [15, 15, 2, 2, 2456, 145]]
df_ct= pd.DataFrame(columns=list_columns, data=list_data)

print(df_ext)
print(df_ct)
df_training = pd.DataFrame(columns=list(df_ct))
found_row = False
for index, row1 in df_ext.iterrows():
    for index, row2 in df_ct.iterrows():
        for s in parameters:
            if row1[s] == row2[s]:
                found_row = True
            else:
                found_row = False
                break
        if found_row == True:
            for res in res_classes:
                row2[res] = row2[res] - row1[res]
            df_training = df_training.append(row2, ignore_index=True)
print(df_training)
"""


<filename>extract.py
import argparse
import os
import numpy as np
import scipy.io as sio
import torch

from configs.default import dataset_cfg
from data import get_test_loader
from models.model import Model

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("gpu", type=int)
    parser.add_argument("model_path", type=str)  # TODO compatible for different models
    parser.add_argument("--img-h", type=int, default=256)
    parser.add_argument("--dataset", type=str, default=None)

    args = parser.parse_args()
    model_path = args.model_path
    fname = model_path.split("/")[-1]

    if args.dataset is not None:
        dataset = args.dataset
    else:
        dataset = model_path.split("/")[1]

    prefix = os.path.splitext(fname)[0]

    dataset_config = dataset_cfg.get(dataset)
    image_size = (args.img_h, 128)

    torch.backends.cudnn.benchmark = True
    os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

    model = Model(eval=True, drop_last_stride=True)

    state_dict = torch.load(model_path)

    model.load_state_dict(state_dict, strict=False)
    model.float()
    model.eval()
    model.cuda()

    # extract query feature
    query = get_test_loader(root=os.path.join(dataset_config.root, dataset_config.query),
                            batch_size=512,
                            image_size=image_size,
                            num_workers=16)

    query_feat = []
    query_label = []
    query_cam_id = []
    query_img_path = []
    for data, label, cam_id, img_path, _ in query:
        with torch.autograd.no_grad():
            feat = model(data.cuda(non_blocking=True))

        query_feat.append(feat.data.cpu().numpy())
        query_label.append(label.data.cpu().numpy())
        query_cam_id.append(cam_id.data.cpu().numpy())
        query_img_path.extend(img_path)

    query_feat = np.concatenate(query_feat, axis=0)
    query_label = np.concatenate(query_label, axis=0)
    query_cam_id = np.concatenate(query_cam_id, axis=0)
    print(query_feat.shape)

    dir_name = "features/{}".format(dataset, prefix)
    if not os.path.isdir(dir_name):
        os.makedirs(dir_name)

    save_name = "{}/query-{}.mat".format(dir_name, prefix)
    sio.savemat(save_name,
                {"feat": query_feat,
                 "ids": query_label,
                 "cam_ids": query_cam_id,
                 "img_path": query_img_path})

    # extract gallery feature
    gallery = get_test_loader(root=os.path.join(dataset_config.root, dataset_config.gallery),
                              batch_size=512,
                              image_size=image_size,
                              num_workers=16)

    gallery_feat = []
    gallery_label = []
    gallery_cam_id = []
    gallery_img_path = []
    for data, label, cam_id, img_path, _ in gallery:
        with torch.autograd.no_grad():
            feat = model(data.cuda(non_blocking=True))

        gallery_feat.append(feat.data.cpu().numpy())
        gallery_label.append(label)
        gallery_cam_id.append(cam_id)
        gallery_img_path.extend(img_path)

    gallery_feat = np.concatenate(gallery_feat, axis=0)
    gallery_label = np.concatenate(gallery_label, axis=0)
    gallery_cam_id = np.concatenate(gallery_cam_id, axis=0)
    print(gallery_feat.shape)

    save_name = "{}/gallery-{}.mat".format(dir_name, prefix)
    sio.savemat(save_name,
                {"feat": gallery_feat,
                 "ids": gallery_label,
                 "cam_ids": gallery_cam_id,
                 "img_path": gallery_img_path})


<reponame>jupito/dwilib
#!/usr/bin/python3

"""Get grid-wise features."""

# TODO Also scale lesiontype.

import argparse
from itertools import product
import logging
import os.path

import numpy as np
from scipy import ndimage

import dwi.files
import dwi.texture
import dwi.util

log = logging.getLogger('grid')

def parse_args():
    """Parse command-line arguments."""
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--verbose', '-v', action='count',
                   help='increase verbosity')
    p.add_argument('--image', required=True,
                   help='input image or pmap')
    p.add_argument('--param', type=int,
                   help='parameter index')
    p.add_argument('--prostate', metavar='MASKFILE', required=True,
                   help='prostate mask')
    p.add_argument('--lesions', metavar='MASKFILE', nargs='+', required=True,
                   help='lesion masks')
    p.add_argument('--mbb', type=float,
                   help='minimum bounding box padding in millimeters (try 15)')
    p.add_argument('--voxelsize', type=float,
                   help='rescaled voxel size in millimeters (try 0.25)')
    p.add_argument('--winsize', type=float, default=5,
                   help='window (cube) size in millimeters (default 5)')
    p.add_argument('--voxelspacing', type=float, nargs=3,
                   help='force voxel spacing (leave out to read from image)')
    p.add_argument('--use_centroid', action='store_true',
                   help='align by prostate centroid instead of image corner')
    p.add_argument('--nanbg', action='store_true',
                   help='set non-prostate background to nan')
    p.add_argument('--lesiontypes', metavar='TYPE', nargs='+',
                   help='lesion types in mask order (CZ or PZ)')
    p.add_argument('--output', metavar='FILENAME', required=True,
                   help='output pmap file')
    return p.parse_args()

def get_lesiontype_array(lesiontypes, lesions):
    """Create lesiontype array. It contains -1 or 1 depending on lesion type,
    or zero where no lesion.
    """
    lesiontype = np.zeros_like(lesions[0], dtype=np.int8)
    if lesiontypes is not None:
        for lt, l in zip(lesiontypes, lesions):
            if lt.lower() == 'cz':
                lesiontype[l] = -1
            elif lt.lower() == 'pz':
                lesiontype[l] = 1
            else:
                raise ValueError('Invalid lesiontype: {}'.format(lt))
    log.info('Lesion types: %s, +1: %i, -1: %i', lesiontypes,
             np.count_nonzero(lesiontype == 1),
             np.count_nonzero(lesiontype == -1))
    return lesiontype

def get_mbb(mask, spacing, pad):
    """Get mask minimum bounding box as slices, with minimum padding in mm."""
    padding = [int(np.ceil(pad / x)) for x in spacing]
    physical_padding = [x * y for x, y in zip(padding, spacing)]
    mbb = dwi.util.bounding_box(mask, padding)
    slices = tuple(slice(*x) for x in mbb)
    log.info('Cropping minimum bounding box, padding: %s', pad)
    log.debug('\tVoxel padding: %s', padding)
    log.debug('\tPhysical padding: %s', physical_padding)
    log.debug('\tMinimum bounding box: %s', mbb)
    return slices

def rescale(img, src_spacing, dst_spacing):
    """Rescale image according to voxel spacing sequences (mm per voxel)."""
    factor = [s/d for s, d in zip(src_spacing, dst_spacing)]
    log.info('Scaling by factor: %s', factor)
    output = ndimage.interpolation.zoom(img, factor, order=0)
    return output

def generate_windows(imageshape, winshape, center):
    """Generate slice objects for a grid of windows around given center.

    Float center will be rounded. Yield a tuple with coordinate slices of each
    window, and window position relative to the center.
    """
    center = [int(round(x)) for x in center]
    starts = [i % w for i, w in zip(center, winshape)]
    stops = [i-w+1 for i, w in zip(imageshape, winshape)]
    its = (range(*x) for x in zip(starts, stops, winshape))
    for coords in product(*its):
        slices = tuple(slice(i, i+w) for i, w in zip(coords, winshape))
        relative = tuple(int((i-c)/w) for i, c, w in zip(coords, center,
                                                         winshape))
        yield slices, relative

def get_datapoint(image, prostate, lesion, lesiontype, stat):
    """Extract output datapoint for a cube.

    If stat is None, median is used. Otherwise, see dwi.texture.stats().
    """
    assert image.shape == prostate.shape == lesion.shape == lesiontype.shape
    if np.isnan(image).all():
        value = np.nan
    else:
        image = image[np.isfinite(image)]  # Remove nan values.
        if stat is None:
            value = np.median(image)
        else:
            value = dwi.texture.stats(image)[stat]
    nneg = np.count_nonzero(lesiontype < 0)
    npos = np.count_nonzero(lesiontype > 0)
    # Label as lesiontype -1 or 1 based on majority, or 0 if no lesion.
    lt = 0
    if nneg > 0:
        lt = -1
    if npos > nneg:
        lt = 1
    return (
        np.count_nonzero(prostate) / prostate.size,
        np.count_nonzero(lesion) / prostate.size,
        lt,
        value,
        )

def create_grid_centroid(metric_winshape, metric_gridshape=(100, 150, 150)):
    """Create and fill grid array based on prostate centroid."""
    gridshape = [int(g//w) for g, w in zip(metric_gridshape, metric_winshape)]
    gridshape = [x + x % 2 for x in gridshape]  # Make any odds even.
    grid = np.full(gridshape + [4], np.nan, dtype=np.float32)
    return grid

def create_grid_corner(image, winshape):
    """Create and fill grid array based on corner."""
    gridshape = [i//w for i, w in zip(image.shape, winshape)]
    grid = np.full(gridshape + [4], np.nan, dtype=np.float32)
    return grid

def process(image, spacing, prostate, lesion, lesiontype, metric_winshape,
            stat, voxelsize=None, use_centroid=False):
    """Process one parameter."""
    # TODO: Should do them all at the same time.
    # Rescale image and masks.
    if voxelsize is not None:
        src_spacing = spacing
        spacing = [voxelsize] * 3
        image = rescale(image, src_spacing, spacing)
        prostate = prostate.astype(np.float_)
        prostate = rescale(prostate, src_spacing, spacing)
        prostate = dwi.util.asbool(prostate)
        lesion = lesion.astype(np.float_)
        lesion = rescale(lesion, src_spacing, spacing)
        lesion = dwi.util.asbool(lesion)
        assert image.shape == prostate.shape == lesion.shape
        # TODO Also scale lesiontype.

        phys_size = [x*y for x, y in zip(image.shape, spacing)]
        log.info('Transformed image: %s %s', image.shape, image.dtype)
        log.info('Voxel spacing: %s, physical size: %s', spacing, phys_size)

    # Extract grid datapoints. Grid placing is based either on prostate
    # centroid, or image corner.
    voxel_winshape = [int(round(x/y)) for x, y in zip(metric_winshape,
                                                      spacing)]
    log.debug('Window metric: %s, voxel: %s', metric_winshape, voxel_winshape)

    centroid = [round(x, 2) for x in dwi.util.centroid(prostate)]
    if use_centroid:
        base = centroid
        grid = create_grid_centroid(metric_winshape)
        grid_base = [s//2 for s in grid.shape]
    else:
        base = [0] * 3
        grid = create_grid_corner(image, voxel_winshape)
        grid_base = [0] * 3
    log.debug('Prostate centroid: %s, base: %s', centroid, base)

    windows = list(generate_windows(image.shape, voxel_winshape, base))
    for slices, relative in windows:
        indices = tuple(c+r for c, r in zip(grid_base, relative))
        values = get_datapoint(image[slices], prostate[slices], lesion[slices],
                               lesiontype[slices], stat)
        grid[indices] = values
    return grid

def average_image(image):
    """Do average filtering for image."""
    d = dict(size=(3, 3), mode='nearest')
    for p in range(image.shape[-1]):
        for i in range(image.shape[0]):
            ix = (i, slice(None), slice(None), p)
            image[ix] = ndimage.filters.median_filter(image[ix], **d)

def indexed_path(path, i):
    """Add an index to path before possible extension."""
    root, ext = os.path.splitext(path)
    return '{r}-{i}{e}'.format(r=root, i=i, e=ext)

def set_loggin(verbosity=0):
    """Set up logging."""
    import sys
    loglevel = logging.INFO if verbosity else logging.WARNING
    logging.basicConfig(level=loglevel, stream=sys.stdout)

def main():
    """Main."""
    args = parse_args()
    set_loggin(verbosity=args.verbose)

    image, attrs = dwi.files.read_pmap(args.image, ondisk=True)
    if args.param is not None:
        image = image[..., args.param]
        image.shape += (1,)
        attrs['parameters'] = [attrs['parameters'][args.param]]
    spacing = attrs['voxel_spacing']

    # Read masks.
    prostate = dwi.files.read_mask(args.prostate,
                                   expected_voxel_spacing=spacing)
    lesions = [dwi.files.read_mask(x, expected_voxel_spacing=spacing,
                                   container=prostate) for x in args.lesions]
    lesion = dwi.util.unify_masks(lesions)

    assert image.shape[:3] == prostate.shape == lesion.shape
    if args.voxelspacing is not None:
        spacing = args.voxelspacing

    phys_size = [x*y for x, y in zip(image.shape[:3], spacing)]
    log.info('Image: %s %s', image.shape, image.dtype)
    log.debug('Voxel spacing: %s, physical size: %s', spacing, phys_size)
    log.debug('Lesions: %i', len(args.lesions))

    lesiontype = get_lesiontype_array(args.lesiontypes, lesions)

    # Crop MBB. The remaining image is stored in memory.
    if args.mbb is None:
        slices = tuple(slice(0, x) for x in image.shape[:3])
    else:
        slices = get_mbb(prostate, spacing, args.mbb)
    image = image[slices]
    prostate = prostate[slices]
    lesion = lesion[slices]
    lesiontype = lesiontype[slices]
    assert (image.shape[:3] == prostate.shape == lesion.shape ==
            lesiontype.shape)

    # average_image(image)

    assert image.ndim == 4, image.ndim
    image = image.astype(np.float32)
    if args.nanbg:
        image[-prostate] = np.nan  # Set background to nan.

    basic = ['prostate', 'lesion', 'lesiontype']
    metric_winshape = [args.winsize] * 3
    if args.param is None:
        params = attrs['parameters']  # Use average of each parameter.
    else:
        params = list(dwi.texture.stats([0]).keys())  # Use statistical feats.
    d = dict(voxelsize=args.voxelsize, use_centroid=args.use_centroid)
    grid = None
    for i, param in enumerate(params):
        if args.param is None:
            img = image[..., i]
            stat = None
        else:
            img = image[..., 0]
            stat = param
        a = process(img, spacing, prostate, lesion, lesiontype,
                    metric_winshape, stat, **d)
        if grid is None:
            shape = a.shape[0:-1] + (len(basic) + len(params),)
            grid = np.empty(shape, dtype=a.dtype)
            log.info('Grid shape: %s', grid.shape)
            grid[..., 0:len(basic)] = a[..., 0:-1]  # Init with basic.
        grid[..., len(basic)+i] = a[..., -1]  # Add each feature.
    outfile = args.output
    attrs = dict(n_lesions=len(args.lesions), spacing=metric_winshape)
    attrs['parameters'] = basic + params
    log.info('Writing %s to %s', grid.shape, outfile)
    dwi.files.write_pmap(outfile, grid, attrs)

if __name__ == '__main__':
    main()