Spaces:
Runtime error
Runtime error
update files
Browse files- MT.py +338 -0
- Model_example.pth.tar +3 -0
- PatternCell_2.mp4 +0 -0
- UnclassifiedCell_2.mp4 +0 -0
- V1.py +909 -0
MT.py
ADDED
@@ -0,0 +1,338 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import torch.nn as nn
|
3 |
+
import torch.nn.functional as F
|
4 |
+
import math
|
5 |
+
import numpy as np
|
6 |
+
|
7 |
+
|
8 |
+
class ConvGRU(nn.Module):
|
9 |
+
def __init__(self, hidden_dim=128, input_dim=192 + 128):
|
10 |
+
super(ConvGRU, self).__init__()
|
11 |
+
self.convz = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
|
12 |
+
self.convr = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
|
13 |
+
self.convq = nn.Conv2d(hidden_dim + input_dim, hidden_dim, 3, padding=1)
|
14 |
+
|
15 |
+
def forward(self, h, x):
|
16 |
+
hx = torch.cat([h, x], dim=1)
|
17 |
+
|
18 |
+
z = torch.sigmoid(self.convz(hx))
|
19 |
+
r = torch.sigmoid(self.convr(hx))
|
20 |
+
q = torch.tanh(self.convq(torch.cat([r * h, x], dim=1)))
|
21 |
+
|
22 |
+
h = (1 - z) * h + z * q
|
23 |
+
return h
|
24 |
+
|
25 |
+
|
26 |
+
class SepConvGRU(nn.Module):
|
27 |
+
def __init__(self, hidden_dim=128, input_dim=192 + 128):
|
28 |
+
super(SepConvGRU, self).__init__()
|
29 |
+
self.convz1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2))
|
30 |
+
self.convr1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2))
|
31 |
+
self.convq1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2))
|
32 |
+
|
33 |
+
self.convz2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0))
|
34 |
+
self.convr2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0))
|
35 |
+
self.convq2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0))
|
36 |
+
|
37 |
+
def forward(self, h, x):
|
38 |
+
# horizontal
|
39 |
+
hx = torch.cat([h, x], dim=1)
|
40 |
+
z = torch.sigmoid(self.convz1(hx))
|
41 |
+
r = torch.sigmoid(self.convr1(hx))
|
42 |
+
q = torch.tanh(self.convq1(torch.cat([r * h, x], dim=1)))
|
43 |
+
h = (1 - z) * h + z * q
|
44 |
+
|
45 |
+
# vertical
|
46 |
+
hx = torch.cat([h, x], dim=1)
|
47 |
+
z = torch.sigmoid(self.convz2(hx))
|
48 |
+
r = torch.sigmoid(self.convr2(hx))
|
49 |
+
q = torch.tanh(self.convq2(torch.cat([r * h, x], dim=1)))
|
50 |
+
h = (1 - z) * h + z * q
|
51 |
+
|
52 |
+
return h
|
53 |
+
|
54 |
+
|
55 |
+
class GRU(nn.Module):
|
56 |
+
def __init__(self, hidden_dim=128, input_dim=192 + 128):
|
57 |
+
super(GRU, self).__init__()
|
58 |
+
self.convz1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2))
|
59 |
+
self.convr1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2))
|
60 |
+
self.convq1 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (1, 5), padding=(0, 2))
|
61 |
+
|
62 |
+
self.convz2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0))
|
63 |
+
self.convr2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0))
|
64 |
+
self.convq2 = nn.Conv2d(hidden_dim + input_dim, hidden_dim, (5, 1), padding=(2, 0))
|
65 |
+
|
66 |
+
def forward(self, hidden, x, shape):
|
67 |
+
# horizontal
|
68 |
+
b, l, c = hidden.shape
|
69 |
+
h, w = shape
|
70 |
+
hidden = hidden.view(b, h, w, c).permute(0, 3, 1, 2).contiguous()
|
71 |
+
x = x.view(b, h, w, c).permute(0, 3, 1, 2).contiguous()
|
72 |
+
|
73 |
+
hx = torch.cat([hidden, x], dim=1)
|
74 |
+
z = torch.sigmoid(self.convz1(hx))
|
75 |
+
r = torch.sigmoid(self.convr1(hx))
|
76 |
+
q = torch.tanh(self.convq1(torch.cat([r * hidden, x], dim=1)))
|
77 |
+
hidden = (1 - z) * hidden + z * q
|
78 |
+
|
79 |
+
# vertical
|
80 |
+
hx = torch.cat([hidden, x], dim=1)
|
81 |
+
z = torch.sigmoid(self.convz2(hx))
|
82 |
+
r = torch.sigmoid(self.convr2(hx))
|
83 |
+
q = torch.tanh(self.convq2(torch.cat([r * hidden, x], dim=1)))
|
84 |
+
hidden = (1 - z) * hidden + z * q
|
85 |
+
|
86 |
+
return hidden.flatten(-2).permute(0, 2, 1)
|
87 |
+
|
88 |
+
|
89 |
+
class PositionEmbeddingSine(nn.Module):
|
90 |
+
"""
|
91 |
+
This is a more standard version of the position embedding, very similar to the one
|
92 |
+
used by the Attention is all you need paper, generalized to work on images.
|
93 |
+
"""
|
94 |
+
|
95 |
+
def __init__(self, num_pos_feats=64, temperature=10000, normalize=True, scale=None):
|
96 |
+
super().__init__()
|
97 |
+
self.num_pos_feats = num_pos_feats
|
98 |
+
self.temperature = temperature
|
99 |
+
self.normalize = normalize
|
100 |
+
if scale is not None and normalize is False:
|
101 |
+
raise ValueError("normalize should be True if scale is passed")
|
102 |
+
if scale is None:
|
103 |
+
scale = 2 * math.pi
|
104 |
+
self.scale = scale
|
105 |
+
|
106 |
+
def forward(self, x):
|
107 |
+
# x = tensor_list.tensors # [B, C, H, W]
|
108 |
+
# mask = tensor_list.mask # [B, H, W], input with padding, valid as 0
|
109 |
+
b, c, h, w = x.size()
|
110 |
+
mask = torch.ones((b, h, w), device=x.device) # [B, H, W]
|
111 |
+
y_embed = mask.cumsum(1, dtype=torch.float32)
|
112 |
+
x_embed = mask.cumsum(2, dtype=torch.float32)
|
113 |
+
#
|
114 |
+
# y_embed = (y_embed / 2) ** 2
|
115 |
+
# x_embed = (x_embed / 2) ** 2
|
116 |
+
|
117 |
+
if self.normalize:
|
118 |
+
eps = 1e-6
|
119 |
+
y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
|
120 |
+
x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
|
121 |
+
|
122 |
+
# using an exponential
|
123 |
+
dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
|
124 |
+
dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
|
125 |
+
|
126 |
+
pos_x = x_embed[:, :, :, None] / dim_t
|
127 |
+
pos_y = y_embed[:, :, :, None] / dim_t
|
128 |
+
pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
|
129 |
+
pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
|
130 |
+
pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
|
131 |
+
return pos
|
132 |
+
|
133 |
+
|
134 |
+
def feature_add_position(feature0, feature_channels, scale=1.0):
|
135 |
+
temp = torch.mean(abs(feature0))
|
136 |
+
pos_enc = PositionEmbeddingSine(num_pos_feats=feature_channels // 2)
|
137 |
+
# position = PositionalEncodingPermute2D(feature_channels)(feature0)
|
138 |
+
position = pos_enc(feature0)
|
139 |
+
feature0 = feature0 + (temp * position / position.mean()) * scale * torch.pi
|
140 |
+
feature0 = feature0 * temp / torch.mean(abs(feature0), dim=(1, 2, 3), keepdim=True)
|
141 |
+
return feature0
|
142 |
+
|
143 |
+
|
144 |
+
def feature_add_image_content(feature0, add_fea, scale=0.4):
|
145 |
+
temp = torch.mean(abs(feature0))
|
146 |
+
position = add_fea
|
147 |
+
feature0 = feature0 + (temp * position / position.mean()) * scale * torch.pi
|
148 |
+
feature0 = feature0 * temp / torch.mean(abs(feature0), dim=(1, 2, 3), keepdim=True)
|
149 |
+
return feature0
|
150 |
+
|
151 |
+
|
152 |
+
class AttUp(nn.Module):
|
153 |
+
def __init__(self,
|
154 |
+
c=512
|
155 |
+
):
|
156 |
+
super(AttUp, self).__init__()
|
157 |
+
self.proj = nn.Linear(c, c, bias=False)
|
158 |
+
self.norm = nn.LayerNorm(c)
|
159 |
+
self.conv = nn.Sequential(nn.Conv2d(2 * c, c, kernel_size=1, stride=1, padding=0),
|
160 |
+
nn.GELU(),
|
161 |
+
nn.Conv2d(c, c, kernel_size=3, stride=1, padding=1),
|
162 |
+
nn.GELU(),
|
163 |
+
nn.Conv2d(c, c, kernel_size=3, stride=1, padding=1),
|
164 |
+
nn.GELU()
|
165 |
+
)
|
166 |
+
self.gru = SepConvGRU(c, c)
|
167 |
+
|
168 |
+
def forward(self, att, message, shape):
|
169 |
+
# q, k, v: [B, L, C]
|
170 |
+
b, l, c = att.shape
|
171 |
+
h, w = shape
|
172 |
+
message = self.norm(self.proj(message)).view(b, h, w, c).permute(0, 3, 1, 2).contiguous()
|
173 |
+
att = att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous()
|
174 |
+
message = self.conv(torch.cat([att, message], dim=1))
|
175 |
+
att = self.gru(att, message).flatten(-2).permute(0, 2, 1)
|
176 |
+
# [B, H*W, C]
|
177 |
+
return att
|
178 |
+
|
179 |
+
|
180 |
+
class TransformerLayer(nn.Module):
|
181 |
+
def __init__(self,
|
182 |
+
d_model=256,
|
183 |
+
nhead=1,
|
184 |
+
no_ffn=False,
|
185 |
+
ffn_dim_expansion=4
|
186 |
+
):
|
187 |
+
super(TransformerLayer, self).__init__()
|
188 |
+
|
189 |
+
self.dim = d_model
|
190 |
+
self.nhead = nhead
|
191 |
+
self.no_ffn = no_ffn
|
192 |
+
# multi-head attention
|
193 |
+
self.att_proj = nn.Sequential(nn.Linear(d_model, d_model, bias=False), nn.ReLU(inplace=True),
|
194 |
+
nn.Linear(d_model, d_model, bias=False))
|
195 |
+
self.v_proj = nn.Linear(d_model, d_model, bias=False)
|
196 |
+
self.merge = nn.Linear(d_model, d_model, bias=False)
|
197 |
+
self.gru = GRU(d_model, d_model)
|
198 |
+
self.attn_updater = AttUp(d_model)
|
199 |
+
self.drop = nn.Dropout(p=0.8)
|
200 |
+
|
201 |
+
self.norm1 = nn.LayerNorm(d_model)
|
202 |
+
|
203 |
+
# no ffn after self-attn, with ffn after cross-attn
|
204 |
+
if not self.no_ffn:
|
205 |
+
in_channels = d_model * 2
|
206 |
+
self.mlp = nn.Sequential(
|
207 |
+
nn.Linear(in_channels, in_channels * ffn_dim_expansion, bias=False),
|
208 |
+
nn.GELU(),
|
209 |
+
nn.Linear(in_channels * ffn_dim_expansion, in_channels * ffn_dim_expansion, bias=False),
|
210 |
+
nn.GELU(),
|
211 |
+
nn.Linear(in_channels * ffn_dim_expansion, d_model, bias=False),
|
212 |
+
)
|
213 |
+
|
214 |
+
self.norm2 = nn.LayerNorm(d_model)
|
215 |
+
|
216 |
+
def forward(self, att, value,
|
217 |
+
shape, iteration=0):
|
218 |
+
# source, target: [B, L, C]
|
219 |
+
max_exp_scale = 3 * torch.pi
|
220 |
+
# single-head attention
|
221 |
+
B, L, C = value.shape
|
222 |
+
if iteration == 0:
|
223 |
+
att = feature_add_position(att.transpose(-1, -2).view(
|
224 |
+
B, C, shape[0], shape[1]), C).reshape(B, C, -1).transpose(-1, -2)
|
225 |
+
|
226 |
+
# att = feature_add_position(att.transpose(-1, -2).view(
|
227 |
+
# B, C, shape[0], shape[1]), C).reshape(B, C, -1).transpose(-1, -2)
|
228 |
+
val_proj = self.v_proj(value)
|
229 |
+
att_proj = self.att_proj(att) # [B, L, C]
|
230 |
+
norm_fac = torch.sum(att_proj ** 2, dim=-1, keepdim=True) ** 0.5
|
231 |
+
scale = max_exp_scale * torch.sigmoid(torch.mean(att_proj, dim=[-1, -2], keepdim=True)) + 1
|
232 |
+
A = torch.exp(scale * torch.matmul(att_proj / norm_fac, att_proj.permute(0, 2, 1) / norm_fac.permute(0, 2, 1)))
|
233 |
+
A = A / A.max()
|
234 |
+
# I = torch.eye(A.shape[-1], device=A.device).unsqueeze(0)
|
235 |
+
# # A[I.repeat(B, 1, 1) == 1] = 1e-6 # remove self-prop
|
236 |
+
D = torch.sum(A, dim=-1, keepdim=True)
|
237 |
+
D = 1 / (torch.sqrt(D) + 1e-6) # normalized node degrees
|
238 |
+
A = D * A * D.transpose(-1, -2)
|
239 |
+
|
240 |
+
# A = torch.softmax(A , dim=2) # [B, L, L]
|
241 |
+
message = torch.matmul(A, val_proj) # [B, L, C]
|
242 |
+
|
243 |
+
message = self.merge(message) # [B, L, C]
|
244 |
+
message = self.norm1(message)
|
245 |
+
if not self.no_ffn:
|
246 |
+
message = self.mlp(torch.cat([value, message], dim=-1))
|
247 |
+
message = self.norm2(message)
|
248 |
+
|
249 |
+
# if iteration > 2:
|
250 |
+
# message = self.drop(message)
|
251 |
+
|
252 |
+
att = self.attn_updater(att, message, shape)
|
253 |
+
value = self.gru(value, message, shape)
|
254 |
+
return value, att, A
|
255 |
+
|
256 |
+
|
257 |
+
class FeatureTransformer(nn.Module):
|
258 |
+
def __init__(self,
|
259 |
+
num_layers=6,
|
260 |
+
d_model=128
|
261 |
+
):
|
262 |
+
super(FeatureTransformer, self).__init__()
|
263 |
+
self.d_model = d_model
|
264 |
+
# self.layers = nn.ModuleList([TransformerLayer(self.d_model, no_ffn=False, ffn_dim_expansion=2)
|
265 |
+
# for i in range(num_layers)])
|
266 |
+
self.layers = TransformerLayer(self.d_model, no_ffn=False, ffn_dim_expansion=2)
|
267 |
+
self.re_proj = nn.Sequential(nn.Linear(d_model, d_model), nn.GELU(), nn.Linear(d_model, d_model))
|
268 |
+
self.num_layers = num_layers
|
269 |
+
self.norm_sigma = nn.Parameter(torch.tensor(1.0, requires_grad=True), requires_grad=True)
|
270 |
+
self.norm_k = nn.Parameter(torch.tensor(1.8, requires_grad=True), requires_grad=True)
|
271 |
+
|
272 |
+
for p in self.parameters():
|
273 |
+
if p.dim() > 1:
|
274 |
+
nn.init.xavier_uniform_(p)
|
275 |
+
|
276 |
+
def normalize(self, x): # TODO
|
277 |
+
sum_activation = torch.mean(x, dim=[1, 2], keepdim=True) + torch.square(self.norm_sigma)
|
278 |
+
x = self.norm_k.abs() * x / sum_activation
|
279 |
+
return x
|
280 |
+
|
281 |
+
def forward(self, feature0):
|
282 |
+
|
283 |
+
feature_list = []
|
284 |
+
attn_list = []
|
285 |
+
attn_viz_list = []
|
286 |
+
b, c, h, w = feature0.shape
|
287 |
+
assert self.d_model == c
|
288 |
+
value = feature0.flatten(-2).permute(0, 2, 1) # [B, H*W, C]
|
289 |
+
att = feature0
|
290 |
+
att = att.flatten(-2).permute(0, 2, 1) # [B, H*W, C]
|
291 |
+
for i in range(self.num_layers):
|
292 |
+
value, att, attn_viz = self.layers(att=att, value=value, shape=[h, w], iteration=i)
|
293 |
+
attn_viz = attn_viz.reshape(b, h, w, h, w)
|
294 |
+
attn_viz_list.append(attn_viz)
|
295 |
+
value_decode = self.normalize(
|
296 |
+
torch.square(self.re_proj(value))) # map to motion energy, Do use normalization here
|
297 |
+
# print("value_decode",value_decode.abs().mean())
|
298 |
+
attn_list.append(att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
|
299 |
+
feature_list.append(value_decode.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
|
300 |
+
# reshape back
|
301 |
+
return feature_list, attn_list, attn_viz_list
|
302 |
+
|
303 |
+
def forward_save_mem(self, feature0, add_position_embedding=True):
|
304 |
+
feature_list = []
|
305 |
+
attn_list = []
|
306 |
+
attn_viz_list = []
|
307 |
+
b, c, h, w = feature0.shape
|
308 |
+
assert self.d_model == c
|
309 |
+
value = feature0.flatten(-2).permute(0, 2, 1) # [B, H*W, C]
|
310 |
+
att = feature0
|
311 |
+
att = att.flatten(-2).permute(0, 2, 1) # [B, H*W, C]
|
312 |
+
for i in range(self.num_layers):
|
313 |
+
value, att, _ = self.layers(att=att, value=value, shape=[h, w], iteration=i)
|
314 |
+
value_decode = self.normalize(
|
315 |
+
torch.square(self.re_proj(value))) # map to motion energy, Do use normalization here
|
316 |
+
# print("value_decode",value_decode.abs().mean())
|
317 |
+
attn_list.append(att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
|
318 |
+
feature_list.append(value_decode.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
|
319 |
+
# reshape back
|
320 |
+
return feature_list, attn_list
|
321 |
+
|
322 |
+
@staticmethod
|
323 |
+
def demo():
|
324 |
+
import time
|
325 |
+
frame_list = torch.randn([4, 256, 64, 64], device="cuda")
|
326 |
+
model = FeatureTransformer(6, 256).cuda()
|
327 |
+
for i in range(100):
|
328 |
+
start = time.time()
|
329 |
+
output = model(frame_list)
|
330 |
+
|
331 |
+
torch.mean(output[-1][-1]).backward()
|
332 |
+
end = time.time()
|
333 |
+
print(end - start)
|
334 |
+
print("#================================++#")
|
335 |
+
|
336 |
+
|
337 |
+
if __name__ == '__main__':
|
338 |
+
FeatureTransformer.demo()
|
Model_example.pth.tar
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
version https://git-lfs.github.com/spec/v1
|
2 |
+
oid sha256:666be808823bae9a29493cb967e8ba233304ff7eaf962133bf6e6499e9c42346
|
3 |
+
size 58749697
|
PatternCell_2.mp4
ADDED
Binary file (138 kB). View file
|
|
UnclassifiedCell_2.mp4
ADDED
Binary file (186 kB). View file
|
|
V1.py
ADDED
@@ -0,0 +1,909 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import numpy as np
|
2 |
+
import math
|
3 |
+
import torch
|
4 |
+
from io import BytesIO
|
5 |
+
import numpy
|
6 |
+
from torch import nn
|
7 |
+
from torch.nn import functional as F
|
8 |
+
import matplotlib.pyplot as plt
|
9 |
+
import os
|
10 |
+
import pandas as pd
|
11 |
+
import imageio
|
12 |
+
from torch.cuda.amp import autocast as autocast
|
13 |
+
|
14 |
+
|
15 |
+
def cart2pol(x, y):
|
16 |
+
rho = np.sqrt(x ** 2 + y ** 2)
|
17 |
+
phi = np.arctan2(y, x)
|
18 |
+
return (rho, phi)
|
19 |
+
|
20 |
+
|
21 |
+
def pol2cart(rho, phi):
|
22 |
+
x = rho * np.cos(phi)
|
23 |
+
y = rho * np.sin(phi)
|
24 |
+
return (x, y)
|
25 |
+
|
26 |
+
|
27 |
+
def inverse_sigmoid(p):
|
28 |
+
return np.log(p / (1 - p))
|
29 |
+
|
30 |
+
|
31 |
+
def artanh(y):
|
32 |
+
return 0.5 * np.log((1 + y) / (1 - y))
|
33 |
+
|
34 |
+
|
35 |
+
class V1(nn.Module):
|
36 |
+
"""each input includes 10 frame with 25 frame/sec sampling rate
|
37 |
+
temporal window size = 5 frame(200ms)
|
38 |
+
spatial window size = 5*2 + 1 = 11
|
39 |
+
spatial filter is
|
40 |
+
lambda is frequency of cos wave
|
41 |
+
"""
|
42 |
+
|
43 |
+
def __init__(self, spatial_num=32, scale_num=8, scale_factor=16, kernel_radius=7, num_ft=32,
|
44 |
+
kernel_size=6, average_time=True):
|
45 |
+
super(V1, self).__init__()
|
46 |
+
|
47 |
+
def make_param(in_channels, values, requires_grad=True, dtype=None):
|
48 |
+
if dtype is None:
|
49 |
+
dtype = 'float32'
|
50 |
+
values = numpy.require(values, dtype=dtype)
|
51 |
+
n = in_channels * len(values)
|
52 |
+
data = torch.from_numpy(values).view(1, -1)
|
53 |
+
data = data.repeat(in_channels, 1)
|
54 |
+
return torch.nn.Parameter(data=data, requires_grad=requires_grad)
|
55 |
+
|
56 |
+
assert spatial_num == num_ft
|
57 |
+
scale_each_level = np.exp(1 / (scale_num - 1) * np.log(1 / scale_factor))
|
58 |
+
self.scale_each_level = scale_each_level
|
59 |
+
self.scale_num = scale_num
|
60 |
+
self.cell_index = 0
|
61 |
+
self.spatial_filter = nn.ModuleList([GaborFilters(kernel_radius=kernel_radius, num_units=spatial_num,random=False)
|
62 |
+
for i in range(scale_num)])
|
63 |
+
self.temporal_decay = 0.2
|
64 |
+
self.spatial_decay = 0.2
|
65 |
+
|
66 |
+
self.spatial_radius = kernel_radius
|
67 |
+
self.spatial_kernel_size = kernel_radius * 2 + 1
|
68 |
+
self.spatial_num = spatial_num
|
69 |
+
self.temporal_filter = nn.ModuleList([TemporalFilter(num_ft=num_ft, kernel_size=kernel_size, random=False)
|
70 |
+
for i in range(scale_num)]) # 16 filter
|
71 |
+
|
72 |
+
self.n_frames = 11
|
73 |
+
self._num_after_st = spatial_num * scale_num
|
74 |
+
if not average_time:
|
75 |
+
self._num_after_st = self._num_after_st * (self.n_frames - kernel_size + 1)
|
76 |
+
if average_time:
|
77 |
+
self.temporal_pooling = make_param(self._num_after_st, np.ones((self.n_frames - kernel_size + 1)),
|
78 |
+
requires_grad=True)
|
79 |
+
# TODO: concentrate on middle frame
|
80 |
+
|
81 |
+
self.temporal_pooling = make_param(self._num_after_st, [0.05, 0.1, 0.4, 0.4, 0.1, 0.05],
|
82 |
+
requires_grad=True)
|
83 |
+
|
84 |
+
self.norm_sigma = make_param(1, np.array([0.2]), requires_grad=True)
|
85 |
+
self.spontaneous_firing = make_param(1, np.array([0.3]), requires_grad=True)
|
86 |
+
self.norm_k = make_param(1, np.array([4.0]), requires_grad=True)
|
87 |
+
self._average_time = average_time
|
88 |
+
self.t_sin = None
|
89 |
+
self.t_cos = None
|
90 |
+
self.s_sin = None
|
91 |
+
self.s_cos = None
|
92 |
+
|
93 |
+
def infer_scale(self, x, scale): # x should be list of B,1,H,W
|
94 |
+
energy_list = []
|
95 |
+
n = len(x)
|
96 |
+
B, C, H, W = x[0].shape
|
97 |
+
x = [img.unsqueeze(0) for img in x]
|
98 |
+
x = torch.cat(x, dim=0).reshape(n * B, C, H, W)
|
99 |
+
|
100 |
+
sy = x.size(2)
|
101 |
+
sx = x.size(3)
|
102 |
+
s_sin = self.s_sin
|
103 |
+
s_cos = self.s_cos
|
104 |
+
|
105 |
+
gb_sin = s_sin.view(self.spatial_num, 1, self.spatial_kernel_size, self.spatial_kernel_size)
|
106 |
+
gb_cos = s_cos.view(self.spatial_num, 1, self.spatial_kernel_size, self.spatial_kernel_size)
|
107 |
+
|
108 |
+
# flip kernel
|
109 |
+
gb_sin = torch.flip(gb_sin, dims=[-1, -2])
|
110 |
+
gb_cos = torch.flip(gb_cos, dims=[-1, -2])
|
111 |
+
|
112 |
+
res_sin = F.conv2d(input=x, weight=gb_sin,
|
113 |
+
padding=self.spatial_radius, groups=1)
|
114 |
+
res_cos = F.conv2d(input=x, weight=gb_cos,
|
115 |
+
padding=self.spatial_radius, groups=1)
|
116 |
+
|
117 |
+
res_sin = res_sin.view(B, -1, sy, sx)
|
118 |
+
res_cos = res_cos.view(B, -1, sy, sx)
|
119 |
+
g_asin_list = res_sin.reshape(n, B, -1, H, W)
|
120 |
+
g_acos_list = res_cos.reshape(n, B, -1, H, W)
|
121 |
+
|
122 |
+
for channel in range(self.spatial_filter[0].n_channels_post_conv):
|
123 |
+
k_sin = self.t_sin[channel, ...][None]
|
124 |
+
k_cos = self.t_cos[channel, ...][None]
|
125 |
+
# spatial filter
|
126 |
+
g_asin, g_acos = g_asin_list[:, :, channel, :, :], g_acos_list[:, :, channel, :, :] # n,b,h,w
|
127 |
+
g_asin = g_asin.reshape(n, B * H * W, 1).permute(1, 2, 0) # bhw,1,n
|
128 |
+
g_acos = g_acos.reshape(n, B * H * W, 1).permute(1, 2, 0)
|
129 |
+
|
130 |
+
# reverse the impulse response
|
131 |
+
k_sin = torch.flip(k_sin, dims=(-1,))
|
132 |
+
k_cos = torch.flip(k_cos, dims=(-1,))
|
133 |
+
#
|
134 |
+
a = F.conv1d(g_acos, k_sin, padding="valid", bias=None)
|
135 |
+
b = F.conv1d(g_asin, k_cos, padding="valid", bias=None)
|
136 |
+
g_o = a + b
|
137 |
+
a = F.conv1d(g_acos, k_cos, padding="valid", bias=None)
|
138 |
+
b = F.conv1d(g_asin, k_sin, padding="valid", bias=None)
|
139 |
+
g_e = a - b
|
140 |
+
energy_component = g_o ** 2 + g_e ** 2 + self.spontaneous_firing.square()
|
141 |
+
energy_component = energy_component.reshape(B, H, W, a.size(-1)).permute(0, 3, 1, 2)
|
142 |
+
if self._average_time: # average motion energy across time
|
143 |
+
total_channel = scale * self.spatial_num + channel
|
144 |
+
pooling = self.temporal_pooling[total_channel][None, ..., None, None]
|
145 |
+
energy_component = abs(torch.mean(energy_component * pooling, dim=1, keepdim=True))
|
146 |
+
energy_list.append(energy_component)
|
147 |
+
energy_list = torch.cat(energy_list, dim=1)
|
148 |
+
return energy_list
|
149 |
+
|
150 |
+
def forward(self, image_list):
|
151 |
+
_, _, H, W = image_list[0].shape
|
152 |
+
MT_size = (H // 8, W // 8)
|
153 |
+
self.cell_index = 0
|
154 |
+
with torch.no_grad():
|
155 |
+
if image_list[0].max() > 10:
|
156 |
+
image_list = [img / 255.0 for img in image_list] # [B, 1, H, W] 0-1
|
157 |
+
# I_mean = torch.cat(image_list, dim=0).mean()
|
158 |
+
# image_list = [(image - I_mean) for image in image_list]
|
159 |
+
|
160 |
+
ms_com = []
|
161 |
+
for scale in range(self.scale_num):
|
162 |
+
self.t_sin, self.t_cos = self.temporal_filter[scale].make_temporal_filter()
|
163 |
+
self.s_sin, self.s_cos = self.spatial_filter[scale].make_gabor_filters(quadrature=True)
|
164 |
+
st_component = self.infer_scale(image_list, scale)
|
165 |
+
st_component = F.interpolate(st_component, size=MT_size, mode="bilinear", align_corners=True)
|
166 |
+
ms_com.append(st_component)
|
167 |
+
image_list = [F.interpolate(img, scale_factor=self.scale_each_level, mode="bilinear") for img in image_list]
|
168 |
+
motion_energy = self.normalize(torch.cat(ms_com, dim=1))
|
169 |
+
# self.visualize_activation(motion_energy)
|
170 |
+
return motion_energy
|
171 |
+
|
172 |
+
def normalize(self, x): # TODO
|
173 |
+
sum_activation = torch.mean(x, dim=[1], keepdim=True) + torch.square(self.norm_sigma)
|
174 |
+
x = self.norm_k.abs() * x / sum_activation
|
175 |
+
return x
|
176 |
+
|
177 |
+
def _get_v1_order(self):
|
178 |
+
thetas = [gabor_scale.thetas for gabor_scale in self.spatial_filter]
|
179 |
+
fss = [gabor_scale.fs for gabor_scale in self.spatial_filter]
|
180 |
+
fts = [temporal_scale.ft for temporal_scale in self.temporal_filter]
|
181 |
+
scale_each_level = self.scale_each_level
|
182 |
+
|
183 |
+
scale_num = self.scale_num
|
184 |
+
neural_representation = []
|
185 |
+
index = 0
|
186 |
+
for scale_idx in range(len(thetas)):
|
187 |
+
theta_scale = thetas[scale_idx]
|
188 |
+
theta_scale = torch.sigmoid(theta_scale) * 2 * torch.pi # spatial orientation constrain to 0-pi
|
189 |
+
fs_scale = fss[scale_idx]
|
190 |
+
fs_scale = torch.sigmoid(fs_scale) * 0.25
|
191 |
+
fs_scale = fs_scale * (scale_each_level ** scale_idx)
|
192 |
+
|
193 |
+
ft_scale = fts[scale_idx]
|
194 |
+
ft_scale = torch.sigmoid(ft_scale) * 0.25
|
195 |
+
|
196 |
+
theta_scale = theta_scale.squeeze().cpu().detach().numpy()
|
197 |
+
fs_scale = fs_scale.squeeze().cpu().detach().numpy()
|
198 |
+
ft_scale = ft_scale.squeeze().cpu().detach().numpy()
|
199 |
+
for gabor_idx in range(len(theta_scale)):
|
200 |
+
speed = ft_scale[gabor_idx] / fs_scale[gabor_idx]
|
201 |
+
assert speed >= 0
|
202 |
+
angle = theta_scale[gabor_idx]
|
203 |
+
a = {"theta": -angle + np.pi, "fs": fs_scale[gabor_idx], "ft": ft_scale[gabor_idx], "speed": speed,
|
204 |
+
"index": index}
|
205 |
+
index = index + 1
|
206 |
+
neural_representation.append(a)
|
207 |
+
return neural_representation
|
208 |
+
|
209 |
+
def visualize_activation(self, activation, if_log=True):
|
210 |
+
neural_representation = self._get_v1_order()
|
211 |
+
activation = activation[:, :, 14:-14, 14:-14] # eliminate boundary
|
212 |
+
activation = torch.mean(activation, dim=[2, 3], keepdim=False)[0]
|
213 |
+
ax1 = plt.subplot(111, projection='polar')
|
214 |
+
theta_list = []
|
215 |
+
v_list = []
|
216 |
+
energy_list = []
|
217 |
+
for index in range(len(neural_representation)):
|
218 |
+
v = neural_representation[index]["speed"]
|
219 |
+
theta = neural_representation[index]["theta"]
|
220 |
+
location = neural_representation[index]["index"]
|
221 |
+
energy = activation.squeeze()[location].cpu().detach().numpy()
|
222 |
+
theta_list.append(theta)
|
223 |
+
v_list.append(v)
|
224 |
+
energy_list.append(energy)
|
225 |
+
v_list, theta_list, energy_list = np.array(v_list), np.array(theta_list), np.array(energy_list)
|
226 |
+
x, y = pol2cart(v_list, theta_list)
|
227 |
+
plt.scatter(theta_list, v_list, c=energy_list, cmap="rainbow", s=(energy_list + 20), alpha=0.5)
|
228 |
+
plt.axis('on')
|
229 |
+
if if_log:
|
230 |
+
ax1.set_rscale('symlog')
|
231 |
+
plt.colorbar()
|
232 |
+
energy_list = np.expand_dims(energy_list, 0).repeat(len(theta_list), 0)
|
233 |
+
|
234 |
+
buf = BytesIO()
|
235 |
+
plt.savefig(buf, format='png')
|
236 |
+
buf.seek(0)
|
237 |
+
# read the buffer and convert to an image
|
238 |
+
image = imageio.imread(buf)
|
239 |
+
buf.close()
|
240 |
+
plt.close()
|
241 |
+
plt.clf()
|
242 |
+
return image
|
243 |
+
|
244 |
+
|
245 |
+
@staticmethod
|
246 |
+
def demo():
|
247 |
+
input = [torch.ones(2, 1, 256, 256).cuda() for k in range(11)]
|
248 |
+
model = V1(spatial_num=16, scale_num=16, scale_factor=16, kernel_radius=7, num_ft=16,
|
249 |
+
kernel_size=6, average_time=True).cuda()
|
250 |
+
for i in range(100):
|
251 |
+
import time
|
252 |
+
start = time.time()
|
253 |
+
with autocast(enabled=True):
|
254 |
+
x = model(input)
|
255 |
+
print(x.shape)
|
256 |
+
torch.mean(x).backward()
|
257 |
+
end = time.time()
|
258 |
+
print(end - start)
|
259 |
+
print("#================================++#")
|
260 |
+
|
261 |
+
@property
|
262 |
+
def num_after_st(self):
|
263 |
+
return self._num_after_st
|
264 |
+
|
265 |
+
|
266 |
+
class TemporalFilter(nn.Module):
|
267 |
+
def __init__(self, in_channels=1, num_ft=8, kernel_size=6, random=True):
|
268 |
+
# 40ms per time unit, 200ms -> 5+1 frames
|
269 |
+
# use exponential decay plus sin wave
|
270 |
+
super().__init__()
|
271 |
+
self.kernel_size = kernel_size
|
272 |
+
|
273 |
+
def make_param(in_channels, values, requires_grad=True, dtype=None):
|
274 |
+
if dtype is None:
|
275 |
+
dtype = 'float32'
|
276 |
+
values = numpy.require(values, dtype=dtype)
|
277 |
+
n = in_channels * len(values)
|
278 |
+
data = torch.from_numpy(values).view(1, -1)
|
279 |
+
data = data.repeat(in_channels, 1)
|
280 |
+
return torch.nn.Parameter(data=data, requires_grad=requires_grad)
|
281 |
+
|
282 |
+
indices = torch.arange(kernel_size, dtype=torch.float32)
|
283 |
+
self.register_buffer('indices', indices)
|
284 |
+
if random:
|
285 |
+
self.ft = make_param(in_channels, values=inverse_sigmoid(numpy.random.uniform(0.01, 0.99, num_ft)),
|
286 |
+
requires_grad=True)
|
287 |
+
self.tao = make_param(in_channels, values=numpy.arange(num_ft) / 2 + 1, requires_grad=True)
|
288 |
+
else: # evenly distributed
|
289 |
+
self.ft = make_param(in_channels, values=inverse_sigmoid(numpy.linspace(0.01, 0.99, num_ft)),
|
290 |
+
requires_grad=True)
|
291 |
+
self.tao = make_param(in_channels, values=numpy.arange(num_ft) / 2 + 1, requires_grad=True)
|
292 |
+
self.feat_dim = num_ft
|
293 |
+
self.temporal_decay = 0.2
|
294 |
+
|
295 |
+
def make_temporal_filter(self):
|
296 |
+
fts = torch.sigmoid(self.ft) * 0.25
|
297 |
+
tao = torch.sigmoid(self.tao) * (-self.kernel_size / np.log(self.temporal_decay))
|
298 |
+
t = self.indices
|
299 |
+
|
300 |
+
fts = fts.view(1, fts.shape[1], 1)
|
301 |
+
tao = tao.view(1, tao.shape[1], 1)
|
302 |
+
t = t.view(1, 1, t.shape[0])
|
303 |
+
|
304 |
+
temporal_sin = torch.exp(-t / tao) * torch.sin(2 * torch.pi * fts * t)
|
305 |
+
temporal_cos = torch.exp(-t / tao) * torch.cos(2 * torch.pi * fts * t)
|
306 |
+
temporal_sin = temporal_sin.view(-1, self.kernel_size)
|
307 |
+
temporal_cos = temporal_cos.view(-1, self.kernel_size)
|
308 |
+
|
309 |
+
temporal_sin = temporal_sin.view(self.feat_dim, 1, self.kernel_size)
|
310 |
+
temporal_cos = temporal_cos.view(self.feat_dim, 1, self.kernel_size)
|
311 |
+
# temporal_sin = torch.chunk(temporal_sin, dim=0, chunks=self._feat_dim)
|
312 |
+
# temporal_cos = torch.chunk(temporal_cos, dim=0, chunks=self._feat_dim)
|
313 |
+
|
314 |
+
return temporal_sin, temporal_cos # 1,kz
|
315 |
+
|
316 |
+
def demo_temporal_filter(self, points=100):
|
317 |
+
fts = torch.sigmoid(self.ft) * 0.25
|
318 |
+
tao = torch.sigmoid(self.tao) * (-(self.kernel_size - 1) / np.log(self.temporal_decay))
|
319 |
+
t = torch.linspace(self.indices[0], self.indices[-1], steps=points)
|
320 |
+
|
321 |
+
fts = fts.view(1, fts.shape[1], 1)
|
322 |
+
tao = tao.view(1, tao.shape[1], 1)
|
323 |
+
t = t.view(1, 1, t.shape[0])
|
324 |
+
print("ft:" + str(fts))
|
325 |
+
print("tao:" + str(tao))
|
326 |
+
|
327 |
+
temporal_sin = torch.exp(-t / tao) * torch.sin(2 * torch.pi * fts * t)
|
328 |
+
temporal_cos = torch.exp(-t / tao) * torch.cos(2 * torch.pi * fts * t)
|
329 |
+
temporal_sin = temporal_sin.view(-1, points)
|
330 |
+
temporal_cos = temporal_cos.view(-1, points)
|
331 |
+
|
332 |
+
temporal_sin = temporal_sin.view(self.feat_dim, 1, points)
|
333 |
+
temporal_cos = temporal_cos.view(self.feat_dim, 1, points)
|
334 |
+
# temporal_sin = torch.chunk(temporal_sin, dim=0, chunks=self._feat_dim)
|
335 |
+
# temporal_cos = torch.chunk(temporal_cos, dim=0, chunks=self._feat_dim)
|
336 |
+
|
337 |
+
return temporal_sin, temporal_cos # 1,kz
|
338 |
+
|
339 |
+
def forward(self, x_sin, x_cos):
|
340 |
+
in_channels = x_sin.size(1)
|
341 |
+
n = x_sin.size(2)
|
342 |
+
# batch, c, sequence
|
343 |
+
me = []
|
344 |
+
t_sin, t_cos = self.make_temporal_filter()
|
345 |
+
for n_t in range(self.feat_dim):
|
346 |
+
k_sin = t_sin[n_t, ...].expand(in_channels, -1, -1)
|
347 |
+
k_cos = t_cos[n_t, ...].expand(in_channels, -1, -1)
|
348 |
+
|
349 |
+
a = F.conv1d(x_sin, weight=k_cos, padding="same", groups=in_channels, bias=None)
|
350 |
+
b = F.conv1d(x_cos, weight=k_sin, padding="same", groups=in_channels, bias=None)
|
351 |
+
g_o = a + b
|
352 |
+
|
353 |
+
a = F.conv1d(x_sin, weight=k_sin, padding="same", groups=in_channels, bias=None)
|
354 |
+
b = F.conv1d(x_cos, weight=k_cos, padding="same", groups=in_channels, bias=None)
|
355 |
+
g_e = a - b
|
356 |
+
|
357 |
+
energy_component = g_o ** 2 + g_e ** 2
|
358 |
+
me.append(energy_component)
|
359 |
+
|
360 |
+
return me
|
361 |
+
|
362 |
+
|
363 |
+
class GaborFilters(nn.Module):
|
364 |
+
def __init__(self,
|
365 |
+
in_channels=1,
|
366 |
+
kernel_radius=7,
|
367 |
+
num_units=512,
|
368 |
+
random=True
|
369 |
+
):
|
370 |
+
# the total number of or units for each scale
|
371 |
+
super().__init__()
|
372 |
+
self.in_channels = in_channels
|
373 |
+
kernel_size = kernel_radius * 2 + 1
|
374 |
+
self.kernel_size = kernel_size
|
375 |
+
self.kernel_radius = kernel_radius
|
376 |
+
|
377 |
+
def make_param(in_channels, values, requires_grad=True, dtype=None):
|
378 |
+
if dtype is None:
|
379 |
+
dtype = 'float32'
|
380 |
+
values = numpy.require(values, dtype=dtype)
|
381 |
+
n = in_channels * len(values)
|
382 |
+
data = torch.from_numpy(values).view(1, -1)
|
383 |
+
data = data.repeat(in_channels, 1)
|
384 |
+
return torch.nn.Parameter(data=data, requires_grad=requires_grad)
|
385 |
+
|
386 |
+
# build all learnable parameters
|
387 |
+
# random distribution
|
388 |
+
if random:
|
389 |
+
self.sigmas = make_param(in_channels, inverse_sigmoid(np.random.uniform(0.8, 0.99, num_units)))
|
390 |
+
self.fs = make_param(in_channels, values=inverse_sigmoid(numpy.random.uniform(0.2, 0.8, num_units)))
|
391 |
+
# maximun is 0.25 cycle/frame
|
392 |
+
self.gammas = make_param(in_channels, numpy.ones(num_units)) # TODO: fix gamma or not
|
393 |
+
self.psis = make_param(in_channels, np.zeros(num_units), requires_grad=False) # fix phase
|
394 |
+
self.thetas = make_param(in_channels, values=inverse_sigmoid(numpy.random.uniform(0.01, 0.99, num_units)),
|
395 |
+
requires_grad=True)
|
396 |
+
else: # evenly distribution
|
397 |
+
self.sigmas = make_param(in_channels, inverse_sigmoid(np.linspace(0.8, 0.99, num_units)))
|
398 |
+
self.fs = make_param(in_channels, values=inverse_sigmoid(numpy.linspace(0.01, 0.99, num_units)))
|
399 |
+
# maximun is 0.25 cycle/frame
|
400 |
+
self.gammas = make_param(in_channels, numpy.ones(num_units)) # TODO: fix gamma or not
|
401 |
+
self.psis = make_param(in_channels, np.zeros(num_units), requires_grad=False) # fix phase
|
402 |
+
self.thetas = make_param(in_channels, values=inverse_sigmoid(numpy.linspace(0, 1, num_units)),
|
403 |
+
requires_grad=True)
|
404 |
+
|
405 |
+
indices = torch.arange(kernel_size, dtype=torch.float32) - (kernel_size - 1) / 2
|
406 |
+
self.register_buffer('indices', indices)
|
407 |
+
self.spatial_decay = 0.5
|
408 |
+
# number of channels after the conv
|
409 |
+
self.n_channels_post_conv = num_units
|
410 |
+
|
411 |
+
def make_gabor_filters(self, quadrature=True):
|
412 |
+
sigmas = torch.sigmoid(self.sigmas) * np.sqrt(
|
413 |
+
(self.kernel_radius - 1) ** 2 * 0.5 / np.log(
|
414 |
+
1 / self.spatial_decay)) # std of gauss win decay to 0.2 by log(0.2)
|
415 |
+
fs = torch.sigmoid(self.fs) * 0.25
|
416 |
+
# frequency of cos and sine wave keep positive, must > 2 to avoid aliasing
|
417 |
+
gammas = torch.abs(self.gammas) # shape of gauss win, set as 1 by default
|
418 |
+
psis = self.psis # phase of cos wave
|
419 |
+
thetas = torch.sigmoid(self.thetas) * 2 * torch.pi # spatial orientation constrain to 0-2pi
|
420 |
+
y = self.indices
|
421 |
+
x = self.indices
|
422 |
+
|
423 |
+
in_channels = sigmas.shape[0]
|
424 |
+
assert in_channels == fs.shape[0]
|
425 |
+
assert in_channels == gammas.shape[0]
|
426 |
+
|
427 |
+
kernel_size = y.shape[0], x.shape[0]
|
428 |
+
|
429 |
+
sigmas = sigmas.view(in_channels, sigmas.shape[1], 1, 1)
|
430 |
+
fs = fs.view(in_channels, fs.shape[1], 1, 1)
|
431 |
+
gammas = gammas.view(in_channels, gammas.shape[1], 1, 1)
|
432 |
+
psis = psis.view(in_channels, psis.shape[1], 1, 1)
|
433 |
+
thetas = thetas.view(in_channels, thetas.shape[1], 1, 1)
|
434 |
+
y = y.view(1, 1, y.shape[0], 1)
|
435 |
+
x = x.view(1, 1, 1, x.shape[0])
|
436 |
+
|
437 |
+
sigma_x = sigmas
|
438 |
+
sigma_y = sigmas / gammas
|
439 |
+
|
440 |
+
sin_t = torch.sin(thetas)
|
441 |
+
cos_t = torch.cos(thetas)
|
442 |
+
y_theta = -x * sin_t + y * cos_t
|
443 |
+
x_theta = x * cos_t + y * sin_t
|
444 |
+
|
445 |
+
if quadrature:
|
446 |
+
gb_cos = torch.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 + y_theta ** 2 / sigma_y ** 2)) \
|
447 |
+
* torch.cos(2.0 * math.pi * x_theta * fs + psis)
|
448 |
+
gb_sin = torch.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 + y_theta ** 2 / sigma_y ** 2)) \
|
449 |
+
* torch.sin(2.0 * math.pi * x_theta * fs + psis)
|
450 |
+
gb_cos = gb_cos.reshape(-1, 1, kernel_size[0], kernel_size[1])
|
451 |
+
gb_sin = gb_sin.reshape(-1, 1, kernel_size[0], kernel_size[1])
|
452 |
+
|
453 |
+
# remove DC
|
454 |
+
gb_cos = gb_cos - torch.sum(gb_cos, dim=[-1, -2], keepdim=True) / (kernel_size[0] * kernel_size[1])
|
455 |
+
gb_sin = gb_sin - torch.sum(gb_sin, dim=[-1, -2], keepdim=True) / (kernel_size[0] * kernel_size[1])
|
456 |
+
|
457 |
+
return gb_sin, gb_cos
|
458 |
+
|
459 |
+
else:
|
460 |
+
gb = torch.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 + y_theta ** 2 / sigma_y ** 2)) \
|
461 |
+
* torch.cos(2.0 * math.pi * x_theta * fs + psis)
|
462 |
+
|
463 |
+
gb = gb.view(-1, kernel_size[0], kernel_size[1])
|
464 |
+
return gb
|
465 |
+
|
466 |
+
def forward(self, x):
|
467 |
+
batch_size = x.size(0)
|
468 |
+
sy = x.size(2)
|
469 |
+
sx = x.size(3)
|
470 |
+
gb_sin, gb_cos = self.make_gabor_filters(quadrature=True)
|
471 |
+
assert gb_sin.shape[0] == self.n_channels_post_conv
|
472 |
+
assert gb_sin.shape[2] == self.kernel_size
|
473 |
+
assert gb_sin.shape[3] == self.kernel_size
|
474 |
+
gb_sin = gb_sin.view(self.n_channels_post_conv, 1, self.kernel_size, self.kernel_size)
|
475 |
+
gb_cos = gb_cos.view(self.n_channels_post_conv, 1, self.kernel_size, self.kernel_size)
|
476 |
+
|
477 |
+
# flip ke
|
478 |
+
gb_sin = torch.flip(gb_sin, dims=[-1, -2])
|
479 |
+
gb_cos = torch.flip(gb_cos, dims=[-1, -2])
|
480 |
+
|
481 |
+
res_sin = F.conv2d(input=x, weight=gb_sin,
|
482 |
+
padding=self.kernel_radius, groups=self.in_channels)
|
483 |
+
res_cos = F.conv2d(input=x, weight=gb_cos,
|
484 |
+
padding=self.kernel_radius, groups=self.in_channels)
|
485 |
+
|
486 |
+
if self.rotation_invariant:
|
487 |
+
res_sin = res_sin.view(batch_size, self.in_channels, -1, self.n_thetas, sy, sx)
|
488 |
+
res_sin, _ = res_sin.max(dim=3)
|
489 |
+
res_cos = res_cos.view(batch_size, self.in_channels, -1, self.n_thetas, sy, sx)
|
490 |
+
res_cos, _ = res_cos.max(dim=3)
|
491 |
+
|
492 |
+
res_sin = res_sin.view(batch_size, -1, sy, sx)
|
493 |
+
res_cos = res_cos.view(batch_size, -1, sy, sx)
|
494 |
+
|
495 |
+
return res_sin, res_cos
|
496 |
+
|
497 |
+
def demo_gabor_filters(self, quadrature=True, points=100):
|
498 |
+
|
499 |
+
sigmas = torch.sigmoid(self.sigmas) * np.sqrt(
|
500 |
+
(self.kernel_radius - 1) ** 2 * 0.5 / np.log(
|
501 |
+
1 / self.spatial_decay)) # std of gauss win decay to 0.2 by log(0.2)
|
502 |
+
fs = torch.sigmoid(self.fs) * 0.25
|
503 |
+
# frequency of cos and sine wave keep positive, must > 2 to avoid aliasing
|
504 |
+
gammas = torch.abs(self.gammas) # shape of gauss win, set as 1 by default
|
505 |
+
thetas = torch.sigmoid(self.thetas) * 2 * torch.pi # spatial orientation constrain to 0-2pi
|
506 |
+
psis = self.psis # phase of cos wave
|
507 |
+
print("theta:" + str(thetas))
|
508 |
+
print("fs:" + str(fs))
|
509 |
+
|
510 |
+
x = torch.linspace(self.indices[0], self.indices[-1], points)
|
511 |
+
y = torch.linspace(self.indices[0], self.indices[-1], points)
|
512 |
+
|
513 |
+
in_channels = sigmas.shape[0]
|
514 |
+
assert in_channels == fs.shape[0]
|
515 |
+
assert in_channels == gammas.shape[0]
|
516 |
+
kernel_size = y.shape[0], x.shape[0]
|
517 |
+
|
518 |
+
sigmas = sigmas.view(in_channels, sigmas.shape[1], 1, 1)
|
519 |
+
fs = fs.view(in_channels, fs.shape[1], 1, 1)
|
520 |
+
gammas = gammas.view(in_channels, gammas.shape[1], 1, 1)
|
521 |
+
psis = psis.view(in_channels, psis.shape[1], 1, 1)
|
522 |
+
thetas = thetas.view(in_channels, thetas.shape[1], 1, 1)
|
523 |
+
y = y.view(1, 1, y.shape[0], 1)
|
524 |
+
x = x.view(1, 1, 1, x.shape[0])
|
525 |
+
|
526 |
+
sigma_x = sigmas
|
527 |
+
sigma_y = sigmas / gammas
|
528 |
+
|
529 |
+
sin_t = torch.sin(thetas)
|
530 |
+
cos_t = torch.cos(thetas)
|
531 |
+
y_theta = -x * sin_t + y * cos_t
|
532 |
+
x_theta = x * cos_t + y * sin_t
|
533 |
+
|
534 |
+
if quadrature:
|
535 |
+
gb_cos = torch.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 + y_theta ** 2 / sigma_y ** 2)) \
|
536 |
+
* torch.cos(2.0 * math.pi * x_theta * fs + psis)
|
537 |
+
gb_sin = torch.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 + y_theta ** 2 / sigma_y ** 2)) \
|
538 |
+
* torch.sin(2.0 * math.pi * x_theta * fs + psis)
|
539 |
+
gb_cos = gb_cos.reshape(-1, 1, points, points)
|
540 |
+
gb_sin = gb_sin.reshape(-1, 1, points, points)
|
541 |
+
|
542 |
+
# remove DC
|
543 |
+
gb_cos = gb_cos - torch.sum(gb_cos, dim=[-1, -2], keepdim=True) / (points * points)
|
544 |
+
gb_sin = gb_sin - torch.sum(gb_sin, dim=[-1, -2], keepdim=True) / (points * points)
|
545 |
+
|
546 |
+
return gb_sin, gb_cos
|
547 |
+
|
548 |
+
else:
|
549 |
+
gb = torch.exp(-.5 * (x_theta ** 2 / sigma_x ** 2 + y_theta ** 2 / sigma_y ** 2)) \
|
550 |
+
* torch.cos(2.0 * math.pi * x_theta * fs + psis)
|
551 |
+
|
552 |
+
gb = gb.view(-1, kernel_size[0], kernel_size[1])
|
553 |
+
return gb
|
554 |
+
|
555 |
+
|
556 |
+
def te_gabor_(num_units=48):
|
557 |
+
s_point = 100
|
558 |
+
s_kz = 7
|
559 |
+
gb_sin, gb_cos = GaborFilters(num_units=num_units, kernel_radius=s_kz).demo_gabor_filters(points=s_point)
|
560 |
+
gb = gb_sin ** 2 + gb_cos ** 2
|
561 |
+
|
562 |
+
print(gb_sin.shape)
|
563 |
+
|
564 |
+
for c in range(gb_sin.size(0)):
|
565 |
+
plt.subplot(1, 3, 1)
|
566 |
+
curve = gb_cos[c].detach().cpu().squeeze().numpy()
|
567 |
+
plt.imshow(curve)
|
568 |
+
plt.subplot(1, 3, 2)
|
569 |
+
curve = gb_sin[c].detach().cpu().squeeze().numpy()
|
570 |
+
plt.imshow(curve)
|
571 |
+
|
572 |
+
plt.subplot(1, 3, 3)
|
573 |
+
curve = gb[c].detach().cpu().squeeze().numpy()
|
574 |
+
plt.imshow(curve)
|
575 |
+
plt.show()
|
576 |
+
|
577 |
+
|
578 |
+
def te_spatial_temporal():
|
579 |
+
t_point = 6 * 100
|
580 |
+
s_point = 14 * 100
|
581 |
+
s_kz = 7
|
582 |
+
t_kz = 6
|
583 |
+
filenames = []
|
584 |
+
gb_sin_b, gb_cos_b = GaborFilters(num_units=48, kernel_radius=s_kz).demo_gabor_filters(points=s_point)
|
585 |
+
temporal = TemporalFilter(num_ft=2, kernel_size=t_kz)
|
586 |
+
t_sin, t_cos = temporal.demo_temporal_filter(points=t_point)
|
587 |
+
x = np.linspace(0, t_kz, t_point)
|
588 |
+
index = 0
|
589 |
+
for i in range(gb_sin_b.size(0)):
|
590 |
+
for j in range(t_sin.size(0)):
|
591 |
+
plt.figure(figsize=(14, 9), dpi=80)
|
592 |
+
plt.subplot(2, 3, 1)
|
593 |
+
curve = gb_sin_b[i].squeeze().detach().numpy()
|
594 |
+
plt.imshow(curve)
|
595 |
+
plt.title("Gabor Sin")
|
596 |
+
plt.subplot(2, 3, 2)
|
597 |
+
curve = gb_cos_b[i].squeeze().detach().numpy()
|
598 |
+
plt.imshow(curve)
|
599 |
+
plt.title("Gabor Cos")
|
600 |
+
|
601 |
+
plt.subplot(2, 3, 3)
|
602 |
+
curve = t_sin[j].squeeze().detach().numpy()
|
603 |
+
plt.plot(x, curve, label='sin')
|
604 |
+
plt.title("Temporal Sin")
|
605 |
+
|
606 |
+
curve = t_cos[j].squeeze().detach().numpy()
|
607 |
+
plt.plot(x, curve, label='cos')
|
608 |
+
plt.xlabel('Time (s)')
|
609 |
+
plt.ylabel('Response to pulse at t=0')
|
610 |
+
plt.legend()
|
611 |
+
plt.title("Temporal filter")
|
612 |
+
|
613 |
+
gb_sin = gb_sin_b[i].squeeze().detach()[5, :]
|
614 |
+
gb_cos = gb_cos_b[i].squeeze().detach()[5, :]
|
615 |
+
|
616 |
+
a = np.outer(t_cos[j].detach(), gb_sin)
|
617 |
+
b = np.outer(t_sin[j].detach(), gb_cos)
|
618 |
+
g_o = a + b
|
619 |
+
|
620 |
+
a = np.outer(t_sin[j].detach(), gb_sin)
|
621 |
+
b = np.outer(t_cos[j].detach(), gb_cos)
|
622 |
+
g_e = a - b
|
623 |
+
energy_component = g_o ** 2 + g_e ** 2
|
624 |
+
|
625 |
+
plt.subplot(2, 3, 4)
|
626 |
+
curve = g_o
|
627 |
+
plt.imshow(curve, cmap="gray")
|
628 |
+
plt.title("Spatial Temporal even")
|
629 |
+
plt.subplot(2, 3, 5)
|
630 |
+
curve = g_e
|
631 |
+
plt.imshow(curve, cmap="gray")
|
632 |
+
plt.title("Spatial Temporal odd")
|
633 |
+
|
634 |
+
plt.subplot(2, 3, 6)
|
635 |
+
curve = energy_component
|
636 |
+
plt.imshow(curve, cmap="gray")
|
637 |
+
plt.title("energy")
|
638 |
+
plt.savefig('filter_%d.png' % (index))
|
639 |
+
filenames.append('filter_%d.png' % (index))
|
640 |
+
index += 1
|
641 |
+
plt.show()
|
642 |
+
# build gif
|
643 |
+
with imageio.get_writer('filters_orientation.gif', mode='I') as writer:
|
644 |
+
for filename in filenames:
|
645 |
+
image = imageio.imread(filename)
|
646 |
+
writer.append_data(image)
|
647 |
+
|
648 |
+
# Remove files
|
649 |
+
for filename in set(filenames):
|
650 |
+
os.remove(filename)
|
651 |
+
|
652 |
+
|
653 |
+
def te_temporal_():
|
654 |
+
k_size = 6
|
655 |
+
temporal = TemporalFilter(n_tao=2, num_ft=8, kernel_size=k_size)
|
656 |
+
sin, cos = temporal.demo_temporal_filter()
|
657 |
+
print(sin.shape)
|
658 |
+
x = np.linspace(0, k_size, k_size * 100)
|
659 |
+
|
660 |
+
# plot temporal filters to illustrate what they look like.
|
661 |
+
for c in range(sin.size(0)):
|
662 |
+
curve = cos[c].detach().cpu().squeeze().numpy()
|
663 |
+
plt.plot(x, curve, label='cos')
|
664 |
+
curve = sin[c].detach().cpu().squeeze().numpy()
|
665 |
+
plt.plot(x, curve, label='sin')
|
666 |
+
|
667 |
+
plt.xlabel('Time (s)')
|
668 |
+
plt.ylabel('Response to pulse at t=0')
|
669 |
+
plt.legend()
|
670 |
+
plt.show()
|
671 |
+
|
672 |
+
|
673 |
+
def circular_hist(ax, x, bins=16, density=True, offset=0, gaps=True):
|
674 |
+
"""
|
675 |
+
Produce a circular histogram of angles on ax.
|
676 |
+
|
677 |
+
Parameters
|
678 |
+
----------
|
679 |
+
ax : matplotlib.axes._subplots.PolarAxesSubplot
|
680 |
+
axis instance created with subplot_kw=dict(projection='polar').
|
681 |
+
|
682 |
+
x : array
|
683 |
+
Angles to plot, expected in units of radians.
|
684 |
+
|
685 |
+
bins : int, optional
|
686 |
+
Defines the number of equal-width bins in the range. The default is 16.
|
687 |
+
|
688 |
+
density : bool, optional
|
689 |
+
If True plot frequency proportional to area. If False plot frequency
|
690 |
+
proportional to radius. The default is True.
|
691 |
+
|
692 |
+
offset : float, optional
|
693 |
+
Sets the offset for the location of the 0 direction in units of
|
694 |
+
radians. The default is 0.
|
695 |
+
|
696 |
+
gaps : bool, optional
|
697 |
+
Whether to allow gaps between bins. When gaps = False the bins are
|
698 |
+
forced to partition the entire [-pi, pi] range. The default is True.
|
699 |
+
|
700 |
+
Returns
|
701 |
+
-------
|
702 |
+
n : array or list of arrays
|
703 |
+
The number of values in each bin.
|
704 |
+
|
705 |
+
bins : array
|
706 |
+
The edges of the bins.
|
707 |
+
|
708 |
+
patches : `.BarContainer` or list of a single `.Polygon`
|
709 |
+
Container of individual artists used to create the histogram
|
710 |
+
or list of such containers if there are multiple input datasets.
|
711 |
+
"""
|
712 |
+
# Wrap angles to [-pi, pi)
|
713 |
+
x = (x + np.pi) % (2 * np.pi) - np.pi
|
714 |
+
|
715 |
+
# Force bins to partition entire circle
|
716 |
+
if not gaps:
|
717 |
+
bins = np.linspace(-np.pi, np.pi, num=bins + 1)
|
718 |
+
|
719 |
+
# Bin data and record counts
|
720 |
+
n, bins = np.histogram(x, bins=bins)
|
721 |
+
|
722 |
+
# Compute width of each bin
|
723 |
+
widths = np.diff(bins)
|
724 |
+
|
725 |
+
# By default plot frequency proportional to area
|
726 |
+
if density:
|
727 |
+
# Area to assign each bin
|
728 |
+
area = n / x.size
|
729 |
+
# Calculate corresponding bin radius
|
730 |
+
radius = (area / np.pi) ** .5
|
731 |
+
# Otherwise plot frequency proportional to radius
|
732 |
+
else:
|
733 |
+
radius = n
|
734 |
+
|
735 |
+
# Plot data on ax
|
736 |
+
patches = ax.bar(bins[:-1], radius, zorder=1, align='edge', width=widths,
|
737 |
+
edgecolor='C0', fill=False, linewidth=1)
|
738 |
+
|
739 |
+
# Set the direction of the zero angle
|
740 |
+
ax.set_theta_offset(offset)
|
741 |
+
|
742 |
+
# Remove ylabels for area plots (they are mostly obstructive)
|
743 |
+
if density:
|
744 |
+
ax.set_yticks([])
|
745 |
+
|
746 |
+
return n, bins, patches
|
747 |
+
|
748 |
+
|
749 |
+
def show_trained_model(file_name="/home/2TSSD/experiment/FFMEDNN/Sintel_fixv1_10.62_ckpt.pth.tar"):
|
750 |
+
import utils.torch_utils as utils
|
751 |
+
from model.fle_version_2_3.FFV1MT_MS import FFV1DNN
|
752 |
+
model = FFV1DNN(num_scales=8,
|
753 |
+
num_cells=256,
|
754 |
+
upsample_factor=8,
|
755 |
+
feature_channels=256,
|
756 |
+
scale_factor=16,
|
757 |
+
num_layers=6)
|
758 |
+
# model = utils.restore_model(model, file_name)
|
759 |
+
model = model.ffv1
|
760 |
+
t_point = 100
|
761 |
+
s_point = 100
|
762 |
+
t_kz = 6
|
763 |
+
filenames = []
|
764 |
+
x = np.arange(0, 6) * 40
|
765 |
+
x = np.repeat(x[None], axis=0, repeats=256)
|
766 |
+
temporal = model.temporal_pooling.data.cpu().squeeze().numpy()
|
767 |
+
mean = np.mean(temporal, axis=0)
|
768 |
+
plt.figure(figsize=(10, 10))
|
769 |
+
plt.subplot(2, 1, 1)
|
770 |
+
for idx in range(0, 256):
|
771 |
+
plt.plot(x[idx], temporal[idx])
|
772 |
+
plt.subplot(2, 1, 2)
|
773 |
+
plt.plot(x[0], mean, label="mean")
|
774 |
+
|
775 |
+
plt.xlabel("times (ms)")
|
776 |
+
plt.ylabel("temporal pooling weight")
|
777 |
+
plt.legend()
|
778 |
+
plt.grid(True)
|
779 |
+
plt.show()
|
780 |
+
neural_representation = model._get_v1_order()
|
781 |
+
|
782 |
+
fs = np.array([ne["fs"] for ne in neural_representation])
|
783 |
+
ft = np.array([ne["ft"] for ne in neural_representation])
|
784 |
+
|
785 |
+
ax1 = plt.subplot(131, projection='polar')
|
786 |
+
theta_list = []
|
787 |
+
v_list = []
|
788 |
+
energy_list = []
|
789 |
+
for index in range(len(neural_representation)):
|
790 |
+
v = neural_representation[index]["speed"]
|
791 |
+
theta = neural_representation[index]["theta"]
|
792 |
+
theta_list.append(theta)
|
793 |
+
v_list.append(v)
|
794 |
+
|
795 |
+
v_list, theta_list = np.array(v_list), np.array(theta_list)
|
796 |
+
x, y = pol2cart(v_list, theta_list)
|
797 |
+
plt.scatter(theta_list, v_list, c=v_list, cmap="rainbow", s=(v_list + 20), alpha=0.8)
|
798 |
+
plt.axis('on')
|
799 |
+
# plt.colorbar()
|
800 |
+
plt.grid(True)
|
801 |
+
# plt.subplot(132, projection="polar")
|
802 |
+
# plt.scatter(theta_list, np.ones_like(theta_list))
|
803 |
+
plt.subplot(132, projection='polar')
|
804 |
+
plt.scatter(theta_list, np.ones_like(v_list))
|
805 |
+
lst = []
|
806 |
+
for scale in range(8):
|
807 |
+
lst += ["scale %d" % scale] * 32
|
808 |
+
data = {"Spatial Frequency": fs, 'Temporal Frequency': ft, "Class": lst}
|
809 |
+
df = pd.DataFrame(data=data)
|
810 |
+
ax = plt.subplot(133, projection='polar')
|
811 |
+
# theta_list = theta_list[v_list > (ft * v_list.mean())]
|
812 |
+
print(len(theta_list))
|
813 |
+
bins_number = 8 # the [0, 360) interval will be subdivided into this
|
814 |
+
# number of equal bins
|
815 |
+
zone = np.pi / 8
|
816 |
+
theta_list[theta_list < (-np.pi + zone)] = theta_list[theta_list < (-np.pi + zone)] + np.pi * 2
|
817 |
+
bins = np.linspace(-np.pi + zone, np.pi + zone, bins_number + 1)
|
818 |
+
n, _, _ = plt.hist(theta_list, bins, edgecolor="black")
|
819 |
+
# ax.set_theta_offset(-np.pi / 8 - np.pi)
|
820 |
+
ax.set_yticklabels([])
|
821 |
+
plt.grid(True)
|
822 |
+
import seaborn as sns
|
823 |
+
sns.jointplot(data=df, x="Spatial Frequency", y="Temporal Frequency", hue="Class", xlim=[0, 0.3], ylim=[0, 0.3])
|
824 |
+
plt.grid(True)
|
825 |
+
g = sns.jointplot(data=df, x="Spatial Frequency", y="Temporal Frequency", xlim=[0, 0.25], ylim=[0, 0.25])
|
826 |
+
# g.plot_joint(sns.kdeplot, color="r", zorder=0, levels=6)
|
827 |
+
|
828 |
+
plt.grid(True)
|
829 |
+
plt.show()
|
830 |
+
|
831 |
+
# show spatial frequency preference and temporal frequency preference.
|
832 |
+
|
833 |
+
x = np.linspace(0, t_kz, t_point)
|
834 |
+
index = 0
|
835 |
+
for scale in range(len(model.spatial_filter)):
|
836 |
+
t_sin, t_cos = model.temporal_filter[scale].demo_temporal_filter(points=t_point)
|
837 |
+
gb_sin_b, gb_cos_b = model.spatial_filter[scale].demo_gabor_filters(points=s_point)
|
838 |
+
for i in range(gb_sin_b.size(0)):
|
839 |
+
plt.figure(figsize=(14, 9), dpi=80)
|
840 |
+
plt.subplot(2, 3, 1)
|
841 |
+
curve = gb_sin_b[i].squeeze().detach().numpy()
|
842 |
+
plt.imshow(curve)
|
843 |
+
plt.title("Gabor Sin")
|
844 |
+
plt.subplot(2, 3, 2)
|
845 |
+
curve = gb_cos_b[i].squeeze().detach().numpy()
|
846 |
+
plt.imshow(curve)
|
847 |
+
plt.title("Gabor Cos")
|
848 |
+
|
849 |
+
plt.subplot(2, 3, 3)
|
850 |
+
curve = t_sin[i].squeeze().detach().numpy()
|
851 |
+
plt.plot(x, curve, label='sin')
|
852 |
+
plt.title("Temporal Sin")
|
853 |
+
|
854 |
+
curve = t_cos[i].squeeze().detach().numpy()
|
855 |
+
plt.plot(x, curve, label='cos')
|
856 |
+
plt.xlabel('Time (s)')
|
857 |
+
plt.ylabel('Response to pulse at t=0')
|
858 |
+
plt.legend()
|
859 |
+
plt.title("Temporal filter")
|
860 |
+
|
861 |
+
gb_sin = gb_sin_b[i].squeeze().detach()[5, :]
|
862 |
+
gb_cos = gb_cos_b[i].squeeze().detach()[5, :]
|
863 |
+
|
864 |
+
a = np.outer(t_cos[i].detach(), gb_sin)
|
865 |
+
b = np.outer(t_sin[i].detach(), gb_cos)
|
866 |
+
g_o = a + b
|
867 |
+
|
868 |
+
a = np.outer(t_sin[i].detach(), gb_sin)
|
869 |
+
b = np.outer(t_cos[i].detach(), gb_cos)
|
870 |
+
g_e = a - b
|
871 |
+
energy_component = g_o ** 2 + g_e ** 2
|
872 |
+
|
873 |
+
plt.subplot(2, 3, 4)
|
874 |
+
curve = g_o
|
875 |
+
plt.imshow(curve, cmap="gray")
|
876 |
+
plt.title("Spatial Temporal even")
|
877 |
+
plt.subplot(2, 3, 5)
|
878 |
+
curve = g_e
|
879 |
+
plt.imshow(curve, cmap="gray")
|
880 |
+
plt.title("Spatial Temporal odd")
|
881 |
+
|
882 |
+
plt.subplot(2, 3, 6)
|
883 |
+
curve = energy_component
|
884 |
+
plt.imshow(curve, cmap="gray")
|
885 |
+
plt.title("energy")
|
886 |
+
plt.savefig('filter_%d.png' % (index))
|
887 |
+
filenames.append('filter_%d.png' % (index))
|
888 |
+
index += 1
|
889 |
+
# plt.show()
|
890 |
+
|
891 |
+
# build gif
|
892 |
+
with imageio.get_writer('filters_orientation.gif', mode='I') as writer:
|
893 |
+
for filename in filenames:
|
894 |
+
image = imageio.imread(filename)
|
895 |
+
writer.append_data(image)
|
896 |
+
|
897 |
+
# Remove files
|
898 |
+
for filename in set(filenames):
|
899 |
+
os.remove(filename)
|
900 |
+
|
901 |
+
|
902 |
+
if __name__ == "__main__":
|
903 |
+
show_trained_model()
|
904 |
+
# V1.demo()
|
905 |
+
# draw_polar()
|
906 |
+
# # V1.demo()
|
907 |
+
# # draw_polar()
|
908 |
+
show_trained_model()
|
909 |
+
# te_spatial_temporal()
|