Spaces:
Running
Running
File size: 1,805 Bytes
8682ae7 ae983b9 8682ae7 f36990d 8682ae7 7ac5a56 8682ae7 ae983b9 f36990d 8682ae7 7ac5a56 f36990d 8682ae7 f36990d 8682ae7 f36990d 7ac5a56 ae983b9 f36990d 7ac5a56 8682ae7 f36990d 8682ae7 f36990d 7ac5a56 |
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 |
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)
# Create list of masks
mask_list = []
for s in segments:
if s['label'] in clothes:
mask_list.append(np.array(s['mask'], dtype=np.uint8)) # Convert to numpy array and ensure it's uint8
if not mask_list:
return img # Return original image if no clothes found
# Initialize final mask with zeros
final_mask = np.zeros_like(mask_list[0], dtype=np.uint8)
# Combine masks into one
for mask in mask_list:
final_mask = np.maximum(final_mask, mask)
# Expand clothing boundaries using morphological dilation
height, width = final_mask.shape
kernel_size = max(1, int(0.05 * min(height, width))) # Ensure at least 1 pixel
kernel = np.ones((kernel_size, kernel_size), np.uint8)
# Dilate mask to expand clothing area
final_mask = cv2.dilate(final_mask, kernel, iterations=1)
# 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
|