import numpy as np
from PIL import Image, ImageDraw
import time
import gc
# --- Configuration ---
WIDTH = 4096
HEIGHT = 4096
TOTAL_PIXELS = WIDTH * HEIGHT
if TOTAL_PIXELS != 2**24:
raise ValueError(f"WIDTH x HEIGHT ({WIDTH}x{HEIGHT}={TOTAL_PIXELS}) must be 2^24 ({2**24})")
# --- Region Definitions (Target Colors and Labels) ---
REGION_DEFINITIONS = [
# Chevron (Labels 0-4)
("Black", (0, 0, 0), 0),
("Brown", (96, 51, 22), 1), # Slightly richer brown: (102, 60, 36)
("Light Blue", (91, 206, 250), 2),
("Pink", (241, 188, 232), 3),
("White", (255, 255, 255), 4), # Innermost triangle
# Main Stripes (Labels 5-10)
("Red", (228, 3, 3), 5),
("Orange", (255, 140, 0), 6),
("Yellow", (255, 237, 0), 7),
("Green", (0, 128, 38), 8),
("Blue", (30, 90, 255), 9),
("Purple", (118, 40, 160), 10),
]
NUM_REGIONS = len(REGION_DEFINITIONS)
TARGET_COLORS_RGB_ORDERED = [region[1] for region in REGION_DEFINITIONS]
def generate_all_rgb_colors_shuffled():
print(f"Generating and shuffling all {TOTAL_PIXELS} colors...")
s = time.time()
color_indices = np.arange(TOTAL_PIXELS, dtype=np.uint32)
r = np.right_shift(color_indices, 16) & 0xFF
g = np.right_shift(color_indices, 8) & 0xFF
b = color_indices & 0xFF
all_colors = np.stack((r, g, b), axis=-1).astype(np.uint8)
# Crucial step: Shuffle the colors
np.random.shuffle(all_colors)
print(f"Generated and shuffled all colors in {time.time() - s:.2f} seconds.")
return all_colors
def create_label_image():
print("Creating label image for flag regions...")
s = time.time()
label_img_pil = Image.new('L', (WIDTH, HEIGHT), 0) # 'L' mode for 8-bit pixels
draw = ImageDraw.Draw(label_img_pil)
# --- Main 6 Stripes ---
# Labels 5 (Red) to 10 (Purple)
stripe_height = HEIGHT // 6
for i in range(6):
y_start = i * stripe_height
y_end = (i + 1) * stripe_height if i < 5 else HEIGHT # Ensure last stripe fills to bottom
region_label = REGION_DEFINITIONS[5+i][2] # Get label ID (5 to 10)
draw.rectangle([0, y_start, WIDTH, y_end], fill=region_label)
# --- Chevron ---
# Point of the chevron
X_TIP = int(WIDTH * 0.38) # How far the chevron extends into the flag
Y_TIP = HEIGHT // 2
# Thickness of each chevron band at the hoist (left edge)
# These are not actual "thicknesses" but define the y-coordinates of the base of each triangle
# We define 5 triangles, Black being the largest, White the smallest.
# They share a common tip (X_TIP, Y_TIP).
# Their bases are on the left edge (x=0), or slightly offset for visual appeal.
# We define 5 layers for the chevron. Layer 0 is Black, Layer 4 is White.
# Each layer k is a triangle. We draw Black first, then Brown on top, etc.
# The base of triangle k (for Black, Brown, ..., White) starts at x_base_k
# and extends from y_top_k to y_bottom_k.
# Let's define the triangles from outermost (Black) to innermost (White)
# The x-coordinates of the bases of the triangles.
# Black triangle base is at x=0. White triangle base is most indented.
# Let base_x_coords define the x-position of the *vertical line* forming the base of each triangle.
# For simplicity, let all chevron triangle bases be at x=0, and their y-extents define their "thickness"
num_chevron_bands = 5 # Black, Brown, LBlue, Pink, White
# These values control the "spread" of the chevron bands at the hoist edge (x=0)
# And how much "thinner" each inner band gets relative to the outer one.
# `offsets_from_border` defines how far from (0,0) and (0,H) the base points of each triangle are.
# Black (k=0): base points (0,0) and (0,H). Triangle: (0,0)-(X_TIP,Y_TIP)-(0,H-1)
# Brown (k=1): base points (0, Y_OFFSET_1) and (0, H-Y_OFFSET_1). Triangle: (0,Y_OFFSET_1)-(X_TIP,Y_TIP)-(0,H-1-Y_OFFSET_1)
# ...
# White (k=4)
# y_offsets defines the y-coordinate for the top point of the base and (H - y_offset) for bottom point.
# A simple linear spacing for the y_offsets of the triangle bases at x=0
y_offsets = np.linspace(0, HEIGHT * 0.18, num_chevron_bands).astype(int) # Fine-tune 0.18 for thickness
for k in range(num_chevron_bands): # 0=Black, 1=Brown, ..., 4=White
region_label = REGION_DEFINITIONS[k][2] # Get label ID (0 to 4)
y_top_base = y_offsets[k]
y_bottom_base = HEIGHT - 1 - y_offsets[k]
# Define the polygon for the current chevron triangle layer
# (Base Point 1), (Tip), (Base Point 2)
polygon_points = [
(0, y_top_base),
(X_TIP, Y_TIP),
(0, y_bottom_base)
]
draw.polygon(polygon_points, fill=region_label)
label_image_np = np.array(label_img_pil, dtype=np.uint8)
print(f"Created label image in {time.time() - s:.2f}s.")
# # Optional: Save label image for debugging
# label_img_pil.save("debug_label_image.png")
# unique_labels, counts = np.unique(label_image_np, return_counts=True)
# print("Label counts:")
# for label, count in zip(unique_labels, counts):
# region_name = [r[0] for r in REGION_DEFINITIONS if r[2] == label][0]
# print(f" Label {label} ({region_name}): {count} pixels")
# if np.sum(counts) != TOTAL_PIXELS:
# print(f"WARNING: Total pixels in labels ({np.sum(counts)}) != TOTAL_PIXELS ({TOTAL_PIXELS})")
return label_image_np
def color_distance_sq(color_array_int32, target_color_np_int32):
diff = color_array_int32 - target_color_np_int32
return np.sum(diff**2, axis=1)
def create_flag_image_geometric(all_colors_shuffled, label_image):
final_image_array = np.empty((HEIGHT, WIDTH, 3), dtype=np.uint8)
current_color_idx = 0
# Verify total pixels in label map matches expected
total_labeled_pixels = 0
pixel_counts_per_region = []
for i in range(NUM_REGIONS):
count = np.sum(label_image == i)
pixel_counts_per_region.append(count)
total_labeled_pixels += count
if total_labeled_pixels != TOTAL_PIXELS:
print(f"ERROR: Mismatch in pixel counts. Label map has {total_labeled_pixels}, expected {TOTAL_PIXELS}")
# Distribute shortfall/excess (this is a hack, ideally label map is perfect)
diff = TOTAL_PIXELS - total_labeled_pixels
# Add/subtract from the largest region (e.g., first main stripe, Red, Label 5)
# Find index of label 5 in REGION_DEFINITIONS to map to pixel_counts_per_region
label_to_adjust = 5
pixel_counts_per_region[label_to_adjust] += diff
print(f"Adjusted pixel count for region with label {label_to_adjust} by {diff}")
if np.sum(pixel_counts_per_region) != TOTAL_PIXELS:
raise ValueError("Pixel count adjustment failed.")
for i in range(NUM_REGIONS):
region_name, target_color_rgb, region_label = REGION_DEFINITIONS[i]
target_color_np_int32 = np.array(target_color_rgb, dtype=np.int32)
num_pixels_in_region = pixel_counts_per_region[region_label] # Use pre-calculated counts
if num_pixels_in_region == 0:
print(f"Skipping Region {region_label} ({region_name}): 0 pixels.")
continue
print(f"\nProcessing Region {region_label} ({region_name}): Target {target_color_rgb}, Pixels {num_pixels_in_region}")
s_chunk = time.time()
colors_for_this_region = all_colors_shuffled[current_color_idx : current_color_idx + num_pixels_in_region]
# print(f" Extracted {len(colors_for_this_region)} colors in {time.time() - s_chunk:.2f}s.")
s_sort = time.time()
colors_for_this_region_int32 = colors_for_this_region.astype(np.int32)
distances = color_distance_sq(colors_for_this_region_int32, target_color_np_int32)
sorted_indices = np.argsort(distances)
sorted_colors_for_region = colors_for_this_region[sorted_indices]
# print(f" Sorted colors in {time.time() - s_sort:.2f}s.")
s_fill = time.time()
# Get coordinates for the current region
# Important: The order of pixels returned by np.where matters for how sorted colors are placed.
# This will create a gradient/texture within the region.
region_pixels_y, region_pixels_x = np.where(label_image == region_label)
if len(region_pixels_y) != num_pixels_in_region:
print(f" WARNING: Mismatch! np.where found {len(region_pixels_y)} pixels for label {region_label}, expected {num_pixels_in_region}.")
# This can happen if the adjustment logic for pixel counts is triggered and np.where is not re-run
# For now, we proceed but the image might have uncolored pixels or color reuse if counts differ.
# A robust solution would be to ensure label_image perfectly matches TOTAL_PIXELS.
# Or, if counts differ, take min(len(region_pixels_y), num_pixels_in_region) for assignment.
# However, this would violate the "all colors once" if num_pixels_in_region is smaller,
# or leave gaps if num_pixels_in_region is larger.
# The current setup *should* ensure counts match if label image generation is correct.
# The hacky adjustment is the main risk here.
final_image_array[region_pixels_y, region_pixels_x] = sorted_colors_for_region
# print(f" Filled image data in {time.time() - s_fill:.2f}s.")
current_color_idx += num_pixels_in_region
del colors_for_this_region, colors_for_this_region_int32, distances, sorted_indices, sorted_colors_for_region
del region_pixels_y, region_pixels_x
gc.collect()
if current_color_idx != TOTAL_PIXELS:
print(f"WARNING: Processed {current_color_idx} colors, but expected {TOTAL_PIXELS}.")
return final_image_array
if __name__ == "__main__":
overall_start_time = time.time()
all_rgb_colors_shuffled = generate_all_rgb_colors_shuffled()
label_map = create_label_image()
# Verify sum of pixels in label_map regions
unique_labels, counts = np.unique(label_map, return_counts=True)
print("\nPixel counts per label from generated label_map:")
sum_label_pixels = 0
for label_val, count_val in zip(unique_labels, counts):
region_name = [r[0] for r in REGION_DEFINITIONS if r[2] == label_val][0]
print(f" Label {label_val} ({region_name}): {count_val} pixels")
sum_label_pixels += count_val
if sum_label_pixels != TOTAL_PIXELS:
print(f"CRITICAL WARNING: Total pixels in label map ({sum_label_pixels}) does not match TOTAL_PIXELS ({TOTAL_PIXELS}).")
print("This will lead to errors in color allocation or image generation.")
print("The create_label_image() function needs to ensure all pixels are assigned a label exactly once.")
# For now, the code attempts to adjust, but this is a symptom of an issue in label generation.
else:
print(f"Label map total pixels: {sum_label_pixels}. Matches TOTAL_PIXELS.")
final_image_array = create_flag_image_geometric(all_rgb_colors_shuffled, label_map)
s_save = time.time()
print("\nSaving image...")
img = Image.fromarray(final_image_array, 'RGB')
img.save("inclusive_pride_all_colors_geometric.png")
print(f"Image saved as inclusive_pride_all_colors_geometric.png in {time.time() - s_save:.2f}s.")
total_time = time.time() - overall_start_time
print(f"\nTotal execution time: {total_time:.2f} seconds.")
