diff --git a/spoolman/api/v1/filament.py b/spoolman/api/v1/filament.py index fc85fe0..451f446 100644 --- a/spoolman/api/v1/filament.py +++ b/spoolman/api/v1/filament.py @@ -183,6 +183,19 @@ async def find( "Specify an empty string to match filaments with no article number." ), ), + color_hex: Optional[str] = Query( + default=None, + title="Filament Color", + description="Match filament by similar color. Slow operation!", + ), + color_similarity_threshold: float = Query( + default=20.0, + description=( + "The similarity threshold for color matching. " + "A value between 0.0-100.0, where 0 means match only exactly the same color." + ), + example=20.0, + ), sort: Optional[str] = Query( default=None, title="Sort", @@ -217,8 +230,19 @@ async def find( else: vendor_ids = None + if color_hex is not None: + matched_filaments = await filament.find_by_color( + db=db, + color_query_hex=color_hex, + similarity_threshold=color_similarity_threshold, + ) + filter_by_ids = [db_filament.id for db_filament in matched_filaments] + else: + filter_by_ids = None + db_items, total_count = await filament.find( db=db, + ids=filter_by_ids, vendor_name=vendor_name if vendor_name is not None else vendor_name_old, vendor_id=vendor_ids, name=name, diff --git a/spoolman/database/filament.py b/spoolman/database/filament.py index 132471c..05ace5f 100644 --- a/spoolman/database/filament.py +++ b/spoolman/database/filament.py @@ -14,12 +14,14 @@ from spoolman.api.v1.models import EventType, Filament, FilamentEvent from spoolman.database import models, vendor from spoolman.database.utils import ( SortOrder, + add_where_clause_int_in, add_where_clause_int_opt, add_where_clause_str, add_where_clause_str_opt, parse_nested_field, ) from spoolman.exceptions import ItemDeleteError, ItemNotFoundError +from spoolman.math import delta_e, hex_to_rgb, rgb_to_lab from spoolman.ws import websocket_manager @@ -82,6 +84,7 @@ async def get_by_id(db: AsyncSession, filament_id: int) -> models.Filament: async def find( *, db: AsyncSession, + ids: Optional[list[int]] = None, vendor_name: Optional[str] = None, vendor_id: Optional[Union[int, Sequence[int]]] = None, name: Optional[str] = None, @@ -104,6 +107,7 @@ async def find( .join(models.Filament.vendor, isouter=True) ) + stmt = add_where_clause_int_in(stmt, models.Filament.id, ids) stmt = add_where_clause_int_opt(stmt, models.Filament.vendor_id, vendor_id) stmt = add_where_clause_str(stmt, models.Vendor.name, vendor_name) stmt = add_where_clause_str_opt(stmt, models.Filament.name, name) @@ -190,6 +194,33 @@ async def find_article_numbers( return [row[0] for row in rows.all() if row[0] is not None] +async def find_by_color( + *, + db: AsyncSession, + color_query_hex: str, + similarity_threshold: float = 25, +) -> list[models.Filament]: + """Find a list of filament objects by similarity to a color. + + This performs a server-side search, where all filaments are loaded into memory, making it not so efficient. + The similarity threshold is a value between 0 and 100, where 0 means the colors must be identical and 100 means + pretty much all colors are considered similar. + """ + filaments, _ = await find(db=db) + + color_query_lab = rgb_to_lab(hex_to_rgb(color_query_hex)) + + found_filaments: list[models.Filament] = [] + for filament in filaments: + if filament.color_hex is None: + continue + color_lab = rgb_to_lab(hex_to_rgb(filament.color_hex)) + if delta_e(color_query_lab, color_lab) <= similarity_threshold: + found_filaments.append(filament) + + return found_filaments + + async def filament_changed(filament: models.Filament, typ: EventType) -> None: """Notify websocket clients that a filament has changed.""" await websocket_manager.send( diff --git a/spoolman/database/utils.py b/spoolman/database/utils.py index c9b9527..5002e68 100644 --- a/spoolman/database/utils.py +++ b/spoolman/database/utils.py @@ -2,7 +2,7 @@ from collections.abc import Sequence from enum import Enum -from typing import Any, Optional, Union +from typing import Any, Optional, TypeVar, Union import sqlalchemy from sqlalchemy import Select @@ -105,3 +105,17 @@ def add_where_clause_int_opt( statements.append(field == value_part) stmt = stmt.where(sqlalchemy.or_(*statements)) return stmt + + +T = TypeVar("T") + + +def add_where_clause_int_in( + stmt: Select, + field: attributes.InstrumentedAttribute[T], + value: Optional[Sequence[T]], +) -> Select: + """Add a where clause to a select statement for a field.""" + if value is not None and len(value) > 0: + stmt = stmt.where(field.in_(value)) + return stmt diff --git a/spoolman/math.py b/spoolman/math.py index a7940b8..e7383e6 100644 --- a/spoolman/math.py +++ b/spoolman/math.py @@ -1,5 +1,7 @@ """Various math-related functions.""" +# ruff: noqa: PLR2004 + import math @@ -33,3 +35,60 @@ def length_from_weight(*, weight: float, diameter: float, density: float) -> flo volume_cm3 = weight / density volume_mm3 = volume_cm3 * 1000 return volume_mm3 / (math.pi * (diameter / 2) ** 2) + + +def rgb_to_lab(rgb: list[int]) -> list[float]: + """Convert a RGB color to CIELAB. + + Input is of form [r, g, b] where r, g, and b are integers between 0 and 255. + Output is of form [l, a, b] where l, a, and b are floats. + """ + r, g, b = rgb[0] / 255, rgb[1] / 255, rgb[2] / 255 + + r = (r / 12.92) if (r <= 0.04045) else math.pow((r + 0.055) / 1.055, 2.4) + g = (g / 12.92) if (g <= 0.04045) else math.pow((g + 0.055) / 1.055, 2.4) + b = (b / 12.92) if (b <= 0.04045) else math.pow((b + 0.055) / 1.055, 2.4) + + x = (r * 0.4124 + g * 0.3576 + b * 0.1805) / 0.95047 + y = (r * 0.2126 + g * 0.7152 + b * 0.0722) / 1.00000 + z = (r * 0.0193 + g * 0.1192 + b * 0.9505) / 1.08883 + + x = math.pow(x, 1 / 3) if (x > 0.008856) else (7.787 * x) + 16 / 116 + y = math.pow(y, 1 / 3) if (y > 0.008856) else (7.787 * y) + 16 / 116 + z = math.pow(z, 1 / 3) if (z > 0.008856) else (7.787 * z) + 16 / 116 + + return [(116 * y) - 16, 500 * (x - y), 200 * (y - z)] + + +def delta_e(lab_a: list[float], lab_b: list[float]) -> float: + """Calculate the color difference between two CIELAB colors.""" + delta_l = lab_a[0] - lab_b[0] + delta_a = lab_a[1] - lab_b[1] + delta_b = lab_a[2] - lab_b[2] + c1 = math.sqrt(lab_a[1] * lab_a[1] + lab_a[2] * lab_a[2]) + c2 = math.sqrt(lab_b[1] * lab_b[1] + lab_b[2] * lab_b[2]) + delta_c = c1 - c2 + delta_h = delta_a * delta_a + delta_b * delta_b - delta_c * delta_c + delta_h = math.sqrt(delta_h) if delta_h > 0 else 0 + sc = 1.0 + 0.045 * c1 + sh = 1.0 + 0.015 * c1 + delta_l_kl_sl = delta_l / 1.0 + delta_c_kc_sc = delta_c / sc + delta_h_kh_sh = delta_h / sh + i = delta_l_kl_sl * delta_l_kl_sl + delta_c_kc_sc * delta_c_kc_sc + delta_h_kh_sh * delta_h_kh_sh + return math.sqrt(i) if i > 0 else 0 + + +def hex_to_rgb(hex_code: str) -> list[int]: + """Convert a hex color code to RGB. + + Input is of form #RRGGBB where RR, GG, and BB are hexadecimal numbers. + Output is of form [r, g, b] where r, g, and b are integers between 0 and 255. + """ + hex_code = hex_code.lstrip("#") + + r = int(hex_code[0:2], 16) + g = int(hex_code[2:4], 16) + b = int(hex_code[4:6], 16) + + return [r, g, b] diff --git a/tests_integration/tests/filament/test_find.py b/tests_integration/tests/filament/test_find.py index a5920fd..a00c985 100644 --- a/tests_integration/tests/filament/test_find.py +++ b/tests_integration/tests/filament/test_find.py @@ -35,6 +35,7 @@ def filaments(random_vendor_mod: dict[str, Any], random_empty_vendor_mod: dict[s "weight": 1000, "spool_weight": 250, "article_number": "123456789", + "color_hex": "FF0000", "comment": "abcdefghåäö", }, ) @@ -53,6 +54,7 @@ def filaments(random_vendor_mod: dict[str, Any], random_empty_vendor_mod: dict[s "weight": 1000, "spool_weight": 250, "article_number": "987654321", + "color_hex": "EE0000", "comment": "abcdefghåäö", }, ) @@ -70,6 +72,7 @@ def filaments(random_vendor_mod: dict[str, Any], random_empty_vendor_mod: dict[s "weight": 1000, "spool_weight": 250, "article_number": "abc", + "color_hex": "000000", "comment": "abcdefghåäö", }, ) @@ -418,3 +421,38 @@ def test_find_filaments_by_empty_article_number(filaments: Fixture): # Verify filaments_result = result.json() assert filament_lists_equal(filaments_result, filaments.filaments[3:]) + + +def test_find_filaments_by_similar_color(filaments: Fixture): + # Execute + result = httpx.get( + f"{URL}/api/v1/filament", + params={ + "color_hex": "FE0000", + "color_similarity_threshold": 20.0, + }, + ) + result.raise_for_status() + + # Verify + filaments_result = result.json() + assert filament_lists_equal(filaments_result, [filaments.filaments[0], filaments.filaments[1]]) + + +def test_find_filaments_by_similar_color_100(filaments: Fixture): + # Execute + result = httpx.get( + f"{URL}/api/v1/filament", + params={ + "color_hex": "FE0000", + "color_similarity_threshold": 100.0, + }, + ) + result.raise_for_status() + + # Verify + filaments_result = result.json() + assert filament_lists_equal( + filaments_result, + [filaments.filaments[0], filaments.filaments[1], filaments.filaments[2]], + )