testing.py

import cv2
import pickle
import matplotlib.pyplot as plt
import numpy as np
import random

# a = ['SF_E_L_P008-000000_left.png/SF_E_L_P008-000001_left.png']
# rawImage = str(a)
# leftImage = rawImage.split("/")[0]
# rightImage = rawImage.split("/")[1]
# leftImageFolder = leftImage.split("-")[0]
# rightImageFolder = rightImage.split("-")[0]
# leftImageName = leftImage.split("-")[1]
# rightImageName = rightImage.split("-")[1]
# print(f"path left: {leftImageFolder}/{leftImageName}")
# print(f"path right: {rightImageFolder}/{rightImageName}")


def get_color(score):
    if score > 0.75:
        return (0, 255, 0)  # Green for strong matches
    elif score > 0.5:
        return (255, 255, 0)  # Yellow for medium matches
    elif score > 0.1:
        return (255, 0, 0)  # Red for low matches
    return (0, 0, 0)  # Black for scores <= 0.1



def filter_matches(matches, scores, percentage=60):
    filtered_matches = []
    filtered_scores = []

    # Split matches into categories based on scores
    red_matches = [(i, m) for i, m in enumerate(matches) if 0.1 < scores[i] <= 0.5]
    yellow_matches = [(i, m) for i, m in enumerate(matches) if 0.5 < scores[i] <= 0.75]
    green_matches = [(i, m) for i, m in enumerate(matches) if scores[i] > 0.75]

    # Randomly remove percentage of matches from each category
    for category in [red_matches, yellow_matches, green_matches]:
        n_remove = int(len(category) * percentage / 100)
        remove_indices = set(random.sample(range(len(category)), n_remove))
        for i, (idx, m) in enumerate(category):
            if i not in remove_indices:
                filtered_matches.append(m)
                filtered_scores.append(scores[idx])

    return filtered_matches, filtered_scores

def filter_percentage(matches, percentage):
    n_remove = int(len(matches) * percentage / 100)
    remove_indices = set(random.sample(range(len(matches)), n_remove))
    filtered_matches = [m for i, m in enumerate(matches) if i not in remove_indices]
    return filtered_matches

def visualize_matches(img1, img2, kp1, kp2, matches, scores):
    fig, ax = plt.subplots(1, 1, figsize=(15, 10))
    img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2RGB)
    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)
    img3 = cv2.hconcat([img1, img2])

    # Filter and pair matches
    
    valid_matches = [(i, m) for i, m in enumerate(matches) if m >= 0]
    # Filter matches if percentage is set
    # matches, scores = filter_matches(matches, scores, percentage=60)
    valid_matches = filter_percentage(valid_matches, percentage=90)
    
    for i, (idx0, idx1) in enumerate(valid_matches):
        color = get_color(scores[i])
        if scores[i] > 0.1:  # Filter to only strong matches
            pt1 = (int(kp1[idx0][0]), int(kp1[idx0][1]))
            pt2 = (int(kp2[idx1][0] + img1.shape[1]), int(kp2[idx1][1]))  # Offset by width of img1
            cv2.line(img3, pt1, pt2, color, 1)

    ax.imshow(img3)
    ax.axis('off')

    # Create custom legend
    from matplotlib.lines import Line2D
    legend_elements = [
        Line2D([0], [0], color=(1, 0, 0), lw=2, label='Low (0.1 - 0.5)'),
        Line2D([0], [0], color=(1, 1, 0), lw=2, label='Medium (0.5 - 0.75)'),
        Line2D([0], [0], color=(0, 1, 0), lw=2, label='High (> 0.75)')
    ]
    ax.legend(handles=legend_elements, loc='upper right')
    
    savePath = "/homes/tp4618/Documents/bitbucket/SuperGlueThesis/external/glue-factory/data/syntheticForestData/plotsMatchescolour.png"
    plt.savefig(savePath)
    print(f"Saved plot to {savePath}")
    plt.close(fig)
    
# # Create a dictionary with the data to be pickled
# pickle_data = {
#     "rawkp0": kp0,
#     "rawkp1": kp1,
#     "rawm0": m0,
#     "rawscores": scores,
#     "kp0": kp0.numpy(),
#     "kp1": kp1.numpy(),
#     "m0": m0.numpy(),
#     "scores": scores.numpy(),
#     "left_image": left_image,
#     "right_image": right_image
# }
    # Specify the path to save the pickled data
pickle_path = "/homes/tp4618/Documents/bitbucket/SuperGlueThesis/external/glue-factory/eval_pickled_data.pkl"

# Load the pickled data
with open(pickle_path, 'rb') as f:
    data = pickle.load(f)

# Access the variables from the loaded data
rawkp0 = data["rawkp0"]
rawkp1 = data["rawkp1"]
rawm0 = data["rawm0"]
rawscores = data["rawscores"]
kp0 = data["kp0"]
kp1 = data["kp1"]
m0 = data["m0"]
scores = data["scores"]
left_image = data["left_image"]
right_image = data["right_image"]

print("k0 is", kp0)
print("Shape of kp0:", kp0.shape)
print("Shape of kp1:", kp1.shape)
print("Shape of m0:", m0.shape)
print("Shape of scores:", scores.shape)


print("Summary:")
print("Number of keypoints in image 0:", len(kp0))
print("Number of keypoints in image 1:", len(kp1))
print("Number of matches:", len(m0))
print("Number of scores:", len(scores))

visualize_matches(left_image, right_image, kp0, kp1, m0, scores)