Untitled

 avatar
unknown
python
a year ago
2.7 kB
2
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_SAD
Editor is loading...
Leave a Comment