Blame view
build/lib.linux-x86_64-2.7/complexnn/pool.py
4.76 KB
f2d3bd141
Initial commit wi...
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Authors: Olexa Bilaniuk
#
import keras.backend as KB
import keras.engine as KE
import keras.layers as KL
import keras.optimizers as KO
import theano as T
import theano.ifelse as TI
import theano.tensor as TT
import theano.tensor.fft as TTF
import numpy as np
#
# Spectral Pooling Layer
#
class SpectralPooling1D(KL.Layer):
def __init__(self, topf=(0,)):
super(SpectralPooling1D, self).__init__()
if "topf" in kwargs:
self.topf = (int (kwargs["topf" ][0]),)
self.topf = (self.topf[0]//2,)
elif "gamma" in kwargs:
self.gamma = (float(kwargs["gamma"][0]),)
self.gamma = (self.gamma[0]/2,)
else:
raise RuntimeError("Must provide either topf= or gamma= !")
def call(self, x, mask=None):
xshape = x._keras_shape
if hasattr(self, "topf"):
topf = self.topf
else:
if KB.image_data_format() == "channels_first":
topf = (int(self.gamma[0]*xshape[2]),)
else:
topf = (int(self.gamma[0]*xshape[1]),)
if KB.image_data_format() == "channels_first":
if topf[0] > 0 and xshape[2] >= 2*topf[0]:
mask = [1]*(topf[0] ) +\
[0]*(xshape[2] - 2*topf[0]) +\
[1]*(topf[0] )
mask = [[mask]]
mask = np.asarray(mask, dtype=KB.floatx()).transpose((0,1,2))
mask = KB.constant(mask)
x *= mask
else:
if topf[0] > 0 and xshape[1] >= 2*topf[0]:
mask = [1]*(topf[0] ) +\
[0]*(xshape[1] - 2*topf[0]) +\
[1]*(topf[0] )
mask = [[mask]]
mask = np.asarray(mask, dtype=KB.floatx()).transpose((0,2,1))
mask = KB.constant(mask)
x *= mask
return x
class SpectralPooling2D(KL.Layer):
def __init__(self, **kwargs):
super(SpectralPooling2D, self).__init__()
if "topf" in kwargs:
self.topf = (int (kwargs["topf" ][0]), int (kwargs["topf" ][1]))
self.topf = (self.topf[0]//2, self.topf[1]//2)
elif "gamma" in kwargs:
self.gamma = (float(kwargs["gamma"][0]), float(kwargs["gamma"][1]))
self.gamma = (self.gamma[0]/2, self.gamma[1]/2)
else:
raise RuntimeError("Must provide either topf= or gamma= !")
def call(self, x, mask=None):
xshape = x._keras_shape
if hasattr(self, "topf"):
topf = self.topf
else:
if KB.image_data_format() == "channels_first":
topf = (int(self.gamma[0]*xshape[2]), int(self.gamma[1]*xshape[3]))
else:
topf = (int(self.gamma[0]*xshape[1]), int(self.gamma[1]*xshape[2]))
if KB.image_data_format() == "channels_first":
if topf[0] > 0 and xshape[2] >= 2*topf[0]:
mask = [1]*(topf[0] ) +\
[0]*(xshape[2] - 2*topf[0]) +\
[1]*(topf[0] )
mask = [[[mask]]]
mask = np.asarray(mask, dtype=KB.floatx()).transpose((0,1,3,2))
mask = KB.constant(mask)
x *= mask
if topf[1] > 0 and xshape[3] >= 2*topf[1]:
mask = [1]*(topf[1] ) +\
[0]*(xshape[3] - 2*topf[1]) +\
[1]*(topf[1] )
mask = [[[mask]]]
mask = np.asarray(mask, dtype=KB.floatx()).transpose((0,1,2,3))
mask = KB.constant(mask)
x *= mask
else:
if topf[0] > 0 and xshape[1] >= 2*topf[0]:
mask = [1]*(topf[0] ) +\
[0]*(xshape[1] - 2*topf[0]) +\
[1]*(topf[0] )
mask = [[[mask]]]
mask = np.asarray(mask, dtype=KB.floatx()).transpose((0,3,1,2))
mask = KB.constant(mask)
x *= mask
if topf[1] > 0 and xshape[2] >= 2*topf[1]:
mask = [1]*(topf[1] ) +\
[0]*(xshape[2] - 2*topf[1]) +\
[1]*(topf[1] )
mask = [[[mask]]]
mask = np.asarray(mask, dtype=KB.floatx()).transpose((0,1,3,2))
mask = KB.constant(mask)
x *= mask
return x
if __name__ == "__main__":
import cv2, sys
import __main__ as SP
import fft as CF
# Build Model
x = i = KL.Input(shape=(6,512,512))
f = CF.FFT2()(x)
p = SP.SpectralPooling2D(gamma=[0.15,0.15])(f)
o = CF.IFFT2()(p)
model = KE.Model([i], [f,p,o])
model.compile("sgd", "mse")
# Use it
img = cv2.imread(sys.argv[1])
imgBatch = img[np.newaxis,...].transpose((0,3,1,2))
imgBatch = np.concatenate([imgBatch, np.zeros_like(imgBatch)], axis=1)
f,p,o = model.predict(imgBatch)
ffted = np.sqrt(np.sum(f[:,:3]**2 + f[:,3:]**2, axis=1))
ffted = ffted .transpose((1,2,0))/255
pooled = np.sqrt(np.sum(p[:,:3]**2 + p[:,3:]**2, axis=1))
pooled = pooled.transpose((1,2,0))/255
filtered = np.clip(o,0,255).transpose((0,2,3,1))[0,:,:,:3].astype("uint8")
# Display it
cv2.imshow("Original", img)
cv2.imshow("FFT", ffted)
cv2.imshow("Pooled", pooled)
cv2.imshow("Filtered", filtered)
cv2.waitKey(0)
|