Spaces:
Running
Running
| import cv2 as cv | |
| import numpy as np | |
| from .config import default_config | |
| from .simplify_image import simplify_image, downsample_image | |
| from .gen_islands import GenerateIslands | |
| from .numbered_islands import create_islands, add_numbers_to_image | |
| class ColorByNumber: | |
| def __init__(self, image_path, | |
| color_list = None, num_colors = None, | |
| config = default_config): | |
| """ | |
| Args: | |
| image_path: Path to the image file. | |
| color_list: List of colors in (R, G, B) format. | |
| config: Dictionary of configuration parameters (optional). | |
| """ | |
| assert color_list is not None or num_colors is not None, \ | |
| "Either color_list or num_colors must be provided." | |
| self.image_path = image_path | |
| self.config = config | |
| self.color_list = color_list | |
| self.num_colors = num_colors | |
| image = cv.imread(self.image_path) | |
| image = cv.cvtColor(image, cv.COLOR_BGR2RGB) | |
| image = downsample_image(image) | |
| self.image = image | |
| def create_color_by_number(self): | |
| simplified_image, indices_color_choices, color_list = simplify_image( | |
| image=self.image, | |
| color_list=self.color_list, | |
| num_colors=self.num_colors, | |
| config=self.config | |
| ) | |
| # Assigning the color_list to the one returned by simplify_image | |
| # because if it was initially None, it would have been assigned a value. | |
| self.color_list = color_list | |
| self.simplified_image = simplified_image | |
| generate_islands_obj = GenerateIslands(indices_color_choices) | |
| island_borders_list, centroid_coords_list = generate_islands_obj.get_islands(config=self.config) | |
| self.generate_islands_obj = generate_islands_obj | |
| self.island_borders_list = island_borders_list | |
| self.centroid_coords_list = centroid_coords_list | |
| # Create the islands image | |
| self.islands_image = create_islands( | |
| islands = self.island_borders_list, | |
| image_shape = self.image.shape, | |
| padding = self.config["border_padding"], | |
| border_color = self.config["border_color"] | |
| ) | |
| # Add numbers to the islands image | |
| self.numbered_islands = add_numbers_to_image( | |
| image=self.islands_image, | |
| centroid_coords_list=self.centroid_coords_list, | |
| color_id_list=[color_id for color_id, _ in self.island_borders_list], | |
| font_size=self.config["font_size"], | |
| font_color=self.config["font_color"], | |
| font_thickness=self.config["font_thickness"] | |
| ) | |
| return self.numbered_islands | |
| def generate_color_legend(self, | |
| cols=7, | |
| rows=None, | |
| square_size=100, | |
| margin=10, | |
| gap_horizontal=5, gap_vertical=30, | |
| font=cv.FONT_HERSHEY_SIMPLEX, | |
| font_size=1, | |
| border_color=(0, 0, 0) | |
| ): | |
| """ | |
| Generates a grid of colored squares with labels below them. | |
| Args: | |
| cols: Number of columns in the grid. | |
| rows: Number of rows in the grid. | |
| square_size: Size of each square in the grid. | |
| margin: Margin around the grid. | |
| gap_horizontal: Horizontal gap between squares. | |
| gap_vertical: Vertical gap between squares. | |
| font: Font for the labels. | |
| font_size: Font size for the labels. | |
| border_color: Color of the border around each square. | |
| """ | |
| # Calculate grid dimensions if not provided | |
| if rows is None and cols is None: | |
| num_colors = len(self.color_list) | |
| rows = cols = int(np.sqrt(num_colors)) + 1 | |
| elif rows is None: | |
| cols = min(cols, len(self.color_list)) | |
| rows = int(np.ceil(len(self.color_list) / cols)) | |
| # Calculate total width and height based on margins, gaps, and squares | |
| total_width = 2 * margin + (cols + 1) * square_size + (cols - 1) * gap_horizontal | |
| total_height = 2 * margin + (rows + 1) * square_size + (rows - 1) * gap_vertical | |
| # Create a white image | |
| image = np.ones((total_height, total_width, 3), dtype=np.uint8) * 255 | |
| # Fill squares with colors | |
| for i, color in enumerate(self.color_list): | |
| row = i // cols | |
| col = i % cols | |
| start_col = margin + col * (square_size + gap_horizontal) | |
| end_col = start_col + square_size | |
| start_row = margin + row * (square_size + gap_vertical) | |
| end_row = start_row + square_size | |
| # Fill square with color | |
| image[start_row:end_row, start_col:end_col] = color | |
| # Draw border around that color | |
| image[start_row, start_col:end_col] = border_color # Top Border | |
| image[end_row, start_col:end_col] = border_color # Bottom Border | |
| image[start_row:end_row, start_col] = border_color # Left Border | |
| image[start_row:end_row, end_col] = border_color # Right Border | |
| # Draw text label below the square | |
| text = str(i + 1) | |
| text_size, _ = cv.getTextSize(text, font, font_size, 1) | |
| text_row = (end_row + text_size[1]) + 5 | |
| text_col = start_col + (square_size // 2) - (text_size[0] // 2) | |
| cv.putText(image, text, (text_col, text_row), font, font_size, (0, 0, 0), 1) | |
| return image | |