-
Notifications
You must be signed in to change notification settings - Fork 1
/
block_matching.py
223 lines (178 loc) · 7.49 KB
/
block_matching.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
# Author: Selahaddin HONI | 001honi@github
# March, 2021
# ============================================================================================
from utils import MAD, MSE
import numpy as np
import itertools
import cv2
class Block():
min = None
max = None
max_mv_amp = 0
def __init__(self,x,y,w,h):
self.coord = (x,y,w,h)
self.center = (x+w//2,y+h//2)
self.mv = (0,0)
self.mv_amp = 0
def check_inside_frame(self,x,y):
"""check if the searched box inside the target frame"""
check = True
if x<Block.min[0] or x>Block.max[0] \
or y<Block.min[1] or y>Block.max[1]:
check = False
return check
def calculate_mv_amp(self):
"""calculate L2 norm of motion-vector"""
amp = (self.mv[0]**2 + self.mv[1]**2)**0.5
if amp > Block.max_mv_amp:
Block.max_mv_amp = amp
self.mv_amp = amp
class BlockMatching():
def __init__(self,dfd,blockSize,searchMethod,searchRange,motionIntensity=True):
"""
dfd : {0:MAD, 1:MSE} Displaced frame difference
blockSize : {(sizeH,sizeW)}
searchMethod : {0:Exhaustive, 1:Three-Step}
searchRange : (int) +/- pixelwise range
"""
# given parameters
self.mse = dfd
self.blockSize = blockSize
self.searchMethod = searchMethod
self.searchRange = searchRange
self.motionIntensity = motionIntensity
# given frames
self.anchor = None
self.target = None
self.shape = None
# blocks in anchor
self.blocks = None
# return frames
self.anchorP = None
self.motionField = None
def step(self,anchor,target):
"""One-step run for given frame pair."""
self.anchor = anchor
self.target = target
self.shape = anchor.shape
self.frame2blocks()
if self.searchMethod == 0:
self.EBMA()
elif self.searchMethod == 1:
self.ThreeStepSearch()
else:
print("Search Method does not exist!")
self.plot_motionField()
self.blocks2frame()
def frame2blocks(self):
"""Divides the frame matrix into block objects."""
(H,W) = self.shape
(sizeH,sizeW) = self.blockSize
self.blocks = []
for h in range(H//sizeH):
for w in range(W//sizeW):
# initialize Block() objects with
# upper-left coordinates and block size
x = w*sizeW ; y = h*sizeH
self.blocks.append(Block(x,y,sizeW,sizeH))
# store the upper-left and bottom-right block coordinates
# for future check if the searched block inside the frame
Block.min = self.blocks[0].coord
Block.max = self.blocks[-1].coord
def blocks2frame(self):
"""Construct the predicted"""
frame = np.zeros(self.shape,dtype="uint8")
for block in self.blocks:
# get block coordinates for anchor frame
(x,y,w,h) = block.coord
# extract the block from anchor frame
block_a = self.anchor[y:y+h, x:x+w]
# append motion-vector to prediction coordinates
(x,y) = x+block.mv[0], y+block.mv[1]
# shift the block to new coordinates
frame[y:y+h, x:x+w] = block_a
self.anchorP = frame
def plot_motionField(self):
"""Construct the motion field from motion-vectors"""
frame = np.zeros(self.shape,dtype="uint8")
for block in self.blocks:
intensity = round(255 * block.mv_amp/Block.max_mv_amp) if self.motionIntensity else 255
intensity = 100 if intensity<100 else intensity
(x2,y2) = block.mv[0]+block.center[0], block.mv[1]+block.center[1]
cv2.arrowedLine(frame, block.center, (x2,y2), intensity, 1, tipLength=0.3)
self.motionField = frame
def EBMA(self):
"""Exhaustive Search Algorithm"""
dx = dy = [i for i in range(-self.searchRange,self.searchRange+1)]
searchArea = [r for r in itertools.product(dx,dy)]
for block in self.blocks:
# get block coordinates for anchor frame
(x,y,w,h) = block.coord
# extract the block from anchor frame
block_a = self.anchor[y:y+h, x:x+w]
# displaced frame difference := initially infinity
dfd_norm_min = np.Inf
# search the matched block in target frame in search area
for (dx,dy) in searchArea:
(x,y,w,h) = block.coord
# check if the searched box inside the target frame
if not block.check_inside_frame(x+dx,y+dy):
continue
x = x+dx ; y = y+dy
# extract the block from target frame
block_t = self.target[y:y+h, x:x+w]
# calculate displaced frame distance
if self.mse:
dfd_norm = MSE(block_a,block_t)
else:
dfd_norm = MAD(block_a,block_t)
if dfd_norm < dfd_norm_min:
# set the difference as motion-vector
block.mv = (dx,dy)
block.calculate_mv_amp()
dfd_norm_min = dfd_norm
def ThreeStepSearch(self):
"""Three-Step Search Algorithm"""
searchStep = [self.searchRange//2,self.searchRange//3,self.searchRange//6]
searchAreas = []
for step in searchStep:
dx = dy = [-step, 0, step]
searchAreas.append([r for r in itertools.product(dx,dy)])
for block in self.blocks:
step1 = self.OneStepSearch(block,block.coord,searchAreas[0])
step2 = self.OneStepSearch(block,step1,searchAreas[1])
step3 = self.OneStepSearch(block,step2,searchAreas[2])
# get best-match coordinates
(x,y,w,h) = step3
# set the difference as motion-vector
block.mv = (x-block.coord[0],y-block.coord[1])
block.calculate_mv_amp()
def OneStepSearch(self,block,searchCoord,searchArea):
"""Three-Step Search helper function"""
# get block coordinates for anchor frame
(x,y,w,h) = block.coord
# extract the block from anchor frame
block_a = self.anchor[y:y+h, x:x+w]
# displaced frame difference := initially infinity
dfd_norm_min = np.Inf
# best-match coordinates
coord = (x,y,w,h)
# search the matched block in target frame in search area
for (dx,dy) in searchArea:
(x,y,w,h) = searchCoord
# check if the searched box inside the target frame
if not block.check_inside_frame(x+dx,y+dy):
continue
x = x+dx ; y = y+dy
# extract the block from target frame
block_t = self.target[y:y+h, x:x+w]
# calculate displaced frame distance
if self.mse:
dfd_norm = MSE(block_a,block_t)
else:
dfd_norm = MAD(block_a,block_t)
# update best-match coordinates
if dfd_norm < dfd_norm_min:
dfd_norm_min = dfd_norm
coord = (x,y,w,h)
return coord