"""
.. _module_ColorDetect:
Module ColorDetect
==================
Defines ColorDetect class
For example:
>>> from colordetect import ColorDetect
>>> user_image = ColorDetect("<path_to_image>")
# where color_count is the target most dominant colors to be found. Default set to 5
>>> colors = user_image.get_color_count(color_count=5)
>>> colors
# alternatively, save these RGB values to the image
>>> user_image.write_color_count()
>>> user_image.save_image("<storage_path>","<image_file_name>")
# Image processed and saved successfully
"""
import logging
import tempfile
from pathlib import Path
from urllib.request import urlopen
import cv2
import matplotlib.colors as mcolors
import numpy as np
import webcolors
from sklearn.cluster import KMeans
from . import col_share
LOGGER = logging.getLogger(__name__)
[docs]class ColorDetect:
"""
Detect and recognize the number of colors in an image
"""
def __init__(self, image, resize_h: int = None):
"""Create ColorDetect object by providing an image"""
# check type of data being passed
if isinstance(image, np.ndarray):
self.image = image
elif isinstance(image, str):
if col_share.is_url(image) is True:
img_data = urlopen(image).read()
dir_name = tempfile.mkdtemp()
dir_path = Path(dir_name)
# image saved in a temp dir as img.jpg to be read by colordetect
dist_file = str(dir_path / "img.jpg")
f = open(dist_file, "wb")
f.write(img_data)
f.close()
self.image = cv2.imread(dist_file)
else:
self.image = cv2.imread(image)
else:
raise TypeError(
"The image parameter accepts a numpy array , string file path or string file url argument only"
)
if resize_h is not None:
h0, w0, _ = self.image.shape
h1 = resize_h
w1 = int(w0 * h1 / h0)
self.image = cv2.resize(self.image, (w1, h1))
self.image_original = self.image.copy()
self.color_description = {}
[docs] def get_segmented_image(
self,
lower_bound: tuple,
upper_bound: tuple,
erode_iterations: int = 3,
dilate_iterations: int = 3,
use_grab_cut: bool = True,
gc_iterations: int = 3,
) -> tuple:
"""
.. _get_segmented_image:
get_segmented_image
---------------
Get image masks from an image
Parameters
----------
lower_bound: tuple
A lower color range from which to look from
upper_bound: tuple
The higher RGB color range from which to look from
erode_iterations: int
The number of times to perform erosion of the image
dilate_iterations: int
The number of times dilation is applied.
use_grab_cut: bool
A boolean indicating whether grabCut will be applied to the image. This is True by default.
gc_iterations: int
Number of iterations the algorithm should make before returning the result
:return: output_image, gray, segmented, mask
"""
if not self._validate_rgb(lower_bound):
raise TypeError(
f"lower_bound has to be a tuple of integers. Provided {type(lower_bound)} "
)
if not self._validate_rgb(upper_bound):
raise TypeError(
f"upper_bound has to be a tuple of integers. Provided {type(upper_bound)} "
)
if type(erode_iterations) != int:
raise TypeError(
f"erode_iterations has to be an integer. Provided {type(erode_iterations)} "
)
if type(dilate_iterations) != int:
raise TypeError(
f"dilate_iterations has to be an integer. Provided {type(dilate_iterations)} "
)
if type(gc_iterations) != int:
raise TypeError(
f"gc_iterations has to be a an integer. Provided {type(gc_iterations)} "
)
if type(use_grab_cut) != bool:
raise TypeError(
f"use_grab_cut has to be a boolean. Provided {type(use_grab_cut)} "
)
gray = cv2.cvtColor(self.image_original, cv2.COLOR_BGR2GRAY)
output_image = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
img2 = self.image_original.copy()
img2 = cv2.GaussianBlur(img2, (11, 11), 0)
img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2HSV)
mask = cv2.inRange(img2, lower_bound, upper_bound)
mask = cv2.erode(mask, None, iterations=erode_iterations)
mask = cv2.dilate(mask, None, iterations=dilate_iterations)
if use_grab_cut:
mask[mask == 0] = cv2.GC_BGD
mask[mask > 0] = cv2.GC_PR_FGD
fg_model = np.zeros((1, 65), dtype="float")
bg_model = np.zeros((1, 65), dtype="float")
mask, bg_model, fg_model = cv2.grabCut(
self.image_original,
mask,
None,
fg_model,
bg_model,
iterCount=gc_iterations,
mode=cv2.GC_INIT_WITH_MASK,
)
mask = np.where((mask == cv2.GC_BGD) | (mask == cv2.GC_PR_BGD), 0, 1)
mask = (mask * 255).astype("uint8")
segmented = cv2.bitwise_and(self.image_original, self.image_original, mask=mask)
for i in range(len(mask)):
for j in range(len(mask[i])):
if mask[i][j] != 0:
output_image[i][j] = self.image_original[i][j]
return output_image, gray, segmented, mask
[docs] def get_color_count(
self, color_count: int = 5, color_format: str = "human_readable"
) -> dict:
"""
.. _get_color_count:
get_color_count
---------------
Count the number of different colors
Parameters
----------
color_count: int
The number of most dominant colors to be obtained from the image
color_format:str
The format to return the color in.
Options
* hsv - (60°,100%,100%)
* rgb - rgb(255, 255, 0) for yellow
* hex - #FFFF00 for yellow
* human_readable - yellow for yellow
:return: color description
"""
if type(color_count) != int:
raise TypeError(
f"color_count has to be an integer. Provided {type(color_count)} "
)
# convert image from BGR to RGB for better accuracy
rgb = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
reshape = rgb.reshape((rgb.shape[0] * rgb.shape[1], 3))
cluster = KMeans(n_clusters=color_count).fit(reshape)
unique_colors = self._find_unique_colors(cluster, cluster.cluster_centers_)
color_format_options = ["rgb", "hex", "hsv", "human_readable"]
if color_format not in color_format_options:
raise ValueError(f"Invalid color format: {color_format}")
# round up figures
for percentage, v in unique_colors.items():
rgb_value = list(np.around(v))
if color_format != "rgb":
color_value = self._format_color(v, color_format)
self.color_description[color_value] = round(percentage, 2)
else:
self.color_description[str(rgb_value)] = round(percentage, 2)
return self.color_description
def _format_color(self, rgb_value, color_format: str):
"""
Get the correct color format as specified
:return:
"""
if color_format == "hsv":
# list(np.around(v))
return str(mcolors.rgb_to_hsv(rgb_value).tolist())
elif color_format == "hex":
rgb_value = np.divide(rgb_value, 255) # give a scale from 0-1
return mcolors.to_hex(rgb_value)
elif color_format == "human_readable":
r0, g0, b0 = int(rgb_value[0]), int(rgb_value[1]), int(rgb_value[2])
try:
nearest = webcolors.rgb_to_name((r0, g0, b0))
except ValueError: # Calculate distances between rgb value and CSS3 rgb colours to determine the closest
distances = {}
for k, v in webcolors.CSS3_HEX_TO_NAMES.items():
r1, g1, b1 = webcolors.hex_to_rgb(k)
distances[
((r0 - r1) ** 2 + (g0 - g1) ** 2 + (b0 - b1) ** 2)
] = v # Ignore sqrt as it has no significant effect
nearest = distances[min(distances.keys())]
return nearest
def _find_unique_colors(self, cluster, centroids) -> dict:
# Get the number of different clusters, create histogram, and normalize
labels = np.arange(0, len(np.unique(cluster.labels_)) + 1)
(hist, _) = np.histogram(cluster.labels_, bins=labels)
hist = hist.astype("float")
hist /= hist.sum()
# iterate through each cluster's color and percentage
colors = sorted(
[((percent * 100), color) for (percent, color) in zip(hist, centroids)]
)
for (percent, color) in colors:
color.astype("uint8").tolist()
return dict(colors)
[docs] def write_color_count(
self,
left_margin: int = 10,
top_margin: int = 20,
font: int = cv2.FONT_HERSHEY_SIMPLEX,
font_color: tuple = (0, 0, 0),
font_scale: float = 1.0,
font_thickness: float = 1,
line_type: int = 1,
save: bool = False,
):
"""
.. _write_color_count:
write_color_count
-----------------
Write the number of colors found to the image
Parameters
----------
left_margin: int
Text spacing from the left
top_margin: int
Text spacing from the top
font: int
Font to use in text. Look up acceptable values from python-opencv
font_color:
RGB tuple of text font color
font_scale:
Size of the text to be written
font_thickness:
Thickness of the text
line_type: int = 1,
"""
if not self.color_description:
raise AttributeError(
"No color description found on this object. Perform get_color_count() first."
)
if not self._validate_rgb(font_color):
raise TypeError(
f"font_color has to be a tuple of integers. Provided {font_color} "
)
for k, v in self.color_description.items():
color_values = str(v) + "% :" + k
(text_width, text_height), baseline = cv2.getTextSize(
color_values, font, font_scale, font_thickness
)
# change to BGR color format from RGB tuple
bgr_color_format = (font_color[2], font_color[1], font_color[0])
self.write_text(
text=color_values,
left_margin=left_margin,
top_margin=top_margin,
font=font,
font_color=bgr_color_format,
font_scale=font_scale,
font_thickness=font_thickness,
line_type=line_type,
)
top_margin += text_height
if save:
self.save_image()
[docs] def write_text(
self,
text: str = "",
left_margin: int = 10,
top_margin: int = 20,
font: int = cv2.FONT_HERSHEY_SIMPLEX,
font_color: tuple = (0, 0, 0),
font_scale: float = 1.0,
font_thickness: float = 1.0,
line_type: int = 1,
line_spacing: int = 0,
):
"""
.. _write_text:
write_text
----------
Write text onto an image
Parameters
----------
text: str
The text to be written onto the image
line_spacing:int
The spacing between lines
left_margin: int
Text spacing from the left
top_margin: int
Text spacing from the top
font: int
Font to use in text. Look up acceptable values from python-opencv
font_color:
RGB tuple of text font color
font_scale:
Size of the text to be written
font_thickness: float = 1.0
Thickness of the text
line_type: int = 1,
Space betweeen the lines
:return:
"""
if type(text) != str:
raise TypeError(
f"text should be a string.Provided {text} of type {type(text)}"
)
if text == "":
raise IOError("text should not be empty")
if not self._validate_rgb(font_color):
raise TypeError(
f"font_color has to be a tuple of integers. Provided {font_color} "
)
# change to BGR color format from RGB tuple
font_color = (font_color[2], font_color[1], font_color[0])
cv2.putText(
self.image,
text,
(left_margin, top_margin),
font,
font_scale,
font_color,
font_thickness,
line_type,
line_spacing,
)
[docs] def save_image(self, location: str = ".", file_name: str = "out.jpg"):
"""
.. _save_image:
save_image
----------------
Save the resultant image file to the local directory
Parameters
----------
location: str
The file location of the image
file_name:str
The name of the new image
"""
image_folder = Path(location)
if not image_folder.exists():
raise NotADirectoryError("The storage folder does not exist.")
if type(file_name) != str:
raise TypeError(f"file_name should be a string.Provided {type(file_name)}")
image_to_save = image_folder / file_name
# Save image
cv2.imwrite(str(image_to_save), self.image)
LOGGER.info("Image processed and saved successfully")
def _validate_rgb(self, rgb_tuple: tuple) -> bool:
"""
Validate whether a tuple passed is a valid RGB
Parameters
----------
rgb_tuple: tuple
An RGB tuple color.
:return:
"""
tuple_made_of_three = isinstance(rgb_tuple, tuple) and len(rgb_tuple) == 3
tuple_has_integers_only = (
isinstance(rgb_tuple[0], int)
and isinstance(rgb_tuple[1], int)
and isinstance(rgb_tuple[2], int)
)
color_range = []
for color in rgb_tuple:
if color in range(0, 256):
color_range.append(True)
else:
color_range.append(False)
invalid_color_range = False in color_range
return (
tuple_made_of_three and tuple_has_integers_only and not invalid_color_range
)