├── requirements.txt ├── image.jpg ├── SegCloth.py └── README.md /requirements.txt: -------------------------------------------------------------------------------- 1 | transformers 2 | torch 3 | pillow 4 | numpy -------------------------------------------------------------------------------- /image.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/TonyAssi/Segment-Clothing/HEAD/image.jpg -------------------------------------------------------------------------------- /SegCloth.py: -------------------------------------------------------------------------------- 1 | from transformers import pipeline 2 | from PIL import Image 3 | import numpy as np 4 | import os 5 | 6 | 7 | # Initialize segmentation pipeline 8 | segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes") 9 | 10 | 11 | def segment_clothing(img, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]): 12 | # Segment image 13 | segments = segmenter(img) 14 | 15 | # Create list of masks 16 | mask_list = [] 17 | for s in segments: 18 | if(s['label'] in clothes): 19 | mask_list.append(s['mask']) 20 | 21 | 22 | # Paste all masks on top of eachother 23 | final_mask = np.array(mask_list[0]) 24 | for mask in mask_list: 25 | current_mask = np.array(mask) 26 | final_mask = final_mask + current_mask 27 | 28 | # Convert final mask from np array to PIL image 29 | final_mask = Image.fromarray(final_mask) 30 | 31 | # Apply mask to original image 32 | img.putalpha(final_mask) 33 | 34 | return img 35 | 36 | 37 | def batch_segment_clothing(img_dir, out_dir, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]): 38 | # Create output directory if it doesn't exist 39 | if not os.path.exists(out_dir): 40 | os.makedirs(out_dir) 41 | 42 | # Iterate through each file in the input directory 43 | for filename in os.listdir(img_dir): 44 | if filename.endswith(".jpg") or filename.endswith(".JPG") or filename.endswith(".png") or filename.endswith(".PNG"): 45 | try: 46 | # Load image 47 | img_path = os.path.join(img_dir, filename) 48 | img = Image.open(img_path).convert("RGBA") 49 | 50 | # Segment clothing 51 | segmented_img = segment_clothing(img, clothes) 52 | 53 | # Save segmented image to output directory as PNG 54 | out_path = os.path.join(out_dir, filename.split('.')[0] + ".png") 55 | segmented_img.save(out_path) 56 | 57 | print(f"Segmented {filename} successfully.") 58 | 59 | except Exception as e: 60 | print(f"Error processing {filename}: {e}") 61 | 62 | else: 63 | print(f"Skipping {filename} as it is not a supported image file.") 64 | 65 | 66 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Segment Clothing 2 | Segment clothing with 1 line of code. 3 | 4 | Takes in a PIL image and outputs a PIL image of segmented clothing. Built on top of 🤗 Tranformers using the [mattmdjaga/segformer_b2_clothes](https://huggingface.co/mattmdjaga/segformer_b2_clothes) image segmentation model. 5 | 6 | Try out the Web Demo: [![Hugging Face Spaces](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/tonyassi/clothing-segmentation) 7 | 8 | Screenshot 2025-11-26 at 1 26 53 PM 9 | 10 | ## Installation 11 | ```bash 12 | pip install -r requirements.txt 13 | ``` 14 | ``` 15 | Python 3.9 16 | Transformers 4.37.2 17 | Torch 2.2.0 18 | Pillow 10.2.0 19 | NumPy 1.26.3 20 | ``` 21 | 22 | ## Usage 23 | 24 | Import module 25 | ```python 26 | from SegCloth import segment_clothing 27 | ``` 28 | 29 | Import PIL and open image 30 | ```python 31 | from PIL import Image 32 | image = Image.open('image.jpg') 33 | ``` 34 | 35 | ![](https://cdn.discordapp.com/attachments/1120417968032063538/1202309847287345253/image-1.jpg?ex=65e8accd&is=65d637cd&hm=f42cd1095001982434a3b05907409ef8d3a380a860a7c7e079ab82f558842697&) 36 | --- 37 | 38 | ### Segment Clothing 39 | - **img** input image of type PIL 40 | ```python 41 | result = segment_clothing(img=image) 42 | result.save('segmented.png') 43 | ``` 44 | ![](https://cdn.discordapp.com/attachments/1120417968032063538/1202309847543185499/segmented-1.png?ex=65e8accd&is=65d637cd&hm=eed593adeca5b6d37ae2576499d5e142e4117f9c3f7bbd076d5cb575655e0efc&) 45 | 46 | You can also explicitly specify which clothes to segment 47 | - **img** input image of type PIL 48 | - **clothes** (optional) list of strings. by default ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"] 49 | ```python 50 | result = segment_clothing(img=image, clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]) 51 | result.save('segmented.png') 52 | ``` 53 | 54 | ### Batch Segment Clothing 55 | 56 | - **img_dir** image folder 57 | - **out_dir** output folder where the segmented images will go 58 | - **clothes** (optional) list of strings. by default ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"] 59 | ```python 60 | batch_segment_clothing(img_dir="images", out_dir="output", clothes= ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]) 61 | ``` 62 | --------------------------------------------------------------------------------