This repository focuses on human image segmentation using PyTorch and the U-Net architecture, which is a popular choice for such tasks. The main objective is to accurately identify areas where humans are present and generate corresponding segmentation masks, marking these areas as white while the rest of the image remains black.
The project follows a structured approach, outlined as follows:
-
Configuration Setup: Define configurations essential for the project and visualize sample images to understand the dataset.
-
Dataset Splitting: Split the dataset into training and validation sets to facilitate model training and evaluation.
-
Data Augmentation: Implement data augmentation techniques tailored for both the training and validation sets. In segmentation tasks, augmentation applies to both images and their corresponding masks.
-
Custom Dataset Creation: Create a custom dataset to handle image-mask pairs, ensuring proper alignment between images and their masks.
-
Data Loading: Utilize DataLoader to efficiently load the dataset into batches for training.
-
Model Building: Build the segmentation model using the U-Net architecture. We'll leverage the Segmentation Models library in PyTorch, exploring various loss functions like Dice Loss and BCE with Logits Loss.
-
Training and Validation: Define training and validation functions, and implement the training loop to train the model on the dataset.
-
Model Deployment: Deploy the best-performing model to create an interface for segmentation tasks, enabling easy use in real-world scenarios.
- Python 3.x
- PyTorch
- Segmentation Models Library
- NumPy
- Pandas
Contributions are welcome! Please fork the repository, make your changes, and submit a pull request.
For any inquiries or support, please contact [email protected].