from transformers import pipeline
from PIL import Image, ImageChops
import numpy as np
from io import BytesIO
import base64

# Initialisation du pipeline de segmentation
segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")

def is_image_significant(image, threshold=0.01):
    """
    Vérifie si une image contient suffisamment de contenu non transparent.
    :param image: L'image PIL à vérifier.
    :param threshold: Seuil de pourcentage de pixels non transparents pour considérer l'image comme significative.
    :return: True si l'image est significative, sinon False.
    """
    np_image = np.array(image)
    # Compte le nombre de pixels non transparents
    non_transparent_pixels = np.sum(np_image[:, :, 3] > 0)
    # Calcul du pourcentage de pixels non transparents
    total_pixels = np_image.shape[0] * np_image.shape[1]
    non_transparent_percentage = non_transparent_pixels / total_pixels
    
    return non_transparent_percentage > threshold

def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"], margin=10):
    # Segmentation de l'image
    segments = segmenter(img)
    
    # Liste des images segmentées
    result_images = []

    for s in segments:
        if s['label'] in clothes:
            # Conversion du masque en tableau NumPy
            mask_array = np.array(s['mask'])
            
            # Création d'une image vide avec transparence
            empty_image = Image.new("RGBA", img.size, (0, 0, 0, 0))
            
            # Conversion du masque en image PIL (niveau de gris)
            mask_image = Image.fromarray(mask_array).convert("L")
            
            # Extraction de la partie de l'image correspondant au masque
            segmented_part = ImageChops.multiply(img.convert("RGBA"), Image.merge("RGBA", [mask_image, mask_image, mask_image, mask_image]))

            # Application du masque sur l'image vide
            empty_image.paste(segmented_part, mask=mask_image)
            
            # Vérifier si l'image est significative avant le recadrage
            if is_image_significant(empty_image):
                # Déterminer la bounding box du masque
                bbox = mask_image.getbbox()
                if bbox:
                    # Ajouter la marge autour de la bounding box
                    left, top, right, bottom = bbox
                    left = max(0, left - margin)
                    top = max(0, top - margin)
                    right = min(img.width, right + margin)
                    bottom = min(img.height, bottom + margin)
                    
                    # Recadrer l'image à la taille du masque avec la marge
                    cropped_image = empty_image.crop((left, top, right, bottom))
                    
                    result_images.append(cropped_image)

    return result_images