Split island creation and adding numbers

colorbynumber/main.py

@@ -4,7 +4,7 @@ import numpy as np
 from .config import default_config
 from .simplify_image import simplify_image
 from .gen_islands import GenerateIslands
-from .numbered_islands import create_numbered_islands
+from .numbered_islands import create_islands, add_numbers_to_image
 
 class ColorByNumber:
     def __init__(self, image_path, 
@@ -47,13 +47,23 @@ class ColorByNumber:
         self.island_borders_list = island_borders_list
         self.centroid_coords_list = centroid_coords_list
 
-        numbered_islands = create_numbered_islands(
+        # Create the islands image
+        self.islands_image = create_islands(
             islands = self.island_borders_list, 
-            image_shape = self.image.shape,
-            centroid_coords_list = self.centroid_coords_list,
-            config = self.config
+            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"]
             )
-        self.numbered_islands = numbered_islands
 
         return self.numbered_islands
     

colorbynumber/numbered_islands.py

@@ -38,6 +38,63 @@ def _add_text_to_image(image, text, position, font_size, font_color, font_thickn
         )
 
 
+def add_numbers_to_image(image, 
+                         centroid_coords_list, color_id_list, 
+                         font_size, font_color, font_thickness):
+    """Add numbers to the image.
+    
+    Args:
+        image (np.array): Numpy image.
+        centroid_coords_list (list): A list of centroid coordinates for the islands.
+        color_id_list (list): A list of color ids.
+    Returns:
+        np.array: A new image with the numbers added.
+    """
+    numbered_islands = image.copy()
+    for idx in range(len(centroid_coords_list)):
+        centroid = centroid_coords_list[idx]
+        color_id = color_id_list[idx]
+        if not np.isnan(centroid).any():
+            numbered_islands = _add_text_to_image(
+                image=numbered_islands, 
+                text=str(color_id), 
+                position=centroid,
+                font_size=font_size,
+                font_color=font_color,
+                font_thickness=font_thickness
+            )
+    return numbered_islands
+
+
+def create_islands(islands, image_shape,
+                   padding, border_color, binary = False):
+    """Create a new image with the islands numbered.
+    
+    Args:
+        islands (list): A list of tuples. 
+            Each tuple contains the color id and the coordinates of the island border.
+        image_shape (tuple): The shape of the original image.
+        padding (int): The padding to add to the image.
+        border_color (tuple): The color of the border.
+        binary (bool): If True, the output will be a binary image.
+    """
+    # Create an all white image
+    width, height, channels = image_shape
+    numbered_islands = np.ones((width + padding*2, height + padding*2, channels), 
+                               dtype=np.uint8) * 255
+
+    for color_id, island_coordinates in islands:
+        numbered_islands[island_coordinates] = border_color
+
+    if binary:
+        # Convert numbered_islands to binary using openCV
+        numbered_islands = cv2.cvtColor(numbered_islands, cv2.COLOR_BGR2GRAY)
+        _, numbered_islands = cv2.threshold(numbered_islands, 127, 255, cv2.THRESH_BINARY)
+        return numbered_islands
+
+    return numbered_islands
+
+
 def create_numbered_islands(islands, image_shape, 
                             centroid_coords_list = None,
                             config = default_config, 
@@ -60,33 +117,27 @@ def create_numbered_islands(islands, image_shape,
     font_color = config["font_color"]
     font_thickness = config["font_thickness"]
 
-    # Create an all white image
-    width, height, channels = image_shape
-    numbered_islands = np.ones((width + padding*2, height + padding*2, channels), 
-                               dtype=np.uint8) * 255
-
-    for idx, (color_id, island_coordinates) in enumerate(islands):
-        numbered_islands[island_coordinates] = border_color
-        
-        # Add the number to the centroid of the island
-        if show_numbers:
-            if centroid_coords_list:
-                centroid = centroid_coords_list[idx]
-            else:
-                centroid = _get_centroid(island_coordinates)
-            if not np.isnan(centroid).any():
-                numbered_islands = _add_text_to_image(
-                    image=numbered_islands, 
-                    text=str(color_id), 
-                    position=centroid,
-                    font_size=font_size,
-                    font_color=font_color,
-                    font_thickness=font_thickness
-                )
-    if binary:
-        # Convert numbered_islands to binary using openCV
-        numbered_islands = cv2.cvtColor(numbered_islands, cv2.COLOR_BGR2GRAY)
-        _, numbered_islands = cv2.threshold(numbered_islands, 127, 255, cv2.THRESH_BINARY)
-        return numbered_islands
-
+    islands_image = create_islands(
+        islands, 
+        image_shape, 
+        padding, 
+        border_color, 
+        binary
+        )
+    
+    if not show_numbers:
+        return islands_image
+    
+    if not centroid_coords_list:
+        centroid_coords_list = [
+            _get_centroid(island_coordinates) for _, island_coordinates in islands
+            ]
+    numbered_islands = add_numbers_to_image(
+        islands_image, 
+        centroid_coords_list, 
+        [color_id for color_id, _ in islands], 
+        font_size, 
+        font_color, 
+        font_thickness
+        )
     return numbered_islands