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