Made changes to denoising

colorbynumber/config.py

@@ -1,8 +1,17 @@
 default_config = {
     # If True, the image will be denoised after simplification.
     "denoise": True,
-
-    # Higher values will result in more aggressive denoising.
+    
+    # Determines the order of denoising.
+    # Options: "before_simplify", "after_simplify"
+    "denoise_order": "before_simplify",
+
+    # Type of denoising to be used.
+    # Options: "fastNlMeansDenoisingColored", "gaussianBlur", "blur"
+    "denoise_type": "gaussianBlur",
+    # Size of the kernel used for gaussian blur.
+    "blur_size": 51,
+    # h parameter for fastNlMeansDenoisingColored.
     "denoise_h": 200,
 
     # Padding around the borders of the image.
@@ -26,7 +35,7 @@ default_config = {
     "border_color": (181, 181, 181),
 
     # Font for the numbers shown in color islands.
-    "font_size": 0.3,
+    "font_size": 1,
     "font_color": (140, 140, 140),
-    "font_thickness": 1,
+    "font_thickness": 2,
 }

colorbynumber/main.py

@@ -7,27 +7,38 @@ from .gen_islands import GenerateIslands
 from .numbered_islands import create_numbered_islands
 
 class ColorByNumber:
-    def __init__(self, image_path, color_list, config = default_config):
+    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.color_list = color_list
         self.config = config
+        self.color_list = color_list
+        self.num_colors = num_colors
 
-    def create_color_by_number(self):
         image = cv.imread(self.image_path)
         image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
         self.image = image
 
-        simplified_image, indices_color_choices = simplify_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)

colorbynumber/simplify_image.py

@@ -31,19 +31,51 @@ def _choose_closest_colors(image, color_list):
 
     return simplified_image, indices_color_choices
 
-def _denoise_image(image, h):
-    denoised_image = cv.fastNlMeansDenoisingColored(
-        src = image.astype(np.uint8),
-        dst = None,
-        h = h,
-        hColor = h,
-        templateWindowSize = 7,
-        searchWindowSize = 21
-    )
+def _kmeans_simplify_image(image, num_colors):
+    Z = image.reshape((-1,3))
+ 
+    # convert to np.float32
+    Z = np.float32(Z)
+    
+    # define criteria, number of clusters(K) and apply kmeans()
+    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
+    K = num_colors
+    ret,label,center=cv.kmeans(Z,K,None,criteria,10,cv.KMEANS_RANDOM_CENTERS)
+    
+    # Now convert back into uint8, and make original image
+    center = np.uint8(center)
+    res = center[label.flatten()]
+    res = res.reshape((image.shape))
+
+    simplified_image = res
+    indices_color_choices = label.reshape((image.shape[:2])) + 1
+    color_list = center
+
+    return simplified_image, indices_color_choices, color_list
+
+def _denoise_image(image, h, denoise_type, blur_size = None):
+    if denoise_type == "fastNlMeansDenoisingColored":
+        denoised_image = cv.fastNlMeansDenoisingColored(
+            src = image.astype(np.uint8),
+            dst = None,
+            h = h,
+            hColor = h,
+            templateWindowSize = 7,
+            searchWindowSize = 21
+        )
+    elif denoise_type == "gaussianBlur":
+        kernel = (blur_size, blur_size)
+        denoised_image = cv.GaussianBlur(image, kernel, 0)
+    
+    elif denoise_type == "blur":
+        kernel = (blur_size, blur_size)
+        denoised_image = cv.blur(image, kernel)
+
     return denoised_image
 
 def simplify_image(image, 
-                   color_list, 
+                   color_list = None,
+                   num_colors = None, 
                    config = default_config,
                    ):
     """
@@ -57,16 +89,36 @@ def simplify_image(image,
     Returns:
       A copy of the image with all colors replaced with the closest color in the list.
     """
-    denoise = config["denoise"]
-    denoise_h = config["denoise_h"]
+    if config["denoise"] and (config["denoise_order"] == "before_simplify"):
+        image = _denoise_image(
+            image=image, 
+            h = config["denoise_h"], 
+            denoise_type = config["denoise_type"], 
+            blur_size = config["blur_size"],
+            )
+
+    if color_list is None:
+        # Use kmeans to simplify the image to the specified number of colors.
+        simplified_image, indices_color_choices, color_list = _kmeans_simplify_image(image, num_colors)
+
+    else:
+        simplified_image, indices_color_choices = _choose_closest_colors(image, color_list)
     
-    simplified_image, indices_color_choices = _choose_closest_colors(image, color_list)
-    if denoise:
+
+    if config["denoise"] and (config["denoise_order"] == "after_simplify"):
         # Denoising after simplifying image as denoising original image need not reduce the 
         # number of colors islands and also removed some key features from the original image.
-        simplified_image = _denoise_image(simplified_image, denoise_h)
+        simplified_image = _denoise_image(
+            simplified_image, 
+            h = config["denoise_h"], 
+            denoise_type = config["denoise_type"], 
+            blur_size = config["blur_size"],
+            )
     
         # Simplying image again as denoising may have introduced new colors.
-        simplified_image, indices_color_choices = _choose_closest_colors(simplified_image, color_list)
+        simplified_image, indices_color_choices = _choose_closest_colors(
+            simplified_image, 
+            color_list
+            )
 
-    return simplified_image, indices_color_choices
+    return simplified_image, indices_color_choices, color_list