299 lines
No EOL
10 KiB
Python
299 lines
No EOL
10 KiB
Python
import hashlib
|
|
from itertools import permutations
|
|
import json
|
|
import os.path
|
|
import math
|
|
from scipy.optimize import linear_sum_assignment
|
|
import numpy as np
|
|
from PIL import Image
|
|
import fastproot
|
|
|
|
|
|
# location to store array cache
|
|
CACHE_FILE_PATH = "./cache//point_array_cache.json"
|
|
LOADED_IMAGE_PATH = "./graphics/loaded.png"
|
|
|
|
CACHE_FILE_PATH = os.path.normpath(CACHE_FILE_PATH)
|
|
LOADED_IMAGE_PATH = os.path.normpath(LOADED_IMAGE_PATH)
|
|
|
|
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 divide_points_into_groups(points):
|
|
print("divide_points_into_groups: ", len(points))
|
|
result = fastproot.divide_into_groups(points)
|
|
res = [[Point2D(p[0], p[1], p[2]) for p in arr] for arr in result]
|
|
return res
|
|
|
|
def _divide_points_into_groups(points):
|
|
def flood_fill(point, group):
|
|
if point not in group:
|
|
group.append(point)
|
|
|
|
neighbors = [(point.x + 1, point.y), (point.x - 1, point.y), (point.x, point.y + 1), (point.x, point.y - 1)]
|
|
for neighbor_coords in neighbors:
|
|
neighbor = next((p for p in points if p.x == neighbor_coords[0] and p.y == neighbor_coords[1]), None)
|
|
if neighbor and neighbor not in group:
|
|
flood_fill(neighbor, group)
|
|
|
|
groups = []
|
|
remaining_points = [point for point in points if point.x < 64] # Filter points with x < 64
|
|
|
|
while remaining_points:
|
|
group = []
|
|
flood_fill(remaining_points[0], group)
|
|
groups.append(group)
|
|
|
|
# Remove points in the group from the remaining points
|
|
remaining_points = [point for point in remaining_points if point not in group]
|
|
|
|
return groups
|
|
|
|
|
|
def compute_bounding_box(points: list[Point2D]) -> tuple[float, float, float, float]:
|
|
min_x = min(point.x for point in points)
|
|
max_x = max(point.x for point in points)
|
|
min_y = min(point.y for point in points)
|
|
max_y = max(point.y for point in points)
|
|
return min_x, max_x, min_y, max_y
|
|
|
|
def calculate_distance(points1: list[Point2D], points2: list[Point2D]) -> float:
|
|
# Calculate the distance between the centers of two groups
|
|
center1_x = (points1[0].x + points1[-1].x) / 2
|
|
center1_y = (points1[0].y + points1[-1].y) / 2
|
|
center2_x = (points2[0].x + points2[-1].x) / 2
|
|
center2_y = (points2[0].y + points2[-1].y) / 2
|
|
return ((center1_x - center2_x) ** 2 + (center1_y - center2_y) ** 2) ** 0.5
|
|
|
|
|
|
def pair_groups(set_a: list[list[Point2D]], set_b: list[list[Point2D]]) -> list[tuple[list[Point2D], list[Point2D]]]:
|
|
# print("set_a", [[(b.x, b.y) for b in i] for i in set_a])
|
|
# print("set_b", [[(b.x, b.y) for b in i] for i in set_b])
|
|
result = []
|
|
if len(set_a) <= len(set_b):
|
|
result = fastproot.pair_groups(set_a, set_b)
|
|
else:
|
|
for (b, a) in fastproot.pair_groups(set_b, set_a):
|
|
result.append((a,b))
|
|
res = [
|
|
(
|
|
[Point2D(a1[0], a1[1], a1[2]) for a1 in a],
|
|
[Point2D(b1[0], b1[1], b1[2]) for b1 in b]
|
|
) for (a, b) in result
|
|
]
|
|
#res = [[Point2D(p[0], p[1], p[2]) for p in arr] for arr in result]
|
|
return res
|
|
|
|
|
|
def _pair_groups(set_a: list[list[Point2D]], set_b: list[list[Point2D]]) -> list[tuple[list[Point2D], list[Point2D]]]:
|
|
pairs = []
|
|
|
|
# Create dictionaries to store bounding boxes for each group
|
|
bounding_boxes_a: dict[int, tuple[float, float, float, float]] = {}
|
|
bounding_boxes_b: dict[int, tuple[float, float, float, float]] = {}
|
|
|
|
# Calculate bounding boxes for all groups in both sets
|
|
for i, group in enumerate(set_a + set_b):
|
|
bounding_box = compute_bounding_box(group)
|
|
if i < len(set_a):
|
|
bounding_boxes_a[i] = bounding_box
|
|
else:
|
|
bounding_boxes_b[i - len(set_a)] = bounding_box
|
|
|
|
# Check for overlaps and determine pairs
|
|
for i, group_a in enumerate(set_a):
|
|
overlap_detected = False
|
|
for j, group_b in enumerate(set_b):
|
|
bounding_box_a = bounding_boxes_a[i]
|
|
bounding_box_b = bounding_boxes_b[j]
|
|
if (
|
|
bounding_box_a[0] <= bounding_box_b[1] and
|
|
bounding_box_a[1] >= bounding_box_b[0] and
|
|
bounding_box_a[2] <= bounding_box_b[3] and
|
|
bounding_box_a[3] >= bounding_box_b[2]
|
|
):
|
|
pairs.append((group_a, group_b))
|
|
overlap_detected = True
|
|
break
|
|
|
|
if not overlap_detected:
|
|
# Find the nearest neighbor in set B
|
|
nearest_group = min(set_b, key=lambda group: calculate_distance(group_a, group))
|
|
pairs.append((group_a, nearest_group))
|
|
|
|
return pairs
|
|
|
|
def mirror_points(points: list[Point2D]) -> list[Point2D]:
|
|
mirrored_points = []
|
|
for point in points:
|
|
mirrored_x = 127 - 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):
|
|
if not os.path.isfile(CACHE_FILE_PATH):
|
|
open(CACHE_FILE_PATH, "x").close()
|
|
|
|
with open(CACHE_FILE_PATH, "w") as file:
|
|
json.dump(cached_point_arrays, file)
|
|
|
|
|
|
def generate_point_array_from_image(image):
|
|
image.load()
|
|
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. 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 found. Generating now.")
|
|
width, height = image.size
|
|
pixel_array = []
|
|
|
|
#image.save(LOADED_IMAGE_PATH)
|
|
|
|
for y in range(height):
|
|
for x in range(width):
|
|
pixel = image.getpixel((x, y))
|
|
if sum(pixel) > 15: # any pixel which total color value is greater than 15.
|
|
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]:
|
|
result = fastproot.interpolate_point_pairs(pairs, percentage)
|
|
res = [Point2D(p[0], p[1], p[2]) for p in result]
|
|
return res
|
|
|
|
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)
|
|
if not interpolated_point in interpolated_points:
|
|
interpolated_points.append(interpolated_point)
|
|
return interpolated_points
|
|
|
|
|
|
def pair_points(points1: list[Point2D], points2: list[Point2D]) -> list[tuple[Point2D, Point2D]]:
|
|
# Update the size of the point arrays
|
|
size1 = len(points1)
|
|
size2 = len(points2)
|
|
|
|
# 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
|
|
|
|
row_ind, col_ind = (fastproot.solve(points1, points2))
|
|
|
|
# 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 _pair_points(points1: list[Point2D], points2: list[Point2D]) -> list[tuple[Point2D, Point2D]]:
|
|
# Update the size of the point arrays
|
|
size1 = len(points1)
|
|
size2 = len(points2)
|
|
|
|
# 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
|
|
|
|
# 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 |