Spaces:
Runtime error
Runtime error
File size: 2,343 Bytes
61a0586 c8c83ea 8f2a481 871269e 8f2a481 1c0cda0 c89d4d6 8f2a481 0bac47f 0703834 1fceef8 eb456a2 c89d4d6 8f2a481 c89d4d6 8f2a481 187d444 a05ee1f c8c83ea c89d4d6 c8c83ea 0bac47f c89d4d6 eb456a2 7046d86 c8c83ea 5eb5c61 c11222a eb456a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
from transformers import pipeline
from PIL import Image, ImageChops, ImageOps
import numpy as np
from io import BytesIO
import base64
# Initialisation du pipeline de segmentation
segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
def encode_image_to_base64(image):
buffered = BytesIO()
image.save(buffered, format="PNG")
return base64.b64encode(buffered.getvalue()).decode('utf-8')
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)
# 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))
# Encodage de l'image recadrée en base64
imageBase64 = encode_image_to_base64(cropped_image)
#result_images.append((s['label'], imageBase64))
result_images.append(imageBase64)
return result_images
|