src/components/color_extraction.py

"""
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