Mercurial > hg > pymctf

comparison pymctf.py @ 0:4214d9245f8e

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
import
author Peter Meerwald <pmeerw@cosy.sbg.ac.at>
date Thu, 06 Sep 2007 13:45:48 +0200
parents
children b67c5ec1a9f0
comparison
equal deleted inserted replaced
-1:000000000000 0:4214d9245f8e
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

Repositories maintained by Peter Meerwald, pmeerw@pmeerw.net.