Untitled
unknown
python
a year ago
2.7 kB
3
Indexable
def twoD_log_search(ref_img, target_block, corner_pt, search_range, block_size):
# '''
# In :
# ref_img : reference image
# target_block : target block
# corner_pt : top-left corner coordinate of the target block
# search_range : search range
# block size : size of macroblock
# Out :
# predicted_block : predicted block
# motion vector : motion vector
# SAD : SAD value
# '''
# # TODO calculate the predicted block, motion vector, and SAD value
# Extract target block parameters
y, x = corner_pt
height, width, _ = target_block.shape
# Initialize variables
min_SAD = float('inf')
motion_vector = np.zeros((2, 1))
# Initial step size for logarithmic search
step_size = search_range // 2
i=x
j=y
# Perform logarithmic search
while step_size > 1:
search_points = [(0,0), (-step_size, 0), (step_size, 0), (0,-step_size), (0,step_size)]
dx=0
dy=0
# Extract reference block
for points in search_points:
if i+points[0] < 0 or i+points[0]+block_size > ref_img.shape[0] or j+points[1] < 0 or j+points[1]+block_size > ref_img.shape[1]:
continue
ref_block = ref_img[i+points[0]:i+block_size+points[0], j+points[1]:j+block_size+points[1]]
# Calculate SAD
current_SAD = calculate_SAD(ref_block, target_block)
if current_SAD < min_SAD:
min_SAD = current_SAD
dx = points[0]
dy = points[1]
motion_vector[0] = i+dx-x
motion_vector[1] = j+dy-y
predicted_block = ref_block.copy()
# Halve the step size
if dx == 0 and dy == 0:
step_size //= 2
else:
i += dx
j += dy
# motion_vector = np.zeros((2,1))
for n in range(-1,2):
for m in range(-1, 2):
if n+i < 0 or n+i+block_size > ref_img.shape[0] or m+j < 0 or m+j+block_size > ref_img.shape[1]:
continue
# print(y,j)
ref_block = ref_img[n+i:n+i+height, m+j:m+j+width]
current_SAD = calculate_SAD(ref_block, target_block)
if current_SAD < min_SAD:
# print("found")
min_SAD = current_SAD
motion_vector[0] = n+i - x
motion_vector[1] = m+j - y
predicted_block = ref_block
return predicted_block, motion_vector, min_SADEditor is loading...
Leave a Comment