feat: add color extraction

setup.py

@@ -1,17 +1,20 @@
 from setuptools import setup, find_packages
 
 HYPER_E_DOT = "-e ."
+
+
 def getRequirements(requirementsPath: str) -> list[str]:
     with open(requirementsPath) as file:
         requirements = file.read().split("\n")
     requirements.remove(HYPER_E_DOT)
     return requirements
 
+
 setup(
-    name = "Magical-Mirror",
-    author = "Subramani Sivakumar",
-    author_email = "bwsubbu@gmail.com",
-    version = "0.1",
-    packages = find_packages(),
-    install_requires = getRequirements(requirementsPath = "./requirements.txt")
-)
+    name="Jewel Mirror",
+    author="Ishwor Subedi",
+    author_email="ishworr.subedi@gmail.com",
+    version="0.1",
+    packages=find_packages(),
+    install_requires=getRequirements(requirementsPath="./requirements.txt")
+)

src/api/image_prep_api.py

@@ -6,6 +6,8 @@ author  @ github.com/ishworrsubedii
 import base64
 import os
 from io import BytesIO
+
+import cv2
 import numpy as np
 import replicate
 import requests
@@ -14,12 +16,14 @@ from fastapi import APIRouter, UploadFile, File, HTTPException
 from fastapi.responses import JSONResponse
 
 from src.components.auto_crop import crop_transparent_image
+from src.components.color_extraction import ColorExtractionRMBG
 
 preprocessing_router = APIRouter()
 
 rmbg: str = os.getenv("RMBG")
 
 enhancer: str = os.getenv("ENHANCER")
+color_extraction_rmbg = ColorExtractionRMBG()
 
 
 def replicate_bg(input):
@@ -131,7 +135,7 @@ async def crop_transparent(image: UploadFile):
 
 
 @preprocessing_router.post("/background_replace")
-async def nto_id(image: UploadFile = File(...), bg_image: UploadFile = File(...)):
+async def bg_replace(image: UploadFile = File(...), bg_image: UploadFile = File(...)):
     image_bytes = await image.read()
     bg_image = await bg_image.read()
 
@@ -156,3 +160,30 @@ async def nto_id(image: UploadFile = File(...), bg_image: UploadFile = File(...)
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")
+
+
+@preprocessing_router.post("/rem_bg_color_extraction")
+async def remove_background_color_extraction(image: UploadFile = File(...), hex_color: str = "#FFFFFF",
+                                             threshold: int = 30):
+    image_bytes = await image.read()
+    image = Image.open(BytesIO(image_bytes)).convert("RGBA")
+    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+
+    result = color_extraction_rmbg.extract_color(image, hex_color, threshold)
+    result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_RGB2BGR)).convert("RGBA")
+    act_img_base_64 = BytesIO()
+    result.save(act_img_base_64, format="PNG")
+    image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+
+    image_data_uri = f"data:image/png;base64,{image_bytes_}"
+
+    try:
+        response = {
+            "output": f"{image_data_uri}",
+            'code': 200
+        }
+
+        return JSONResponse(content=response, status_code=200)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")

src/components/color_extraction.py

@@ -0,0 +1,38 @@
+"""
+project @ NTO-TCP-HF
+created @ 2024-10-28
+author  @ github.com/ishworrsubedii
+"""
+
+import cv2
+import numpy as np
+
+
+class ColorExtractionRMBG:
+    def __init__(self):
+        self.HSV_LOWER = None
+        self.HSV_UPPER = None
+
+    def hex_to_rgb(self, hex_color: str):
+        hex_color = hex_color.lstrip("#")
+        return tuple(int(hex_color[i:i + 2], 16) for i in (0, 2, 4))
+
+    def rgb_to_hsv(self, rgb_color):
+        rgb_array = np.uint8([[rgb_color]])
+        hsv_array = cv2.cvtColor(rgb_array, cv2.COLOR_RGB2HSV)
+        return hsv_array[0][0]
+
+    def set_thresholds(self, hex_color: str, threshold: int):
+        hsv_color = self.rgb_to_hsv(self.hex_to_rgb(hex_color))
+        lower_bound = np.clip([hsv_color[0] - threshold, 50, 50], 0, 255)
+        upper_bound = np.clip([hsv_color[0] + threshold, 255, 255], 0, 255)
+        return lower_bound, upper_bound
+
+    def extract_color(self, image: np.ndarray, hex_color: str, threshold: int):
+        self.HSV_LOWER, self.HSV_UPPER = self.set_thresholds(hex_color, threshold)
+        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+        mask = cv2.inRange(hsv, self.HSV_LOWER, self.HSV_UPPER)
+        result = cv2.bitwise_and(image, image, mask=mask)
+        result = cv2.cvtColor(result, cv2.COLOR_BGR2BGRA)
+        result[:, :, 3] = mask
+        return result