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 | |