Added backend filament search by similar color
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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]],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user