Spaces:
Running
Running
File size: 2,626 Bytes
8682ae7 ae983b9 8682ae7 f36990d 8682ae7 2f63795 7ac5a56 8682ae7 a72dd3e 2f63795 a72dd3e 8682ae7 a72dd3e 8682ae7 a72dd3e ae983b9 a72dd3e 8682ae7 a72dd3e 7ac5a56 f36990d a72dd3e 8682ae7 a72dd3e 8682ae7 f36990d 2f63795 a72dd3e 73cc779 e01b37c 73cc779 a72dd3e 2f63795 a72dd3e 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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
|