from transformers import pipeline
from PIL import Image
import numpy as np
import cv2  # OpenCV for better mask processing

# Initialize segmentation pipeline
segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")


def segment_clothing(img, clothes):
    # Segment image
    segments = segmenter(img)

    # Define clothing items to expand
    EXPAND_CLOTHING = {"Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe"}

    # Create list of masks
    mask_list = []
    expand_mask_list = []  # Separate list for clothes that need expansion

    for s in segments:
        mask = np.array(s['mask'], dtype=np.uint8)  # Convert mask to numpy array
        if s['label'] in clothes:
            if s['label'] in EXPAND_CLOTHING:
                expand_mask_list.append(mask)  # Store separately for expansion
            else:
                mask_list.append(mask)  # Keep others as they are

    if not mask_list and not expand_mask_list:
        return img  # Return original image if no relevant items found

    # Initialize final mask with zeros
    final_mask = np.zeros_like(mask_list[0] if mask_list else expand_mask_list[0], dtype=np.uint8)

    # Combine normal masks into one
    for mask in mask_list:
        final_mask = np.maximum(final_mask, mask)

    # Expand selected clothing masks using closing + dilation
    for mask in expand_mask_list:
        height, width = mask.shape
        kernel_size = max(20, int(0.02 * min(height, width)))  # 5% expansion
        print(kernel_size)
        print(height)
        print(width)
        kernel = np.ones((kernel_size, kernel_size), np.uint8)

        # **Step 1: Fill gaps using Closing (Dilation + Erosion)**
        closed_mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel, iterations=1)

        # **Step 2: Expand using Dilation**
        dilated_mask = cv2.dilate(closed_mask, kernel, iterations=1)

        # Merge into final mask
        final_mask = np.maximum(final_mask, dilated_mask)

    # Optional: Use contour filling to ensure all areas within contours are filled
    contours, _ = cv2.findContours(final_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cv2.drawContours(final_mask, contours, -1, (255), thickness=cv2.FILLED)

    # Convert mask to binary (0 or 255) if needed for alpha channel
    _, final_mask = cv2.threshold(final_mask, 127, 255, cv2.THRESH_BINARY)

    # Convert final mask from numpy array to PIL image
    final_mask = Image.fromarray(final_mask)

    # Apply mask to original image (convert to RGBA first)
    img = img.convert("RGBA")
    img.putalpha(final_mask)

    return img