import numpy as np
from PIL import Image
from sklearn.decomposition import PCA
from scipy.spatial import cKDTree
from skimage import color
import itertools
import colorsys
# --- CONFIG ---
IMAGE_SIZE = 4096
TOTAL_PIXELS = IMAGE_SIZE * IMAGE_SIZE
# --- STEP 1: Generate E8 root system ---
def generate_e8_roots():
roots = []
for i, j in itertools.combinations(range(8), 2):
for s1 in [-1, 1]:
for s2 in [-1, 1]:
if s1 == s2:
vec = [0.0] * 8
vec[i], vec[j] = s1, s2
roots.append(vec)
for signs in itertools.product([-0.5, 0.5], repeat=8):
if sum(signs) == 0 and signs.count(-0.5) % 2 == 0:
roots.append(list(signs))
return np.array(roots)
# --- STEP 2: Project E8 roots to 2D via PCA ---
e8_roots = generate_e8_roots()
pca = PCA(n_components=2)
projected = pca.fit_transform(e8_roots)
projected -= projected.min(axis=0)
projected /= projected.max(axis=0)
projected *= (IMAGE_SIZE - 1)
e8_centers = projected.astype(np.int32)
NUM_MARKERS = len(e8_centers)
# --- STEP 3: Generate all 24-bit RGB colors and sort them ---
indices = np.arange(TOTAL_PIXELS, dtype=np.uint32)
np.random.seed(42)
np.random.shuffle(indices)
r = (indices >> 16) & 0xFF
g = (indices >> 8) & 0xFF
b = indices & 0xFF
rgb_colors = np.stack([r, g, b], axis=1).astype(np.uint8)
# --- STEP 4: Select 240 brightest colors as marker colors ---
brightness = 0.2126 * r + 0.7152 * g + 0.0722 * b
brightest_indices = np.argsort(-brightness)
marker_indices = brightest_indices[:NUM_MARKERS]
marker_colors = rgb_colors[marker_indices]
# Remove marker colors from pool
mask = np.ones(TOTAL_PIXELS, dtype=bool)
mask[marker_indices] = False
remaining_rgb = rgb_colors[mask]
# Convert to LAB and HSV for hybrid sorting
norm = remaining_rgb / 255.0
lab = color.rgb2lab(norm.reshape(-1, 1, 3)).reshape(-1, 3)
hsv = np.array([colorsys.rgb_to_hsv(*c) for c in norm])
# Sort remaining colors by hybrid key
hybrid_key = hsv[:, 0] * 100000 + lab[:, 0] * 100 + lab[:, 1]
sorted_indices = np.argsort(hybrid_key)
remaining_colors = remaining_rgb[sorted_indices]
# --- STEP 5: Assign each pixel to nearest E8 root ---
print("Computing pixel regions...")
coords = np.indices((IMAGE_SIZE, IMAGE_SIZE)).reshape(2, -1).T.astype(np.float32)
tree = cKDTree(e8_centers)
_, regions = tree.query(coords, k=1)
# --- STEP 6: Radially sort pixels within each region ---
coords_by_region = [[] for _ in range(NUM_MARKERS)]
for i, region in enumerate(regions):
coords_by_region[region].append(coords[i])
ordered_coords = []
for i, coord_list in enumerate(coords_by_region):
coord_list = np.array(coord_list)
center = e8_centers[i]
if len(coord_list) == 0:
continue
deltas = coord_list - center
angles = np.arctan2(deltas[:, 1], deltas[:, 0])
distances = np.linalg.norm(deltas, axis=1)
sort_idx = np.lexsort((distances, angles))
sorted_coords = coord_list[sort_idx]
ordered_coords.append(sorted_coords)
coords_ordered = np.vstack(ordered_coords).astype(np.int32)
# --- STEP 7: Reserve marker coordinates (nearest to centers) ---
marker_coords = []
used = set()
for cx, cy in e8_centers:
best = None
best_dist = float('inf')
for dx in range(-2, 3):
for dy in range(-2, 3):
x = int(np.clip(cx + dx, 0, IMAGE_SIZE - 1))
y = int(np.clip(cy + dy, 0, IMAGE_SIZE - 1))
coord = (x, y)
if coord not in used:
d = np.hypot(dx, dy)
if d < best_dist:
best_dist = d
best = coord
used.add(best)
marker_coords.append(best)
# Remove marker_coords from ordered_coords and insert at front
used_set = set(marker_coords)
coords_filtered = [tuple(x) for x in coords_ordered if tuple(x) not in used_set]
final_coords = np.array(marker_coords + coords_filtered, dtype=np.int32)
# --- STEP 8: Build image ---
print("Mapping colors to pixels...")
image_array = np.zeros((IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8)
all_colors = np.vstack((marker_colors, remaining_colors))
assert len(final_coords) == len(all_colors), "Mismatch in coordinate/color count"
for i, (x, y) in enumerate(final_coords):
image_array[x, y] = all_colors[i]
# --- STEP 9: Save image ---
Image.fromarray(image_array, mode='RGB').save("e8_rgb_exact_4096.png")
print("✅ Saved: e8_rgb_exact_4096.png")
