mean color utils functions

justfile

@@ -1,5 +1,7 @@
 separate-compose := "command -v docker-compose 2>/dev/null"
 
+# TODO: add compose variable to avoid branching in every recipe
+
 init:
     if [ ! -f requirements.lock ]; then rye sync; fi
     rye run pre-commit install

scraper/src/utils.py

@@ -72,7 +72,18 @@ async def process_image(url: str, info: ScraperInfo) -> None:
         img.save(output_path)
 
     await ctx.image_repo.add(
-        entities.Image(id=id, url=url, hash=hash_value, processed=0)
+        entities.Image(
+            id=id,
+            url=url,
+            hash=hash_value,
+            mean_h=0,
+            mean_s=0,
+            mean_v=0,
+            mean_l=0,
+            mean_a=0,
+            mean_b=0,
+            processed=0,
+        )
     )
     info.images_scraped += 1
 

shared/color.py

@@ -0,0 +1,26 @@
+from enum import Enum
+
+import numpy as np
+from PIL import Image
+from skimage import color
+
+
+class ColorModel(int, Enum):
+    LAB = 0
+    HSV = 1
+
+
+def mean_color(image_path, color_model):
+    image = Image.open(image_path).convert("RGBA")
+    np_image = np.array(image)
+
+    mask = np_image[..., 3] > 0
+    np_image = np_image[..., :3][mask].reshape(-1, 3) / 255.0
+
+    if color_model == ColorModel.LAB:
+        converted_image = color.rgb2lab(np_image)
+    else:
+        converted_image = color.rgb2hsv(np_image)
+
+    mean_color = np.mean(converted_image, axis=0)
+    return mean_color

shared/db.py

@@ -91,7 +91,9 @@ class SqliteRepository(AbstractRepository):
     async def create_table(self) -> None:
         await self._db.execute(
             query=f"""CREATE TABLE IF NOT EXISTS {self._table_name}
-                (id TEXT, url TEXT, hash TEXT, processed INTEGER);"""
+                (id TEXT, url TEXT, hash TEXT, mean_h REAL, mean_s REAL,
+                 mean_v REAL, mean_l REAL, mean_a REAL, mean_b REAL,
+                 processed INTEGER);"""
         )
 
     async def get_many_from_list(self, field, values) -> List[Entity]:

shared/entities.py

@@ -7,6 +7,12 @@ class Image(Entity):
     id: str
     url: str
     hash: str
+    mean_h: float
+    mean_s: float
+    mean_v: float
+    mean_l: float
+    mean_a: float
+    mean_b: float
     processed: int  # bool but SQLite moment
 
     _pk: ClassVar[str] = "id"