ColorMask.py

import numpy as np
from PIL import Image
import cv2


def hex_to_rgb(hex_color):
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
def rgb_to_bgr(rgb_color):
    return rgb_color[::-1]
def create_mask(image_np, color_to_replace, tolerance=20):
    # Convert hex colors to RGB and then to BGR
    color_to_replace_rgb = hex_to_rgb(color_to_replace)
    # replacement_color_rgb = hex_to_rgb(replacement_color)
    color_to_replace_bgr = rgb_to_bgr(color_to_replace_rgb)
    # replacement_color_bgr = rgb_to_bgr(replacement_color_rgb)
    # Calculate the lower and upper bounds for color tolerance in BGR
    lower_bound = np.array([max(c - tolerance, 0) for c in color_to_replace_bgr])
    upper_bound = np.array([min(c + tolerance, 255) for c in color_to_replace_bgr])
    # Convert the image to BGR format for OpenCV processing
    image_bgr = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
    # Create a mask for pixels within the tolerance range using cv2.inRange
    mask = cv2.inRange(image_bgr, lower_bound, upper_bound)
    # Replace color in the masked area
    # image_bgr[mask > 0] = replacement_color_bgr
    # Convert back to RGB
    # result_image_np = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    # Save the result
    result_image = Image.fromarray(mask)
    result_image.save('./result.jpg')
    return result_image