Untitled

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:
                # Extract reference block
        ref_block = ref_img[i:i+height, j:j+width]
                # Calculate SAD
        current_SAD = calculate_SAD(ref_block, target_block)
                # Check if current SAD is less than minimum SAD found so far
        if current_SAD < min_SAD:
            min_SAD = current_SAD
            # Update motion vector
            motion_vector[0] = i - x
            motion_vector[1] = j - y
            # Update predicted block
            predicted_block = ref_block.copy()
        
        if i-step_size >= 0:
            ref_block = ref_img[i-step_size:i-step_size+height, j:j+width]
            current_SAD = calculate_SAD(ref_block, target_block)
            
            if current_SAD < min_SAD:
                min_SAD = current_SAD
                motion_vector[0] = i - x
                motion_vector[1] = j - y
                predicted_block = ref_block.copy()
                
        if i+step_size+height <= ref_img.shape[0]:
            ref_block = ref_img[i+step_size:i+height+step_size, j:j+width]
            current_SAD = calculate_SAD(ref_block, target_block)
            
            if current_SAD < min_SAD:
                min_SAD = current_SAD
                motion_vector[0] = i - x
                motion_vector[1] = j - y
                predicted_block = ref_block.copy()

        if j-step_size >= 0:
            ref_block = ref_img[i:i+height, j-step_size:j+width-step_size]
            current_SAD = calculate_SAD(ref_block, target_block)
            
            if current_SAD < min_SAD:
                min_SAD = current_SAD
                motion_vector[0] = i - x
                motion_vector[1] = j - y
                predicted_block = ref_block.copy()

        if j+step_size+width <= ref_img.shape[1]:
            ref_block = ref_img[i:i+height, j+step_size:j+width+step_size]
            current_SAD = calculate_SAD(ref_block, target_block)
            
            if current_SAD < min_SAD:
                min_SAD = current_SAD
                motion_vector[0] = i - x
                motion_vector[1] = j - y
                predicted_block = ref_block.copy()

        # Halve the step size
        if motion_vector[0] == 0 and motion_vector[1] == 0:
            step_size //= 2

        else:
            i += motion_vector[0]
            j += motion_vector[1]
        
        motion_vector = np.zeros((2,1))

    for n in range(max(0, i-1), min(ref_img.shape[0]-height, i+2)):
        for m in range(max(0, j-1), min(ref_img.shape[1]-width, j+2)):
            ref_block = ref_img[n:n+height, m:m+width]
            current_SAD = calculate_SAD(ref_block, target_block)
            
            if current_SAD < min_SAD:
                # print("found")
                min_SAD = current_SAD
                motion_vector[0] = n - x
                motion_vector[1] = m - y
                predicted_block = ref_block


    return predicted_block, motion_vector, min_SAD