Spaces:
Runtime error
Runtime error
| 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 | |