Spaces:
Running
Running
File size: 3,293 Bytes
a6258d6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import numpy as np
from scipy import ndimage
from sklearn.cluster import MiniBatchKMeans
def vectorized_edge_preserving_quantization(
voxel_matrix, lego_colors, n_initial_clusters=50
):
"""
Vectorized version of edge-preserving color quantization
Parameters:
voxel_matrix: numpy array of shape (x, y, z, 3) containing RGB values
lego_colors: list of [R, G, B] values for target LEGO colors
n_initial_clusters: number of initial color clusters before mapping to LEGO colors
"""
lego_colors = np.array(lego_colors)
shape = voxel_matrix.shape
# Reshape to 2D array of pixels
pixels = voxel_matrix.reshape(-1, 3)
# Step 1: Initial color clustering using K-means
kmeans = MiniBatchKMeans(
n_clusters=n_initial_clusters, batch_size=1000, random_state=42
)
labels = kmeans.fit_predict(pixels)
cluster_centers = kmeans.cluster_centers_
# Step 2: Create 3D gradient magnitude
gradients = np.zeros(shape[:3])
# Compute gradients along each axis
for axis in range(3):
# Forward difference
forward = np.roll(voxel_matrix, -1, axis=axis)
# Compute color differences
diff = np.sqrt(np.sum((forward - voxel_matrix) ** 2, axis=-1))
# Set boundary differences to 0
slice_idx = [slice(None)] * 3
slice_idx[axis] = -1
diff[tuple(slice_idx)] = 0
gradients += diff
# Step 3: Segment using watershed algorithm
# Reshape labels back to 3D
labels_3d = labels.reshape(shape[:3])
# Find local minima in gradient magnitude
markers = ndimage.label(ndimage.minimum_filter(gradients, size=3) == gradients)[0]
# Apply watershed segmentation
segments = ndimage.watershed_ift(gradients.astype(np.uint8), markers)
# Step 4: Map segments to LEGO colors
# Get mean color for each segment
segment_colors = ndimage.mean(
voxel_matrix.reshape(-1, 3),
labels=segments.ravel(),
index=np.arange(segments.max() + 1),
)
# Find nearest LEGO color for each segment color§
def find_nearest_lego_colors(colors):
# Reshape inputs for broadcasting
colors = colors[:, np.newaxis, :]
lego_colors_r = lego_colors[np.newaxis, :, :]
# Compute distances to all LEGO colors at once
distances = np.sqrt(np.sum((colors - lego_colors_r) ** 2, axis=2))
# Find index of minimum distance for each color
nearest_indices = np.argmin(distances, axis=1)
return lego_colors[nearest_indices]
segment_lego_colors = find_nearest_lego_colors(segment_colors)
# Create output array
result = np.zeros_like(voxel_matrix)
# Map segments to final colors
for i, color in enumerate(segment_lego_colors):
mask = segments == i
result[mask] = color
return result
def analyze_quantization(original, quantized):
"""
Analyze the results of quantization
"""
original_colors = np.unique(original.reshape(-1, 3), axis=0)
quantized_colors = np.unique(quantized.reshape(-1, 3), axis=0)
stats = {
"original_colors": len(original_colors),
"quantized_colors": len(quantized_colors),
"reduction_ratio": len(quantized_colors) / len(original_colors),
}
return stats
|