subpixel_refinement.py

"subpixel_refinement.py"
"Single octave matching with implemented subpixel refinement"
import cv2
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

def load_and_prepare_image(path):
    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE).astype(np.float32)
    return img

# downsample and blur the image according to desired octave level
def prepare_octave_input(img, octave_level, base_sigma=1.6, assumed_blur=0.5):
    # downscale the image if we're not at the base octave
    if octave_level > 0:
        scale = 1 / (2 ** octave_level)
        img = cv2.resize(img, (0, 0), fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    # compute the additional blur needed
    sigma_diff = np.sqrt(max((base_sigma ** 2) - (assumed_blur ** 2), 0.01))
    blurred = cv2.GaussianBlur(img, (0, 0), sigmaX=sigma_diff, sigmaY=sigma_diff)

    return blurred

def generate_single_octave(img, num_levels, base_sigma):
    s = num_levels - 3
    k = 2 ** (1/s)
    blurred_images = []
    for i in range(num_levels):
        if i == 0:
            # First level: apply initial blur to the base image (usually to reach σ = 1.6)
            sigma = base_sigma
            blurred = cv2.GaussianBlur(img, (0, 0), sigmaX=sigma)
            blurred_images.append(blurred)
            prev_sigma = sigma
            prev_img = blurred
        else:
            # Subsequent levels: apply *incremental* blur from previous image
            sigma_total = base_sigma * (k ** i)
            sigma_diff = np.sqrt(sigma_total**2 - prev_sigma**2)
            blurred = cv2.GaussianBlur(prev_img, (0, 0), sigmaX=sigma_diff)
            blurred_images.append(blurred)
            prev_img = blurred
            prev_sigma = sigma_total
    return blurred_images, k

def compute_dog_pyramid(octave_images):
    dog_pyramid = []
    # computes pyramid by subtracting adjacent images to highlight features
    for octave in octave_images:
        dogs = [octave[i] - octave[i - 1] for i in range(1, len(octave))]
        dog_pyramid.append(dogs)

    # visualization of dog
    # fig, axs = plt.subplots(len(dog_pyramid), len(dog_pyramid[0]), figsize=(3 * len(dog_pyramid[0]), 2.5 * len(dog_pyramid)), constrained_layout=True)
    # for o, octave in enumerate(dog_pyramid):
    #     for i, dog in enumerate(octave):
    #         ax = axs[o, i] if len(dog_pyramid) > 1 else axs[i]
    #         ax.imshow(dog, cmap='gray')
    #         ax.set_title(f'DoG O{o} L{i}', fontsize=8)
    #         ax.axis('off')
    # plt.show()

    return dog_pyramid

def compute_hessian(dog, y, x):
    # approximating partial second derivative with taylor series 
    dxx = dog[y, x+1] + dog[y, x-1] - 2 * dog[y, x]
    dyy = dog[y+1, x] + dog[y-1, x] - 2 * dog[y, x]
    dxy = (dog[y+1, x+1] - dog[y+1, x-1] - dog[y-1, x+1] + dog[y-1, x-1]) / 4.0
    
    return dxx, dyy, dxy

# r = 10 is the ratio of eigenvalues (one large one small = edge)
# can this be changed too
def edge_keypoint(dxx, dyy, dxy, r=10):
    Tr = dxx + dyy
    Det = dxx * dyy - dxy ** 2
    if Det < 0:
        return True # unstable with negative determinant
    ratio = (Tr ** 2) / Det
    threshold = ((r+1) ** 2) / r
    if ratio > threshold:
        return True
    return False

def detect_keypoints(dog_pyramid, octave_start_index):
    keypoints = []
    for rel_o, octave in enumerate(dog_pyramid):
        o = rel_o + octave_start_index
        for i in range(1, len(octave) - 1):
            prev_img = octave[i - 1]
            curr_img = octave[i]
            next_img = octave[i + 1]

            # compares current pixel to all of its neighboring pixels to determine
            # whether its a local extremum 
            for y in range(1, curr_img.shape[0] - 1):
                for x in range(1, curr_img.shape[1] - 1):
                    val = curr_img[y, x]
                    patch_prev = prev_img[y - 1:y + 2, x - 1:x + 2]
                    patch_curr = curr_img[y - 1:y + 2, x - 1:x + 2]
                    patch_next = next_img[y - 1:y + 2, x - 1:x + 2]

                    neighbors = np.concatenate([
                        patch_prev.flatten(),
                        patch_curr.flatten(),
                        patch_next.flatten()
                    ])
                    neighbors = np.delete(neighbors, 9 + 9)
                    if val >= np.max(neighbors) or val <= np.min(neighbors):
                        keypoints.append((x, y, o, i))
    return keypoints

def show_keypoints(img, keypoints, k, base_sigma):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(img, cmap='gray')
    ax.set_title("Satellite Keypoints")
    ax.axis('off')

    for x, y, o, level_idx in keypoints:
        sigma = base_sigma * (k ** level_idx) * (2 ** o)
        circ = patches.Circle((x, y), radius=3 * sigma, edgecolor='lime', facecolor='none', linewidth=1.5)
        ax.add_patch(circ)
    plt.show()

def compute_derivatives(curr_img, prev_img, next_img, x, y):
    # First derivatives (gradient)
    dx = (curr_img[y, x + 1] - curr_img[y, x - 1]) / 2.0
    dy = (curr_img[y + 1, x] - curr_img[y - 1, x]) / 2.0
    ds = (next_img[y, x] - prev_img[y, x]) / 2.0

    # Second derivatives (Hessian)
    dxx = curr_img[y, x + 1] + curr_img[y, x - 1] - 2 * curr_img[y, x]
    dyy = curr_img[y + 1, x] + curr_img[y - 1, x] - 2 * curr_img[y, x]
    dss = next_img[y, x] + prev_img[y, x] - 2 * curr_img[y, x]

    dxy = (curr_img[y + 1, x + 1] - curr_img[y + 1, x - 1] -
           curr_img[y - 1, x + 1] + curr_img[y - 1, x - 1]) / 4.0
    dxs = (next_img[y, x + 1] - next_img[y, x - 1] -
           prev_img[y, x + 1] + prev_img[y, x - 1]) / 4.0
    dys = (next_img[y + 1, x] - next_img[y - 1, x] -
           prev_img[y + 1, x] + prev_img[y - 1, x]) / 4.0

    grad = np.array([dx, dy, ds])
    H = np.array([
        [dxx, dxy, dxs],
        [dxy, dyy, dys],
        [dxs, dys, dss]
    ])

    return grad, H, dxx, dyy, dxy

# solve matrix for offset value
def solve_for_offset(H, grad):
    try:
        return -np.linalg.solve(H, grad)
    except np.linalg.LinAlgError:
        return None
    
def contrast_check(contrast_threshold, D, grad, offset):
    contrast = D + 0.5 * np.dot(grad, offset)
    return abs(contrast) >= contrast_threshold

def edge_check(dxx, dyy, dxy, r=10):
    trH = dxx + dyy
    detH = dxx * dyy - dxy * dxy
    if detH <= 0:
        return False
    return (trH ** 2) / detH <= ((r + 1) ** 2) / r

# localizes keypoint location down to subpixel level
def subpixel_keypoint_localization(dog_pyramid, keypoints, start_index, contrast_threshold=0.03, max_iterations=5):
    refined_keypoints = []
    for x, y, o, i in keypoints:
        rel_o = o - start_index  # adjust to relative index
        octave = dog_pyramid[rel_o]
        if i == 0 or i == len(octave) - 1:
            continue  # skip scale boundaries

        curr_img = octave[i]
        prev_img = octave[i - 1]
        next_img = octave[i + 1]
        xi = np.array([0.0, 0.0, 0.0])

        for _ in range(max_iterations):
            sample_x = int(round(x + xi[0]))
            sample_y = int(round(y + xi[1]))

            # bounds check
            if (sample_y < 1 or sample_y >= curr_img.shape[0] - 1 or
                sample_x < 1 or sample_x >= curr_img.shape[1] - 1):
                break  # stop refining, out of bounds

            grad, H, dxx, dyy, dxy = compute_derivatives(curr_img, prev_img, next_img, sample_x, sample_y)

            # find offset to refine keypoint location
            offset = solve_for_offset(H, grad)
            if offset is None:
                break
            xi = offset
            # if offset too small then stop refining
            if np.max(np.abs(xi)) < 0.5:
                break
        final_x = x + xi[0]
        final_y = y + xi[1]

        sample_x = int(round(final_x))
        sample_y = int(round(final_y))

        # check for edges, contrast, and boundaries of image
        height, width = curr_img.shape
        if sample_y < 0 or sample_y >= height or sample_x < 0 or sample_x >= width:
            continue
        if not contrast_check(contrast_threshold, curr_img[sample_y, sample_x], grad, xi):
            continue
        if not edge_check(dxx, dyy, dxy):
            continue
        refined_keypoints.append((final_x, final_y, o, i))

    return refined_keypoints

def assign_orientations(keypoints, octave_images, base_sigma, k):
    num_bins = 36
    bin_width = 360 // num_bins
    radius_factor = 3
    oriented_keypoints = []

    for x, y, o, i in keypoints:
        # scale of feature based on corresponding k value and level in octave
        scale = base_sigma * (k ** i)
        radius = int(radius_factor * scale)
        image = octave_images[o][i]

        # uses sobel filters to compute gradient magnitude and direction
        dx = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
        dy = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)

        grad_mag = np.sqrt(dx ** 2 + dy ** 2)
        grad_ori = (np.degrees(np.arctan2(dy, dx)) + 360) % 360

        # histogram with 36 bins (each bin represents 10 degrees)
        hist = np.zeros(num_bins, dtype=np.float32)

        # uses gaussian weighted neighborhood around keypoint
        for u in range(-radius, radius + 1):
            for v in range(-radius, radius + 1):
                xp = int(round(x + u))
                yp = int(round(y + v))
                if 0 <= xp < image.shape[1] and 0 <= yp < image.shape[0]:
                    weight = np.exp(-(u ** 2 + v ** 2) / (2 * (scale ** 2)))
                    magnitude = grad_mag[yp, xp]
                    orientation = grad_ori[yp, xp]
                    bin_idx = int(round(orientation)) % 360 // bin_width
                    # put in corresponding bin of the histogram
                    hist[bin_idx] += weight * magnitude
        # find dominant orientation of gradient based on histogram results
        max_val = np.max(hist)
        for bin_idx, val in enumerate(hist):
            if val >= 0.8 * max_val:
                angle = bin_idx * bin_width
                oriented_keypoints.append((x, y, o, i, angle))
    return oriented_keypoints

def generate_descriptors(oriented_keypoints, octave_images, base_sigma, k):
    descriptors = []
    for x, y, o, i, angle in oriented_keypoints:
        image = octave_images[o][i]
        dx = cv2.Sobel(image, cv2.CV_32F, 1, 0, ksize=3)
        dy = cv2.Sobel(image, cv2.CV_32F, 0, 1, ksize=3)
        magnitude = np.sqrt(dx ** 2 + dy ** 2)
        orientation = (np.degrees(np.arctan2(dy, dx)) + 360) % 360

        # window of keypoint descriptor
        scale = base_sigma * (k ** i)
        window_size = int(round(16 * scale))
        half_window = window_size // 2
        descriptor = np.zeros((4, 4, 8), dtype=np.float32)

        for u in range(-half_window, half_window):
            for v in range(-half_window, half_window):
                # rotate coordinate system to align with keypoint orientation
                theta = -np.radians(angle)
                rot_u = int(np.round(np.cos(theta) * u - np.sin(theta) * v))
                rot_v = int(np.round(np.sin(theta) * u + np.cos(theta) * v))

                xp = int(round(x + rot_u))
                yp = int(round(y + rot_v))

                if not (0 <= xp < image.shape[1] and 0 <= yp < image.shape[0]):
                    continue

                mag = magnitude[yp, xp]
                ori = orientation[yp, xp]
                relative_ori = (ori - angle + 360) % 360
                bin_idx = int(relative_ori // 45) % 8

                cell_x = (u + half_window) // (window_size // 4)
                cell_y = (v + half_window) // (window_size // 4)

                if 0 <= cell_x < 4 and 0 <= cell_y < 4:
                    descriptor[cell_y, cell_x, bin_idx] += mag

        vec = descriptor.flatten()
        vec /= np.linalg.norm(vec) + 1e-7
        vec = np.clip(vec, 0, 0.2)
        vec /= np.linalg.norm(vec) + 1e-7
        descriptors.append(vec)
    return descriptors

def match_descriptors(desc1, desc2, ratio_thresh=0.85):
    matches = []
    for i, d1 in enumerate(desc1):
        # calculate all distances from desc1 descriptor to desc2
        distances = np.linalg.norm(desc2-d1, axis = 1)
        if len(distances) < 2:
            continue
        # sort the distances in ascending order - first one is best match
        nearest = np.argsort(distances)
        # use lowe's ratio test to remove ambiguous matches
        if distances[nearest[0]] < ratio_thresh * distances[nearest[1]]:
            matches.append((i, nearest[0]))

    return matches

def draw_matches(img1, kp1, img2, kp2, matches):
    h1, w1 = img1.shape
    h2, w2 = img2.shape
    canvas = np.zeros((max(h1, h2), w1 + w2, 3), dtype=np.uint8)

    img1_color = cv2.cvtColor(img1.astype(np.uint8), cv2.COLOR_GRAY2BGR)
    img2_color = cv2.cvtColor(img2.astype(np.uint8), cv2.COLOR_GRAY2BGR)
    canvas[:h1, :w1] = img1_color
    canvas[:h2, w1:] = img2_color

    for i1, i2 in matches:
        x1, y1, o1, l1, _ = kp1[i1]
        x2, y2, o2, l2, _ = kp2[i2]
        pt1 = (int(x1 * (2 ** o1)), int(y1 * (2 ** o1)))
        pt2 = (int(x2 * (2 ** o2)) + w1, int(y2 * (2 ** o2)))
        color = tuple(np.random.randint(0, 255, 3).tolist())
        cv2.line(canvas, pt1, pt2, color, 1)
        cv2.circle(canvas, pt1, 3, color, -1)
        cv2.circle(canvas, pt2, 3, color, -1)

    plt.figure(figsize=(14, 7))
    plt.imshow(canvas[..., ::-1])
    plt.title(f"Matched keypoints: {len(matches)}")
    plt.axis('off')
    plt.show()

satellite = load_and_prepare_image("images/satellite_copy.png")
warped = load_and_prepare_image("images/warped_ground.png")

octave_level1 = 1
octave_level2 = 0
base_sigma = 1.6

octave_input1 = prepare_octave_input(warped, octave_level1, base_sigma)
octave_input2 = prepare_octave_input(satellite, octave_level2, base_sigma)

num_levels = 5
single_octave1, k1 = generate_single_octave(octave_input1, num_levels, base_sigma)
single_octave2, k2 = generate_single_octave(octave_input2, num_levels, base_sigma)

dog_pyramid1 = compute_dog_pyramid([single_octave1])
dog_pyramid2 = compute_dog_pyramid([single_octave2])

keypoints1 = detect_keypoints(dog_pyramid1, octave_level1)
keypoints2 = detect_keypoints(dog_pyramid2, octave_level2)

print(f"Detected {len(keypoints1)} keypoints at octave {octave_level1} for ground")
print(f"Detected {len(keypoints2)} keypoints at octave {octave_level2} for warped")

refined_keypoints1 = subpixel_keypoint_localization(dog_pyramid1, keypoints1, octave_level1)
refined_keypoints2 = subpixel_keypoint_localization(dog_pyramid2, keypoints2, octave_level2)

print(f"Detected {len(refined_keypoints1)} keypoints at octave {octave_level1} for ground")
print(f"Detected {len(refined_keypoints2)} keypoints at octave {octave_level2} for warped")

# force all keypoints to have o = 0 to line up arrays 
keypoints1_fixed = [(x, y, 0, i) for (x, y, o, i) in refined_keypoints1]
oriented_kp1 = assign_orientations(keypoints1_fixed, [single_octave1], base_sigma, k1)
keypoints2_fixed = [(x, y, 0, i) for (x, y, o, i) in refined_keypoints2]
oriented_kp2 = assign_orientations(keypoints2_fixed, [single_octave2], base_sigma, k2)

desc1 = generate_descriptors(oriented_kp1, [single_octave1], base_sigma, k1)
desc2 = generate_descriptors(oriented_kp2, [single_octave2], base_sigma, k2)

matches = match_descriptors(np.array(desc2), np.array(desc1))

draw_matches(octave_input2, oriented_kp2, octave_input1, oriented_kp1, matches)