Spaces:
Build error
Build error
File size: 1,327 Bytes
5754201 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """
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
|