import cv2
import numpy as np


# Convert Original Mask into Custom Mask Contain only clothes mask 
def remap_mask(mask, label_map):
    new_mask = np.zeros_like(mask, dtype=np.uint8)
    for old_id, new_id in label_map.items():
        new_mask[mask == old_id] = new_id # convert non-clothes to background
    return new_mask

# Convert mask to background -> 0 / clothes -> 1
def to_binary_clothes(mask):
    return (mask > 0).astype(np.uint8)

# Add Morphological operations to remove noise and fill gab inside clothes
def morphological_operations(mask, opening_kernel = 3, closing_kernel = 11, opening_iterations=1, closing_iterations=1):
    mask = mask.astype(np.uint8)
    
    # Opening -> Remove Small Noise
    open_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (opening_kernel, opening_kernel))
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, open_kernel, iterations=opening_iterations)
    
    # CLosing -> fill gaps
    close_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (closing_kernel, closing_kernel))
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, close_kernel, iterations=closing_iterations)
    
    return mask

# create overlay 
def create_overlay(image, mask, alpha=0.5):
    overlay = image.copy()
    clothes = mask == 1
    overlay[clothes] = (overlay[clothes] * (1 - alpha) + np.array([0, 255, 0]) * alpha).astype(np.uint8)

    return overlay