From dd3d3f70f475349cba9de91f3469edb6f2068e8f Mon Sep 17 00:00:00 2001
From: ArtemiyGranat <artemiigranat@gmail.com>
Date: Fri, 8 Mar 2024 19:04:14 +0400
Subject: [PATCH] mean color utils functions

---
 justfile             |  2 ++
 scraper/src/utils.py | 13 ++++++++++++-
 shared/color.py      | 26 ++++++++++++++++++++++++++
 shared/db.py         |  4 +++-
 shared/entities.py   |  6 ++++++
 5 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 shared/color.py

diff --git a/justfile b/justfile
index f20b28d..683ad1f 100644
--- a/justfile
+++ b/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
diff --git a/scraper/src/utils.py b/scraper/src/utils.py
index 2638f87..8c53f65 100644
--- a/scraper/src/utils.py
+++ b/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
 
diff --git a/shared/color.py b/shared/color.py
new file mode 100644
index 0000000..c6bb514
--- /dev/null
+++ b/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
diff --git a/shared/db.py b/shared/db.py
index 864f160..980c7df 100644
--- a/shared/db.py
+++ b/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]:
diff --git a/shared/entities.py b/shared/entities.py
index 9e264fe..4c8b9e2 100644
--- a/shared/entities.py
+++ b/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"