annotate pymctf.py @ 0:4214d9245f8e

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

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