CiscoTheProot/rpi/Point2D.py

162 lines
No EOL
5.2 KiB
Python

import hashlib
import json
import os.path
import math
from scipy.optimize import linear_sum_assignment
import numpy as np
from PIL import Image
# location to store array cache
CACHE_FILE_PATH = "/home/cisco/CiscoTheProot/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, point.color)
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:
print("Found existing point array matching png: " + image.filename + " using existing array.")
return [Point2D(point["x"], point["y"], tuple(point["color"])) for point in cached_point_arrays[image_hash]]
print("No existing point array matching png: " + image.filename + " found. Generating now.")
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)
print("Point array generated and stored.")
return [Point2D(point["x"], point["y"], tuple(point["color"])) for point in pixel_array]
def generate_image_from_point_array(points: list[Point2D], width: int, height: int) -> Image:
# 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] = point.color
return image
def interpolate_point_pairs(pairs: list[tuple[Point2D, Point2D]], percentage: float) -> list[Point2D]:
interpolated_points:list[Point2D] = []
for pair in pairs:
point1, point2 = pair
interpolated_point = point1.interpolate(point2, percentage)
interpolated_points.append(interpolated_point)
return interpolated_points
def pair_points(points1: list[Point2D], points2: list[Point2D]) -> list[tuple[Point2D, Point2D]]:
# 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