test.py

from PIL import Image
import numpy as np
import scipy as sp
import os

# Dictionary of known images
KNOWN_IMAGES = {
    "0": "out/0.png",
    "1": "out/1.png",
    "2": "out/2.png",
    "3": "out/3.png",
    "4": "out/4.png",
    "5": "out/5.png",
    "6": "out/6.png",
    "7": "out/7.png",
    "8": "out/8.png",
    "9": "out/9.png",
}

# List of unknown image paths
UNKNOWN_IMAGES = [
    "out/u1.png",
    "out/u2.png",
    "out/u3.png",
    "out/u4.png",
    "out/u5.png",
    "out/u6.png",
    "out/u7.png",
    "out/u8.png",
    "out/u9.png",
    "out/u10.png",
    "out/u11.png",
    "out/u12.png",
    "out/u13.png",
    "out/u14.png",
    "out/u15.png",
    "out/u16.png",
    "out/u17.png",
    "out/u18.png",
]


def fast_box_blur(image, radius, repeat=1):
    img_float = np.array(image)

    is_grayscale = len(img_float.shape) == 2

    if is_grayscale:
        padded = np.pad(img_float, radius, mode="edge")
    else:
        padded = np.pad(
            img_float, ((radius, radius), (radius, radius), (0, 0)), mode="edge"
        )

    kernel_size = 2 * radius + 1

    result = np.copy(padded)
    for _ in range(repeat):
        # Horizontal pass
        kernel = np.ones(kernel_size) / kernel_size
        if is_grayscale:
            for i in range(padded.shape[0]):
                result[i, :] = np.convolve(padded[i, :], kernel, mode="same")
        else:
            for i in range(padded.shape[0]):
                for c in range(padded.shape[2]):
                    result[i, :, c] = np.convolve(padded[i, :, c], kernel, mode="same")

        # Vertical pass
        if is_grayscale:
            for j in range(padded.shape[1]):
                result[:, j] = np.convolve(result[:, j], kernel, mode="same")
        else:
            for j in range(padded.shape[1]):
                for c in range(padded.shape[2]):
                    result[:, j, c] = np.convolve(result[:, j, c], kernel, mode="same")

        padded = result

    # Crop the padding
    if is_grayscale:
        result = result[radius:-radius, radius:-radius]
    else:
        result = result[radius:-radius, radius:-radius, :]

    # Clip values to valid range
    return Image.fromarray(np.clip(result, 0, 255).astype(np.uint8))


def calculate_similarity(img1, img2):
    arr1 = np.array(img1)
    arr2 = np.array(img2)

    # Ensure both images are the same size
    if arr1.shape != arr2.shape:
        # Resize second image to match first image's dimensions just in case i messed up splitting
        img2 = img2.resize(img1.size)
        arr2 = np.array(img2)

    # Calculate Mean Squared Error
    mse = np.mean((arr1 - arr2) ** 2)
    return mse


def find_closest_match(unknown_img, known_images):
    """Find the closest matching known image"""
    similarities = {}

    # Calculate similarity with each known image
    for name, img in known_images.items():
        similarity = calculate_similarity(unknown_img, img)
        similarities[name] = similarity

    # Return the name of the most similar image (lowest MSE)
    return min(similarities.items(), key=lambda x: x[1])


def main():
    os.makedirs("out/blurred", exist_ok=True)
    # Load and blur known images
    known_blurred = {}
    for name, path in KNOWN_IMAGES.items():
        try:
            img = Image.open(path).convert("RGB")
            blurred = fast_box_blur(img, radius=20, repeat=3)
            known_blurred[name] = blurred

            # Save blurred image
            blurred_path = f"out/blurred/{name}_blurred.png"
            blurred.save(blurred_path)
            print(f"Saved blurred image to {blurred_path}")
        except Exception as e:
            print(f"Error loading known image {name}: {e}")
            continue

    for unknown_path in UNKNOWN_IMAGES:
        try:
            unknown_img = Image.open(unknown_path).convert("RGB")

            # Find closest match
            closest_name, similarity = find_closest_match(unknown_img, known_blurred)

            print(f"Unknown image {unknown_path}:")
            print(f"Closest match: {closest_name}")
            print(f"Similarity score (MSE): {similarity}")
            print("-" * 50)

        except Exception as e:
            print(f"Error processing unknown image {unknown_path}: {e}")
            continue


if __name__ == "__main__":
    main()