import cv2 import numpy as np from collections import defaultdict from gradio import Interface, components from PIL import Image from scipy.spatial import cKDTree # Define the Lego color palette lego_colors = [ (190, 148, 116), (86, 142, 186), (207, 220, 146), (160, 57, 140), (214, 135, 64), (130, 151, 150), (120, 72, 65), (121, 190, 211), (25, 31, 56), (195, 77, 151), (72, 132, 51), (134, 75, 44), (64, 37, 81), (169, 181, 178), (219, 159, 68), (162, 46, 49), (51, 27, 39), (229, 194, 114), (121, 50, 123), (27, 51, 45), (187, 120, 47), (230, 214, 180), (167, 221, 219), (22, 29, 40), (171, 119, 88), (167, 204, 90), (235, 237, 233), (76, 43, 50), (40, 57, 94), (9, 10, 20), (213, 182, 149), (17, 20, 31), (115, 34, 56), (58, 74, 80), (64, 93, 139), (32, 18, 61), (94, 44, 44), (64, 28, 49), (31, 28, 57), (36, 20, 39), (89, 114, 119), (33, 46, 55), (118, 169, 69), (203, 86, 62), (39, 87, 46), (222, 213, 254) ] def closest_color(color, palette): """Find the closest color in the palette to the given color.""" tree = cKDTree(palette) _, index = tree.query(color) return palette[index] def legofy_image(image, palette, num_bricks_horizontal, num_bricks_vertical, brick_size): """Legofy the image using the provided palette and brick size.""" # Convert image to PIL format image_pil = Image.fromarray(image.astype('uint8'), 'RGB') # Get original image size original_width, original_height = image_pil.size # Calculate the width and height of the legofied image legofied_width = num_bricks_horizontal * brick_size legofied_height = num_bricks_vertical * brick_size # Resize the image to fit the legofied dimensions image_pil = image_pil.resize((legofied_width, legofied_height)) # Create a blank legofied image legofied_image = Image.new('RGB', (legofied_width, legofied_height), (255, 255, 255)) # Iterate through each brick and replace with the closest lego color for y in range(0, legofied_height, brick_size): for x in range(0, legofied_width, brick_size): # Get the most frequent color of the brick brick_colors = [image_pil.getpixel((x + i, y + j)) for i in range(brick_size) for j in range(brick_size)] most_frequent_color = max(set(brick_colors), key=brick_colors.count) # Find the closest lego color closest_lego_color = closest_color(most_frequent_color, palette) # Fill the legofied image with the closest lego color for i in range(brick_size): for j in range(brick_size): legofied_image.putpixel((x + i, y + j), closest_lego_color) return legofied_image def lego_interface(image, num_bricks_horizontal, num_bricks_vertical, brick_size): if image is None: return None, "No image provided.", "N/A" try: # Perform image processing annotated_image, color_info, total_bricks = process_image(image, num_bricks_horizontal, num_bricks_vertical, brick_size) # Convert annotated image to numpy array annotated_image = np.array(annotated_image) # Convert color info to string output_text_value = '\n'.join([f"Color: {info['Color']}, ID: {info['ID']}, Usage: {info['Usage']}" for info in color_info]) return annotated_image, output_text_value, str(total_bricks) except Exception as e: return None, str(e), "N/A" def process_image(image, num_bricks_horizontal, num_bricks_vertical, brick_size): # Get the dimensions of the image height, width, _ = image.shape # Ensure that the image dimensions are divisible by the brick size if brick_size == 0: raise ValueError("Brick size cannot be zero.") num_rows = height // brick_size num_cols = width // brick_size if num_rows == 0 or num_cols == 0: raise ValueError("Invalid brick size or image dimensions.") image = image[:num_rows * brick_size, :num_cols * brick_size] # Trim to ensure complete brick regions # Initialize dictionaries to store color IDs and corresponding colors color_ids = {} color_counts = defaultdict(int) # Iterate through each brick region for row in range(num_rows): for col in range(num_cols): # Define the region for the current brick start_row = row * brick_size end_row = (row + 1) * brick_size start_col = col * brick_size end_col = (col + 1) * brick_size # Extract the region from the image brick_region = image[start_row:end_row, start_col:end_col] # Calculate the mean color of the brick region mean_color = tuple(np.mean(brick_region, axis=(0, 1)).astype(int)) # Increment the count for the current color color_counts[mean_color] += 1 # If the color is not already assigned a color ID, assign one if mean_color not in color_ids: color_ids[mean_color] = len(color_ids) + 1 # Overlay the color IDs on the bricks annotated_image = image.copy() for row in range(num_rows): for col in range(num_cols): # Define the position to place the color ID text_position = (col * brick_size + 5, (row + 1) * brick_size - 5) # Get the color of the current brick brick_region = image[row * brick_size:(row + 1) * brick_size, col * brick_size:(col + 1) * brick_size] mean_color = tuple(np.mean(brick_region, axis=(0, 1)).astype(int)) # Get the color ID for the current brick color_id = color_ids[mean_color] # Convert the color ID to a string color_id_str = str(color_id) # Draw the color ID on the annotated image # cv2.putText(annotated_image, color_id_str, text_position, cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2, cv2.LINE_AA) # Convert the annotated image back to Gradio format annotated_image = annotated_image.astype(np.uint8) # Assign unique IDs to each unique color and display the usage count color_info = [] total_bricks = 0 for color, usage in color_counts.items(): color_id = color_ids[color] color_info.append({"Color": color, "ID": color_id, "Usage": usage}) total_bricks += usage return annotated_image, color_info, total_bricks # Define the Gradio interface image_input = components.Image(type="numpy", label="Upload your LEGO image") num_bricks_horizontal = components.Number(default=100, label="Number of Bricks Horizontal") num_bricks_vertical = components.Number(default=100, label="Number of Bricks Vertical") brick_size = components.Number(default=30, label="Brick Size") output_image = components.Image(type="numpy", label="Annotated Image with Color IDs") output_text = components.Textbox(label="Color Information") output_total_bricks = components.Textbox(label="Total Bricks Count") # Launch the interface iface = Interface( fn=lego_interface, inputs=[image_input, num_bricks_horizontal, num_bricks_vertical, brick_size], outputs=[output_image, output_text, output_total_bricks], title="LEGO Color Identifier" ) iface.launch()