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