AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-17 11:05:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Add cloth category selection feature to u2net_cloth_seg (#485)

Browse Source
This commit is contained in:
szriru 2023-07-13 15:09:13 +09:00 committed by GitHub
parent a9c2a39213
commit f0019d723b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 28 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
rembg/sessions/u2net_cloth_seg.py
View File

@ -75,21 +75,37 @@ class Unet2ClothSession(BaseSession):
masks = [] masks = []
cloth_category = kwargs.get("cc") or kwargs.get("cloth_category")
def upper_cloth():
mask1 = mask.copy() mask1 = mask.copy()
mask1.putpalette(palette1) mask1.putpalette(palette1)
mask1 = mask1.convert("RGB").convert("L") mask1 = mask1.convert("RGB").convert("L")
masks.append(mask1) masks.append(mask1)
def lower_cloth():
mask2 = mask.copy() mask2 = mask.copy()
mask2.putpalette(palette2) mask2.putpalette(palette2)
mask2 = mask2.convert("RGB").convert("L") mask2 = mask2.convert("RGB").convert("L")
masks.append(mask2) masks.append(mask2)
def full_cloth():
mask3 = mask.copy() mask3 = mask.copy()
mask3.putpalette(palette3) mask3.putpalette(palette3)
mask3 = mask3.convert("RGB").convert("L") mask3 = mask3.convert("RGB").convert("L")
masks.append(mask3) masks.append(mask3)
if cloth_category == "upper":
upper_cloth()
elif cloth_category == "lower":
lower_cloth()
elif cloth_category == "full":
full_cloth()
else:
upper_cloth()
lower_cloth()
full_cloth()
return masks return masks
@classmethod @classmethod