Added backend filament search by similar color

This commit is contained in:
Donkie
2023-12-19 21:37:07 +01:00
parent 8322451f26
commit cf8d7472f5
5 changed files with 167 additions and 1 deletions

View File

@@ -183,6 +183,19 @@ async def find(
"Specify an empty string to match filaments with no article number." "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( sort: Optional[str] = Query(
default=None, default=None,
title="Sort", title="Sort",
@@ -217,8 +230,19 @@ async def find(
else: else:
vendor_ids = None 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_items, total_count = await filament.find(
db=db, db=db,
ids=filter_by_ids,
vendor_name=vendor_name if vendor_name is not None else vendor_name_old, vendor_name=vendor_name if vendor_name is not None else vendor_name_old,
vendor_id=vendor_ids, vendor_id=vendor_ids,
name=name, name=name,

View File

@@ -14,12 +14,14 @@ from spoolman.api.v1.models import EventType, Filament, FilamentEvent
from spoolman.database import models, vendor from spoolman.database import models, vendor
from spoolman.database.utils import ( from spoolman.database.utils import (
SortOrder, SortOrder,
add_where_clause_int_in,
add_where_clause_int_opt, add_where_clause_int_opt,
add_where_clause_str, add_where_clause_str,
add_where_clause_str_opt, add_where_clause_str_opt,
parse_nested_field, parse_nested_field,
) )
from spoolman.exceptions import ItemDeleteError, ItemNotFoundError from spoolman.exceptions import ItemDeleteError, ItemNotFoundError
from spoolman.math import delta_e, hex_to_rgb, rgb_to_lab
from spoolman.ws import websocket_manager 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( async def find(
*, *,
db: AsyncSession, db: AsyncSession,
ids: Optional[list[int]] = None,
vendor_name: Optional[str] = None, vendor_name: Optional[str] = None,
vendor_id: Optional[Union[int, Sequence[int]]] = None, vendor_id: Optional[Union[int, Sequence[int]]] = None,
name: Optional[str] = None, name: Optional[str] = None,
@@ -104,6 +107,7 @@ async def find(
.join(models.Filament.vendor, isouter=True) .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_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(stmt, models.Vendor.name, vendor_name)
stmt = add_where_clause_str_opt(stmt, models.Filament.name, 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] 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: async def filament_changed(filament: models.Filament, typ: EventType) -> None:
"""Notify websocket clients that a filament has changed.""" """Notify websocket clients that a filament has changed."""
await websocket_manager.send( await websocket_manager.send(

View File

@@ -2,7 +2,7 @@
from collections.abc import Sequence from collections.abc import Sequence
from enum import Enum from enum import Enum
from typing import Any, Optional, Union from typing import Any, Optional, TypeVar, Union
import sqlalchemy import sqlalchemy
from sqlalchemy import Select from sqlalchemy import Select
@@ -105,3 +105,17 @@ def add_where_clause_int_opt(
statements.append(field == value_part) statements.append(field == value_part)
stmt = stmt.where(sqlalchemy.or_(*statements)) stmt = stmt.where(sqlalchemy.or_(*statements))
return stmt 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

View File

@@ -1,5 +1,7 @@
"""Various math-related functions.""" """Various math-related functions."""
# ruff: noqa: PLR2004
import math import math
@@ -33,3 +35,60 @@ def length_from_weight(*, weight: float, diameter: float, density: float) -> flo
volume_cm3 = weight / density volume_cm3 = weight / density
volume_mm3 = volume_cm3 * 1000 volume_mm3 = volume_cm3 * 1000
return volume_mm3 / (math.pi * (diameter / 2) ** 2) 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]

View File

@@ -35,6 +35,7 @@ def filaments(random_vendor_mod: dict[str, Any], random_empty_vendor_mod: dict[s
"weight": 1000, "weight": 1000,
"spool_weight": 250, "spool_weight": 250,
"article_number": "123456789", "article_number": "123456789",
"color_hex": "FF0000",
"comment": "abcdefghåäö", "comment": "abcdefghåäö",
}, },
) )
@@ -53,6 +54,7 @@ def filaments(random_vendor_mod: dict[str, Any], random_empty_vendor_mod: dict[s
"weight": 1000, "weight": 1000,
"spool_weight": 250, "spool_weight": 250,
"article_number": "987654321", "article_number": "987654321",
"color_hex": "EE0000",
"comment": "abcdefghåäö", "comment": "abcdefghåäö",
}, },
) )
@@ -70,6 +72,7 @@ def filaments(random_vendor_mod: dict[str, Any], random_empty_vendor_mod: dict[s
"weight": 1000, "weight": 1000,
"spool_weight": 250, "spool_weight": 250,
"article_number": "abc", "article_number": "abc",
"color_hex": "000000",
"comment": "abcdefghåäö", "comment": "abcdefghåäö",
}, },
) )
@@ -418,3 +421,38 @@ def test_find_filaments_by_empty_article_number(filaments: Fixture):
# Verify # Verify
filaments_result = result.json() filaments_result = result.json()
assert filament_lists_equal(filaments_result, filaments.filaments[3:]) 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]],
)