No public description

PiperOrigin-RevId: 597369513

Author: tensorflower-gardener

official/projects/waste_identification_ml/model_inference/Inference.ipynb

@@ -13,13 +13,26 @@
 # limitations under the License.
 
 """Extract properties from each object mask and detect its color."""
-from typing import Optional, Union
+
+from typing import Annotated, Literal, TypeVar, Union
+
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
+from skimage import color as skimage_color
 import skimage.measure
-from sklearn.cluster import KMeans
+from sklearn import cluster as sklearn_cluster
+from sklearn import neighbors as sklearn_neighbors
 import webcolors
 
+DType = TypeVar('DType', bound=np.generic)
+# Color representation as numpy array of 3 elements of float64
+# Those values could be in different scales like
+# RGB ([0.0,255.0], [0.0,255.0], [0.0 to 255.0])
+# LAB ([0.0,100], [-128,127], [-128,127])
+NColor = Annotated[npt.NDArray[DType], Literal[3]][np.float64]
+
+
 PROPERTIES = [
     'area',
     'bbox',
@@ -31,6 +44,68 @@
     'centroid',
 ]
 
+GENERIC_COLORS = [
+    ('black', '#000000'),
+    ('green', '#008000'),
+    ('green', '#00ff00'),  # lime
+    ('green', '#3cb371'),  # mediumseagreen
+    ('green', '#2E8B57'),  # seagreen
+    ('green', '#8FBC8B'),  # darkseagreen
+    ('green', '#adff2f'),  # olive
+    ('green', '#008080'),  # Teal
+    ('green', '#808000'),
+    ('blue', '#000080'),  # navy
+    ('blue', '#00008b'),  # darkblue
+    ('blue', '#4682b4'),  # steelblue
+    ('blue', '#40E0D0'),  # turquoise
+    ('blue', '#00FFFF'),  # cyan
+    ('blue', '#00ffff'),  # aqua
+    ('blue', '#6495ED'),  # cornflowerBlue
+    ('blue', '#4169E1'),  # royalBlue
+    ('blue', '#87CEFA'),  # lightSkyBlue
+    ('blue', '#4682B4'),  # steelBlue
+    ('blue', '#B0C4DE'),  # lightSteelBlue
+    ('blue', '#87CEEB'),  # skyblue
+    ('blue', '#0000CD'),  # mediumBlue
+    ('blue', '#0000ff'),
+    ('purple', '#800080'),
+    ('purple', '#9370db'),  # mediumpurple
+    ('purple', '#8B008B'),  # darkMagenta
+    ('purple', '#4B0082'),  # indigo
+    ('red', '#ff0000'),
+    ('red', '#B22222'),  # fireBrick
+    ('red', '#DC143C'),  # fireBrick
+    ('red', '#8B0000'),  # crimson
+    ('red', '#CD5C5C'),  # indianred
+    ('red', '#F08080'),  # lightCoral
+    ('red', '#FA8072'),  # salmon
+    ('red', '#E9967A'),  # darkSalmon
+    ('red', '#FFA07A'),  # lightSalmon
+    ('gray', '#c0c0c0'),  # silver,
+    ('white', '#ffffff'),
+    ('white', '#F5F5DC'),  # beige
+    ('white', '#FFFAFA'),  # snow
+    ('white', '#F0F8FF'),  # aliceBlue
+    ('white', '#FFE4E1'),  # mistyRose
+    ('yellow', '#ffff00'),
+    ('yellow', '#ffffe0'),  # lightyellow
+    ('yellow', '#8B8000'),  # darkyellow,
+    ('orange', '#ffa500'),
+    ('orange', '#ff8c00'),  # darkorange
+    ('pink', '#ffc0cb'),
+    ('pink', '#ff00ff'),  # fuchsia
+    ('pink', '#C71585'),  # mediumVioletRed
+    ('pink', '#DB7093'),  # paleVioletRed
+    ('pink', '#FFB6C1'),  # lightPink
+    ('pink', '#FF69B4'),  # hotPink
+    ('pink', '#FF1493'),  # deepPink
+    ('pink', '#BC8F8F'),  # rosybrown
+    ('brown', '#a52a2a'),
+    ('brown', '#8b4513'),  # saddlebrown
+    ('brown', '#f4a460'),  # sandybrown
+    ('brown', '#800000'),  # maroon
+]
+
 
 def extract_properties_and_object_masks(
     final_result: dict[str, np.ndarray],
@@ -114,9 +189,9 @@ def find_dominant_color(
   non_black_pixels = pixels[(pixels > black_threshold).any(axis=1)]
 
   if non_black_pixels.size != 0:
-    kmeans = KMeans(n_clusters=1, n_init=10, random_state=0).fit(
-        non_black_pixels
-    )
+    kmeans = sklearn_cluster.KMeans(
+        n_clusters=1, n_init=10, random_state=0
+    ).fit(non_black_pixels)
     dominant_color = kmeans.cluster_centers_[0].astype(int)
 
   else:
@@ -163,7 +238,7 @@ def est_color(requested_color: tuple[int, int, int]) -> str:
   return min_colors[min(min_colors.keys())]
 
 
-def get_color_name(rgb_color: tuple[int, int, int]) -> Optional[str]:
+def get_color_name(rgb_color: tuple[int, int, int]) -> str | None:
   """Retrieves the name of a given RGB color.
 
   If the RGB color exactly matches one of the CSS3 predefined colors, it returns
@@ -188,3 +263,90 @@ def get_color_name(rgb_color: tuple[int, int, int]) -> Optional[str]:
     return closest_color_name
   else:
     return None
+
+
+def rgb_int_to_lab(rgb_int_color: tuple[int, int, int]) -> NColor:
+  """Convert RGB color to LAB color space.
+
+  Args:
+    rgb_int_color: RGB tuple color e.g. (128,128,128)
+
+  Returns:
+    Numpy array of 3 elements that contains LAB color space.
+  """
+  return skimage_color.rgb2lab(
+      (rgb_int_color[0] / 255, rgb_int_color[1] / 255, rgb_int_color[2] / 255)
+  )
+
+
+def color_distance(
+    a: tuple[int, int, int], b: tuple[int, int, int]
+) -> np.ndarray:
+  """The color distance following the ciede2000 formula.
+
+  See: https://en.wikipedia.org/wiki/Color_difference#CIEDE2000
+
+  Args:
+    a: Color a
+    b: Color b
+
+  Returns:
+    The distance between color a and b
+  """
+  return skimage_color.deltaE_ciede2000(a, b, kC=0.6)
+
+
+def build_color_lab_list(
+    generic_colors: list[tuple[str, str]]
+) -> tuple[npt.NDArray[np.str_], list[NColor]]:
+  """Get Simple colors names and lab values.
+
+  Args:
+    generic_colors: List of colors in this format (color_name, rgb_value in hex)
+      e.g. [ ('black', '#000000'), ('green', '#008000'), ]
+
+  Returns:
+    Numpy array of strings that contains color names
+      ['black', 'green']
+    List of color lab values in the format of Numpy array of 3 elements
+    e.g.
+      [
+        np.array([0., 0., 0.]),
+        np.array([ 46.2276577 , -51.69868348,  49.89707556])
+      ]
+  """
+  names: list[str] = []
+  lab_values = []
+  for color_name, color_hex in generic_colors:
+    names.append(color_name)
+    hex_color = webcolors.hex_to_rgb(color_hex)
+    lab_values.append(rgb_int_to_lab(hex_color))
+  color_names = np.array(names)
+  return color_names, lab_values
+
+
+def get_generic_color_name(
+    rgb_colors: list[tuple[int, int, int]],
+    generic_colors: list[tuple[str, str]] | None = None,
+) -> list[str]:
+  """Retrieves generic names of given RGB colors.
+
+  Estimates the closest matching color name.
+
+  Args:
+    rgb_colors: A list of RGB values for which to retrieve the name.
+    generic_colors: A list of color names and their RGB values in hex.
+
+  Returns:
+    The list of closest color names.
+
+    Example: get_generic_color_name([(255, 0, 0), (0,0,0)])
+    ['red','black']
+  """
+  names, rgb_simple_colors = build_color_lab_list(
+      generic_colors or GENERIC_COLORS
+  )
+  tree = sklearn_neighbors.BallTree(rgb_simple_colors, metric=color_distance)
+  rgb_query = [*map(rgb_int_to_lab, rgb_colors)]
+  _, index = tree.query(rgb_query)
+  return [x[0] for x in names[index]]

official/projects/waste_identification_ml/model_inference/color_and_property_extractor.py

@@ -112,6 +112,16 @@ def test_est_color(self):
 
     self.assertEqual(result, 'red')
 
+  def test_generic_color(self):
+    test_colors = np.array(
+        [(255, 0, 0), (55, 118, 171), (73, 128, 41), (231, 112, 13)]
+    )
+    expected_colors = ['red', 'blue', 'green', 'orange']
+
+    result = color_and_property_extractor.get_generic_color_name(test_colors)
+
+    self.assertEqual(result, expected_colors)
+
 
 if __name__ == '__main__':
   unittest.main()