Mercurial > hg > pymctf
annotate pymctf.py @ 2:f22cbbbb6814
enable psyco, add README
| author | Peter Meerwald <pmeerw@cosy.sbg.ac.at> |
|---|---|
| date | Thu, 06 Sep 2007 14:18:51 +0200 |
| parents | b67c5ec1a9f0 |
| children | 0e5a584fd6b2 |
| 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 | |
| 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'): | |
|
1
b67c5ec1a9f0
import only needed modules
Peter Meerwald <pmeerw@cosy.sbg.ac.at>
parents:
0
diff
changeset
|
144 import Image, ImageDraw |
| 0 | 145 im = Image.new('RGB', (mvf.shape[1], mvf.shape[0])) |
| 146 d = ImageDraw.Draw(im) | |
| 147 | |
| 148 for r in xrange(mvf.shape[0]): | |
| 149 for c in xrange(mvf.shape[1]): | |
| 150 mv = mvf[r][c] | |
| 151 | |
| 152 if sr > 0: w = int(math.sqrt(mv[0]**2 + mv[1]**2)*255/(sr*math.sqrt(2.0))) | |
| 153 else: w = 0 | |
| 154 | |
| 155 if mv[2] >= 0 or mv[2] == CONNECTED: color = (0, w, 0) | |
| 156 elif mv[2] == UNCONNECTED: color = (255, 0, 0) | |
| 157 elif mv[2] == MULTIPLE_CONNECTED: color = (0, 0, w) | |
| 158 | |
| 159 d.point((c, r), fill=color) | |
| 160 | |
| 161 del d | |
| 162 im.save('%s-%02d-%04d.png' % (fname, level, idx), 'PNG') | |
| 163 del im | |
| 164 | |
| 165 def show_mvf(mvf, imvf, idx=0, level=0, bs=8, sr=8): | |
| 166 ''' | |
| 167 Visualize the motion field as .png and output motion vectors to .txt. | |
| 168 ''' | |
| 169 | |
| 170 im = Image.new('RGB', (mvf.shape[1]*2, mvf.shape[0]*2)) | |
| 171 d = ImageDraw.Draw(im) | |
| 172 f = open('mv-%02d-%04d.txt' % (level, idx), 'wt') | |
| 173 sad = mvf[:,:,2].ravel() | |
| 174 sad_min = numpy.min(numpy.where(sad < 0.0, 0, sad)) | |
| 175 sad_max = numpy.max(sad) | |
| 176 for r in xrange(0,mvf.shape[0],bs): | |
| 177 for c in xrange(0,mvf.shape[1],bs): | |
| 178 mv = mvf[r][c] | |
| 179 print >>f, '(%d %d)' % (mv[1], mv[0]), | |
| 180 | |
| 181 # fill block according to SAD | |
| 182 if sad_max > 0 and mv[2] > 0: | |
| 183 d.rectangle([(c*2,r*2),(c*2+bs*2,r*2+bs*2)], fill=((mv[2]-sad_min)*255/sad_max,0,0)) | |
| 184 | |
| 185 # draw motion vector | |
| 186 if sr > 0: w = int(math.sqrt(mv[0]**2 + mv[1]**2)/(sr*math.sqrt(2.0))) | |
| 187 else: w = 0 | |
| 188 | |
| 189 d.line([ \ | |
| 190 (c*2+bs, r*2+bs), \ | |
| 191 (c*2+bs+mv[1]*2, r*2+bs+mv[0]*2)], \ | |
| 192 fill=(0,int(32+(255-32)*w),0)) | |
| 193 d.point((c*2+bs, r*2+bs), fill=(255,255,255)) | |
| 194 | |
| 195 print >>f | |
| 196 print >>f | |
| 197 | |
| 198 f.close() | |
| 199 del d | |
| 200 | |
| 201 im.save('mv-%02d-%04d.png' % (level, idx), 'PNG') | |
| 202 del im | |
| 203 | |
| 204 _show_mv_dist(mvf, idx, level, sr, 'mvf_dist') | |
| 205 _show_mv_dist(imvf, idx, level, sr, 'mvi_dist') | |
| 206 | |
| 207 | |
| 208 def inverse_mvf(mvf, bs=8): | |
| 209 ''' | |
| 210 Compute the inverse of the motion field. | |
| 211 ''' | |
| 212 | |
| 213 imvf = numpy.zeros((mvf.shape[0], mvf.shape[1], 3), numpy.int) | |
| 214 imvf[:,:,2] = UNCONNECTED | |
| 215 for r in xrange(0, mvf.shape[0], bs): | |
| 216 for c in xrange(0, mvf.shape[1], bs): | |
| 217 rm = mvf[r,c,0] + r | |
| 218 cm = mvf[r,c,1] + c | |
| 219 | |
| 220 blockmvf = mvf[r:r+bs,c:c+bs] | |
| 221 blockimvf = imvf[rm:rm+bs,cm:cm+bs] | |
| 222 | |
| 223 # mark multiple connected in forward motion field if pixel already connected | |
| 224 numpy.place(blockmvf[:,:,2], blockimvf[:,:,2] > UNCONNECTED, MULTIPLE_CONNECTED) | |
| 225 | |
| 226 # invert motion vector and store in inverse motion field, mark pixel as connected | |
| 227 unconnected = blockimvf[:,:,2] == UNCONNECTED | |
| 228 numpy.place(blockimvf[:,:,0], unconnected, -mvf[r,c,0]) | |
| 229 numpy.place(blockimvf[:,:,1], unconnected, -mvf[r,c,1]) | |
| 230 numpy.place(blockimvf[:,:,2], unconnected, CONNECTED) | |
| 231 | |
| 232 return mvf, imvf | |
| 233 | |
| 234 def decompose_sequence(seq, Hs=[], MVFs=[], bs=8, sr=8, hlevel=2, tlp=MIDDLE, visualize_mvf=False, dlevel=-1): | |
| 235 ''' | |
| 236 Recursively decompose frame sequence using motion-compensated temporal filtering | |
| 237 employing the parameters blocksize, searchrange and hierarchy level for motion estimation. | |
| 238 | |
| 239 Output is [L], [H0, H1, H1, H2, H2, H2, H2], [MVF0, MVF1, MVF1, MVF2, MVF2, MVF2, MVF2] for | |
| 240 a sequence of length 8. | |
| 241 | |
| 242 The tlp argument allows to move the temporal low-pass frame to the left, | |
| 243 middle or right. | |
| 244 ''' | |
| 245 Ls = [] | |
| 246 if dlevel < 0: dlevel = int(math.log(len(seq), 2)) | |
| 247 | |
| 248 if len(seq) == 1: | |
| 249 return seq, Hs, MVFs | |
| 250 | |
| 251 if tlp == RIGHT: left = 0; mid = len(seq); right = 0 | |
| 252 elif tlp == LEFT: left = 0; mid = 0; right = len(seq) | |
| 253 else: left = 0; mid = max(len(seq)/2, 2); right = len(seq) | |
| 254 | |
| 255 for i in xrange(left, mid, 2): | |
| 256 sad, mvf = motion_estimation(seq[i+1], seq[i], bs, sr, hlevel) | |
| 257 mvf, imvf = inverse_mvf(mvf, bs) | |
| 258 if visualize_mvf: | |
| 259 show_mvf(mvf, imvf, i, dlevel-1, bs, sr) | |
| 260 MVFs.insert(i//2, mvf) | |
| 261 L, H = ft_mvf(seq[i], seq[i+1], mvf, imvf, bs) | |
| 262 Ls.append(L) | |
| 263 Hs.insert(i//2, H) | |
| 264 | |
| 265 for i in xrange(mid, right, 2): | |
| 266 sad, mvf = motion_estimation(seq[i], seq[i+1], bs, sr, hlevel) | |
| 267 mvf, imvf = inverse_mvf(mvf, bs) | |
| 268 if visualize_mvf: | |
| 269 show_mvf(mvf, imvf, i, dlevel-1, bs, sr) | |
| 270 MVFs.insert(i//2, mvf) | |
| 271 L, H = ft_mvf(seq[i+1], seq[i], mvf, imvf, bs) | |
| 272 Ls.append(L) | |
| 273 Hs.insert(i//2, H) | |
| 274 | |
| 275 return decompose_sequence(Ls, Hs, MVFs, bs, sr, hlevel, tlp, visualize_mvf, dlevel-1) | |
| 276 | |
| 277 def reconstruct_sequence(seq, Hs, MVFs, bs=8, tlp=MIDDLE): | |
| 278 ''' | |
| 279 Recursively reconstruct a frame sequence from temporal low- and high-pass subbands | |
| 280 and motion fields. | |
| 281 ''' | |
| 282 | |
| 283 Ls = [] | |
| 284 | |
| 285 if len(Hs) == 0: | |
| 286 return seq | |
| 287 | |
| 288 if tlp == RIGHT: left = 0; mid = len(seq); right = 0 | |
| 289 elif tlp == LEFT: left = 0; mid = 0; right = len(seq) | |
| 290 else: left = 0; mid = max(len(seq)/2, 1); right = len(seq) | |
| 291 | |
| 292 for i in xrange(0, mid): | |
| 293 mvf = MVFs[0] | |
| 294 mvf, imvf = inverse_mvf(mvf, bs) | |
| 295 a, b = it_mvf(seq[i], Hs[0], mvf, imvf, bs) | |
| 296 Ls += [a] + [b] | |
| 297 del Hs[0] | |
| 298 del MVFs[0] | |
| 299 | |
| 300 for i in xrange(mid, right): | |
| 301 mvf = MVFs[0] | |
| 302 mvf, imvf = inverse_mvf(mvf, bs) | |
| 303 a, b = it_mvf(seq[i], Hs[0], mvf, imvf, bs) | |
| 304 Ls += [b] + [a] | |
| 305 del Hs[0] | |
| 306 del MVFs[0] | |
| 307 | |
| 308 return reconstruct_sequence(Ls, Hs, MVFs, bs, tlp) | |
| 309 |