CiscoTheProot/minDistance.py

import math
from PIL import Image
import numpy as np
from scipy.optimize import linear_sum_assignment
import json
import os.path
import hashlib


# location to store array cache
CACHE_FILE_PATH = "point_array_cache.json"

    
class Point2D:
    x = 0
    y = 0
    color: tuple[int, int, int] = (0, 0, 0)

    def __init__(self, x, y, color: tuple[int, int, int] = (0, 0, 0)):
        self.x = x
        self.y = y
        self.color = color

    def round(self):
        self.x = round(self.x)
        self.y = round(self.y)
        return self

    def distance(self, other):
        dx = self.x - other.x
        dy = self.y - other.y
        return math.sqrt(dx ** 2 + dy ** 2)

    def interpolate(self, other, percentage):
        new_x = self.x + (other.x - self.x) * percentage
        new_y = self.y + (other.y - self.y) * percentage
        new_color = tuple(int((1 - percentage) * self.color[i] + percentage * other.color[i]) for i in range(3))
        return Point2D(new_x, new_y, new_color)

    def __eq__(self, other):
        return (self.x, self.y) == (other.x, other.y)


def mirror_points(points: list[Point2D]) -> list[Point2D]:
    mirrored_points = []
    for point in points:
        mirrored_x = 128 - point.x  # Calculate the mirrored x-coordinate
        mirrored_point = Point2D(mirrored_x, point.y)
        mirrored_points.append(mirrored_point)
    return mirrored_points



def get_image_hash(image):
    image_hash = hashlib.sha1(image.tobytes()).hexdigest()
    return image_hash


def load_cached_point_arrays():
    cached_point_arrays = {}

    if os.path.isfile(CACHE_FILE_PATH):
        with open(CACHE_FILE_PATH, "r") as file:
            cached_point_arrays = json.load(file)

    return cached_point_arrays


def save_cached_point_arrays(cached_point_arrays):
    with open(CACHE_FILE_PATH, "w") as file:
        json.dump(cached_point_arrays, file)


def generate_point_array_from_image(image):
    image = image.convert("RGB")  # Convert image to RGB color mode
    image_hash = get_image_hash(image)
    cached_point_arrays = load_cached_point_arrays()

    if image_hash in cached_point_arrays:
        return [Point2D(point["x"], point["y"], tuple(point["color"])) for point in cached_point_arrays[image_hash]]

    width, height = image.size
    pixel_array = []

    for y in range(height):
        for x in range(width):
            pixel = image.getpixel((x, y))
            if pixel != (0, 0, 0):  # any non-white pixels
                point = {"x": x, "y": y, "color": pixel}
                pixel_array.append(point)

    cached_point_arrays[image_hash] = pixel_array
    save_cached_point_arrays(cached_point_arrays)

    return [Point2D(point["x"], point["y"], tuple(point["color"])) for point in pixel_array]


def generate_image_from_point_array(points, width, height):
    # Create a new blank image
    image = Image.new("RGB", (width, height), "black")
    
    # Set the pixels corresponding to the points as white
    pixels = image.load()
    for point in points:
        point = point.round()
        x = point.x
        y = point.y
        pixels[x, y] = (255, 255, 255)
    
    return image


def pair_points(points1, points2):
    # Determine the size of the point arrays
    size1 = len(points1)
    size2 = len(points2)
    
    # Create a cost matrix based on the distances between points
    cost_matrix = np.zeros((size1, size2))
    for i in range(size1):
        for j in range(size2):
            cost_matrix[i, j] = points1[i].distance(points2[j])
    
    # Duplicate points in the smaller array to match the size of the larger array
    if size1 > size2:
        num_duplicates = size1 - size2
        duplicated_points = np.random.choice(points2, size=num_duplicates).tolist()
        points2 += duplicated_points
    elif size2 > size1:
        num_duplicates = size2 - size1
        duplicated_points = np.random.choice(points1, size=num_duplicates).tolist()
        points1 += duplicated_points
    
    # Update the size of the point arrays
    size1 = len(points1)
    size2 = len(points2)
    
    # Create a new cost matrix with the updated sizes
    cost_matrix = np.zeros((size1, size2))
    for i in range(size1):
        for j in range(size2):
            cost_matrix[i, j] = points1[i].distance(points2[j])
    
    # Solve the assignment problem using the Hungarian algorithm
    row_ind, col_ind = linear_sum_assignment(cost_matrix)
    
    # Create pairs of points based on the optimal assignment
    pairs = []
    for i, j in zip(row_ind, col_ind):
        pairs.append((points1[i], points2[j]))
    
    return pairs


def interpolate_point_pairs(pairs, percentage):
    interpolated_points = []
    for pair in pairs:
        point1, point2 = pair
        interpolated_point = point1.interpolate(point2, percentage)
        interpolated_points.append(interpolated_point)
    return interpolated_points


Image1 = Image.open("faces/prootface3.png")
Image2 = Image.open("faces/prootface4.png")

pixelArray1 = generate_point_array_from_image(Image1)
pixelArray2 = generate_point_array_from_image(Image2)

pairs = pair_points(pixelArray1, pixelArray2)


generate_image_from_point_array(interpolate_point_pairs(pairs, 0), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, .25), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, .5), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, .75), 128, 32).show()
generate_image_from_point_array(interpolate_point_pairs(pairs, 1), 128, 32).show()

print(pairs)