|
0
|
1 # MCTF following Ohm04a
|
|
|
2
|
|
|
3 import Image
|
|
|
4 import ImageDraw
|
|
|
5 import pywt
|
|
|
6 import numpy
|
|
|
7 import math
|
|
|
8 import sys
|
|
|
9 import time
|
|
|
10 import _me
|
|
|
11
|
|
|
12 # type of motion vectors
|
|
|
13 UNCONNECTED = -(sys.maxint)
|
|
|
14 CONNECTED = -(sys.maxint-1)
|
|
|
15 MULTIPLE_CONNECTED = -(sys.maxint-2)
|
|
|
16
|
|
|
17 # temporal low-pass frame position
|
|
|
18 LEFT = -1
|
|
|
19 MIDDLE = 0
|
|
|
20 RIGHT = 1
|
|
|
21
|
|
|
22 def me(a, refblock, rc, cc, sr):
|
|
|
23 min_sad = sys.maxint
|
|
|
24 min_r, min_c = 0, 0
|
|
|
25 bs = refblock.shape[0]
|
|
|
26 for rs in xrange(max(0,rc-sr),min(a.shape[0]-bs,rc+sr)+1):
|
|
|
27 for cs in xrange(max(0,cc-sr),min(cc+sr,a.shape[1]-bs)+1):
|
|
|
28 sad = numpy.sum(numpy.abs(refblock - a[rs:rs+bs, cs:cs+bs]))
|
|
|
29 if sad < min_sad:
|
|
|
30 # found new local block SAD minimum, store motion vector
|
|
|
31 min_r, min_c, min_sad = rs, cs, sad
|
|
|
32 return min_r, min_c, min_sad
|
|
|
33
|
|
|
34 def motion_estimation(a, b, blocksize=8, searchrange=8, hlevel=2):
|
|
|
35 '''
|
|
|
36 Hierarchical motion estimation from frame a to frame b.
|
|
|
37 Parameters are blocksize, searchrange and search hierarchy level.
|
|
|
38 Precision is full pixel only.
|
|
|
39 Returns the sum-of-absolute-differences (SAD) and the motion
|
|
|
40 vector field (MVF).
|
|
|
41 '''
|
|
|
42
|
|
|
43 mvf = numpy.zeros((b.shape[0], b.shape[1], 3), numpy.int)
|
|
|
44 mvf[:,:,2] = UNCONNECTED
|
|
|
45
|
|
|
46 ref = numpy.asarray(b, numpy.float)
|
|
|
47
|
|
|
48 # downsample frame data using Haar wavelet
|
|
|
49 w = pywt.Wavelet('haar')
|
|
|
50 ha = pywt.wavedec2(a, w, level=hlevel)
|
|
|
51 href = pywt.wavedec2(ref, w, level=hlevel)
|
|
|
52
|
|
|
53 # grows by 2 for every level
|
|
|
54 hbs = blocksize//2**hlevel
|
|
|
55 hsr = searchrange//2**hlevel
|
|
|
56
|
|
|
57 while True:
|
|
|
58 total_sad = 0.0
|
|
|
59 _2hlevel = 2**hlevel
|
|
|
60 for r in xrange(0, href[0].shape[0], hbs):
|
|
|
61 for c in xrange(0, href[0].shape[1], hbs):
|
|
|
62 rm = r * _2hlevel
|
|
|
63 cm = c * _2hlevel
|
|
|
64
|
|
|
65 # set center of new search of previously found vector at higher level
|
|
|
66 if mvf[rm,cm,2] >= 0: rc, cc = mvf[rm,cm,0]*2 + r, mvf[rm,cm,1]*2 + c
|
|
|
67 else: rc, cc = r, c
|
|
|
68 rs, cs, sad = _me.me(ha[0], href[0][r:r+hbs,c:c+hbs], rc, cc, hsr)
|
|
|
69 mvf[rm:rm+blocksize,cm:cm+blocksize,:] = rs - r, cs - c, int(sad)
|
|
|
70 total_sad += sad
|
|
|
71
|
|
|
72 if hlevel == 0: break
|
|
|
73
|
|
|
74 # upsample frame data using Haar wavelet
|
|
|
75 ha = [pywt.waverec2(ha[:2], w)] + ha[2:]
|
|
|
76 href = [pywt.waverec2(href[:2], w)] + href[2:]
|
|
|
77 hbs *= 2
|
|
|
78 hlevel -= 1
|
|
|
79
|
|
|
80 return total_sad, mvf
|
|
|
81
|
|
|
82 def ft_mvf(a, b, mvf, imvf, bs=8):
|
|
|
83 '''
|
|
|
84 Motion-compensated temporal decomposition between frame a and b
|
|
|
85 using Haar wavelet applying a given forward and inverse motion field.
|
|
|
86 '''
|
|
|
87
|
|
|
88 H = numpy.empty(a.shape, numpy.float)
|
|
|
89 L = numpy.empty(a.shape, numpy.float)
|
|
|
90
|
|
|
91 i0 = numpy.indices((bs,bs))[0]
|
|
|
92 i1 = numpy.indices((bs,bs))[1]
|
|
|
93
|
|
|
94 for r in xrange(0, a.shape[0], bs):
|
|
|
95 for c in xrange(0, a.shape[1], bs):
|
|
|
96 rm = mvf[r, c, 0] + r
|
|
|
97 cm = mvf[r, c, 1] + c
|
|
|
98 H[r:r+bs, c:c+bs] = numpy.asarray(a[r:r+bs,c:c+bs], numpy.float) - b[rm:rm+bs,cm:cm+bs]
|
|
|
99 rm = r + imvf[r:r+bs,c:c+bs,0] + i0
|
|
|
100 cm = c + imvf[r:r+bs,c:c+bs,1] + i1
|
|
|
101 _a = a[rm, cm]
|
|
|
102 L[r:r+bs, c:c+bs] = numpy.where( \
|
|
|
103 imvf[r:r+bs, c:c+bs, 2] == UNCONNECTED, \
|
|
|
104 numpy.asarray(b[r:r+bs, c:c+bs], numpy.float), \
|
|
|
105 0.5 * (numpy.asarray(b[r:r+bs, c:c+bs], numpy.float) + _a))
|
|
|
106
|
|
|
107 return L, H
|
|
|
108
|
|
|
109 def it_mvf(L, H, mvf, imvf, bs=8):
|
|
|
110 '''
|
|
|
111 Reconstruction of two frames a and b from temporal low- and high-pass subband
|
|
|
112 using Haar wavelet and applying the given forward and inverse motion field.
|
|
|
113 '''
|
|
|
114
|
|
|
115 i0 = numpy.indices((bs,bs))[0]
|
|
|
116 i1 = numpy.indices((bs,bs))[1]
|
|
|
117
|
|
|
118 b = numpy.empty(L.shape, numpy.float)
|
|
|
119 for r in xrange(0, L.shape[0], bs):
|
|
|
120 for c in xrange(0, L.shape[1], bs):
|
|
|
121 _L = L[r:r+bs,c:c+bs]
|
|
|
122 rm = r + imvf[r:r+bs,c:c+bs,0] + i0
|
|
|
123 cm = c + imvf[r:r+bs,c:c+bs,1] + i1
|
|
|
124 _H = H[rm, cm]
|
|
|
125 b[r:r+bs,c:c+bs] = numpy.where( \
|
|
|
126 imvf[r:r+bs,c:c+bs,2] == UNCONNECTED, \
|
|
|
127 _L, \
|
|
|
128 _L - 0.5 * _H)
|
|
|
129
|
|
|
130 a = numpy.empty(L.shape, numpy.float)
|
|
|
131 for r in xrange(0, L.shape[0], bs):
|
|
|
132 for c in xrange(0, L.shape[1], bs):
|
|
|
133 rm = mvf[r, c, 0] + r
|
|
|
134 cm = mvf[r, c, 1] + c
|
|
|
135 _H = H[r:r+bs,c:c+bs]
|
|
|
136 a[r:r+bs, c:c+bs] = numpy.where( \
|
|
|
137 mvf[r:r+bs,c:c+bs,2] == MULTIPLE_CONNECTED, \
|
|
|
138 b[rm:rm+bs,cm:cm+bs] + _H, \
|
|
|
139 L[rm:rm+bs,cm:cm+bs] + 0.5 * _H)
|
|
|
140
|
|
|
141 return a, b
|
|
|
142
|
|
|
143 def _show_mv_dist(mvf, idx=0, level=0, sr=8, fname='mv_dist'):
|
|
|
144 im = Image.new('RGB', (mvf.shape[1], mvf.shape[0]))
|
|
|
145 d = ImageDraw.Draw(im)
|
|
|
146
|
|
|
147 for r in xrange(mvf.shape[0]):
|
|
|
148 for c in xrange(mvf.shape[1]):
|
|
|
149 mv = mvf[r][c]
|
|
|
150
|
|
|
151 if sr > 0: w = int(math.sqrt(mv[0]**2 + mv[1]**2)*255/(sr*math.sqrt(2.0)))
|
|
|
152 else: w = 0
|
|
|
153
|
|
|
154 if mv[2] >= 0 or mv[2] == CONNECTED: color = (0, w, 0)
|
|
|
155 elif mv[2] == UNCONNECTED: color = (255, 0, 0)
|
|
|
156 elif mv[2] == MULTIPLE_CONNECTED: color = (0, 0, w)
|
|
|
157
|
|
|
158 d.point((c, r), fill=color)
|
|
|
159
|
|
|
160 del d
|
|
|
161 im.save('%s-%02d-%04d.png' % (fname, level, idx), 'PNG')
|
|
|
162 del im
|
|
|
163
|
|
|
164 def show_mvf(mvf, imvf, idx=0, level=0, bs=8, sr=8):
|
|
|
165 '''
|
|
|
166 Visualize the motion field as .png and output motion vectors to .txt.
|
|
|
167 '''
|
|
|
168
|
|
|
169 im = Image.new('RGB', (mvf.shape[1]*2, mvf.shape[0]*2))
|
|
|
170 d = ImageDraw.Draw(im)
|
|
|
171 f = open('mv-%02d-%04d.txt' % (level, idx), 'wt')
|
|
|
172 sad = mvf[:,:,2].ravel()
|
|
|
173 sad_min = numpy.min(numpy.where(sad < 0.0, 0, sad))
|
|
|
174 sad_max = numpy.max(sad)
|
|
|
175 for r in xrange(0,mvf.shape[0],bs):
|
|
|
176 for c in xrange(0,mvf.shape[1],bs):
|
|
|
177 mv = mvf[r][c]
|
|
|
178 print >>f, '(%d %d)' % (mv[1], mv[0]),
|
|
|
179
|
|
|
180 # fill block according to SAD
|
|
|
181 if sad_max > 0 and mv[2] > 0:
|
|
|
182 d.rectangle([(c*2,r*2),(c*2+bs*2,r*2+bs*2)], fill=((mv[2]-sad_min)*255/sad_max,0,0))
|
|
|
183
|
|
|
184 # draw motion vector
|
|
|
185 if sr > 0: w = int(math.sqrt(mv[0]**2 + mv[1]**2)/(sr*math.sqrt(2.0)))
|
|
|
186 else: w = 0
|
|
|
187
|
|
|
188 d.line([ \
|
|
|
189 (c*2+bs, r*2+bs), \
|
|
|
190 (c*2+bs+mv[1]*2, r*2+bs+mv[0]*2)], \
|
|
|
191 fill=(0,int(32+(255-32)*w),0))
|
|
|
192 d.point((c*2+bs, r*2+bs), fill=(255,255,255))
|
|
|
193
|
|
|
194 print >>f
|
|
|
195 print >>f
|
|
|
196
|
|
|
197 f.close()
|
|
|
198 del d
|
|
|
199
|
|
|
200 im.save('mv-%02d-%04d.png' % (level, idx), 'PNG')
|
|
|
201 del im
|
|
|
202
|
|
|
203 _show_mv_dist(mvf, idx, level, sr, 'mvf_dist')
|
|
|
204 _show_mv_dist(imvf, idx, level, sr, 'mvi_dist')
|
|
|
205
|
|
|
206
|
|
|
207 def inverse_mvf(mvf, bs=8):
|
|
|
208 '''
|
|
|
209 Compute the inverse of the motion field.
|
|
|
210 '''
|
|
|
211
|
|
|
212 imvf = numpy.zeros((mvf.shape[0], mvf.shape[1], 3), numpy.int)
|
|
|
213 imvf[:,:,2] = UNCONNECTED
|
|
|
214 for r in xrange(0, mvf.shape[0], bs):
|
|
|
215 for c in xrange(0, mvf.shape[1], bs):
|
|
|
216 rm = mvf[r,c,0] + r
|
|
|
217 cm = mvf[r,c,1] + c
|
|
|
218
|
|
|
219 blockmvf = mvf[r:r+bs,c:c+bs]
|
|
|
220 blockimvf = imvf[rm:rm+bs,cm:cm+bs]
|
|
|
221
|
|
|
222 # mark multiple connected in forward motion field if pixel already connected
|
|
|
223 numpy.place(blockmvf[:,:,2], blockimvf[:,:,2] > UNCONNECTED, MULTIPLE_CONNECTED)
|
|
|
224
|
|
|
225 # invert motion vector and store in inverse motion field, mark pixel as connected
|
|
|
226 unconnected = blockimvf[:,:,2] == UNCONNECTED
|
|
|
227 numpy.place(blockimvf[:,:,0], unconnected, -mvf[r,c,0])
|
|
|
228 numpy.place(blockimvf[:,:,1], unconnected, -mvf[r,c,1])
|
|
|
229 numpy.place(blockimvf[:,:,2], unconnected, CONNECTED)
|
|
|
230
|
|
|
231 return mvf, imvf
|
|
|
232
|
|
|
233 def decompose_sequence(seq, Hs=[], MVFs=[], bs=8, sr=8, hlevel=2, tlp=MIDDLE, visualize_mvf=False, dlevel=-1):
|
|
|
234 '''
|
|
|
235 Recursively decompose frame sequence using motion-compensated temporal filtering
|
|
|
236 employing the parameters blocksize, searchrange and hierarchy level for motion estimation.
|
|
|
237
|
|
|
238 Output is [L], [H0, H1, H1, H2, H2, H2, H2], [MVF0, MVF1, MVF1, MVF2, MVF2, MVF2, MVF2] for
|
|
|
239 a sequence of length 8.
|
|
|
240
|
|
|
241 The tlp argument allows to move the temporal low-pass frame to the left,
|
|
|
242 middle or right.
|
|
|
243 '''
|
|
|
244 Ls = []
|
|
|
245 if dlevel < 0: dlevel = int(math.log(len(seq), 2))
|
|
|
246
|
|
|
247 if len(seq) == 1:
|
|
|
248 return seq, Hs, MVFs
|
|
|
249
|
|
|
250 if tlp == RIGHT: left = 0; mid = len(seq); right = 0
|
|
|
251 elif tlp == LEFT: left = 0; mid = 0; right = len(seq)
|
|
|
252 else: left = 0; mid = max(len(seq)/2, 2); right = len(seq)
|
|
|
253
|
|
|
254 for i in xrange(left, mid, 2):
|
|
|
255 sad, mvf = motion_estimation(seq[i+1], seq[i], bs, sr, hlevel)
|
|
|
256 mvf, imvf = inverse_mvf(mvf, bs)
|
|
|
257 if visualize_mvf:
|
|
|
258 show_mvf(mvf, imvf, i, dlevel-1, bs, sr)
|
|
|
259 MVFs.insert(i//2, mvf)
|
|
|
260 L, H = ft_mvf(seq[i], seq[i+1], mvf, imvf, bs)
|
|
|
261 Ls.append(L)
|
|
|
262 Hs.insert(i//2, H)
|
|
|
263
|
|
|
264 for i in xrange(mid, right, 2):
|
|
|
265 sad, mvf = motion_estimation(seq[i], seq[i+1], bs, sr, hlevel)
|
|
|
266 mvf, imvf = inverse_mvf(mvf, bs)
|
|
|
267 if visualize_mvf:
|
|
|
268 show_mvf(mvf, imvf, i, dlevel-1, bs, sr)
|
|
|
269 MVFs.insert(i//2, mvf)
|
|
|
270 L, H = ft_mvf(seq[i+1], seq[i], mvf, imvf, bs)
|
|
|
271 Ls.append(L)
|
|
|
272 Hs.insert(i//2, H)
|
|
|
273
|
|
|
274 return decompose_sequence(Ls, Hs, MVFs, bs, sr, hlevel, tlp, visualize_mvf, dlevel-1)
|
|
|
275
|
|
|
276 def reconstruct_sequence(seq, Hs, MVFs, bs=8, tlp=MIDDLE):
|
|
|
277 '''
|
|
|
278 Recursively reconstruct a frame sequence from temporal low- and high-pass subbands
|
|
|
279 and motion fields.
|
|
|
280 '''
|
|
|
281
|
|
|
282 Ls = []
|
|
|
283
|
|
|
284 if len(Hs) == 0:
|
|
|
285 return seq
|
|
|
286
|
|
|
287 if tlp == RIGHT: left = 0; mid = len(seq); right = 0
|
|
|
288 elif tlp == LEFT: left = 0; mid = 0; right = len(seq)
|
|
|
289 else: left = 0; mid = max(len(seq)/2, 1); right = len(seq)
|
|
|
290
|
|
|
291 for i in xrange(0, mid):
|
|
|
292 mvf = MVFs[0]
|
|
|
293 mvf, imvf = inverse_mvf(mvf, bs)
|
|
|
294 a, b = it_mvf(seq[i], Hs[0], mvf, imvf, bs)
|
|
|
295 Ls += [a] + [b]
|
|
|
296 del Hs[0]
|
|
|
297 del MVFs[0]
|
|
|
298
|
|
|
299 for i in xrange(mid, right):
|
|
|
300 mvf = MVFs[0]
|
|
|
301 mvf, imvf = inverse_mvf(mvf, bs)
|
|
|
302 a, b = it_mvf(seq[i], Hs[0], mvf, imvf, bs)
|
|
|
303 Ls += [b] + [a]
|
|
|
304 del Hs[0]
|
|
|
305 del MVFs[0]
|
|
|
306
|
|
|
307 return reconstruct_sequence(Ls, Hs, MVFs, bs, tlp)
|
|
|
308
|
