Spaces:
Running
Running
| 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 | |