This project is a Python-based implementation of the U-Net architecture for image segmentation tasks. The U-Net model is a type of convolutional neural network that is widely used for biomedical image segmentation, achieving state-of-the-art results on a wide range of datasets.
These instructions will get you a copy of the project up and running on your local machine for development and testing purposes.
- Python
- PyTorch
- torchvision
- tqdm
Clone the repository to your local machine.
git clone <repository_link>
Install the required packages.
pip install -r requirements.txt
Train the model using the following command:
python train.py
model.py
: This file defines the U-Net model structure. The U-Net model is a fully convolutional network, consisting of an encoder (downsampling) part and a decoder (upsampling) part.data.py
: This file defines aMyDataset
class for loading and processing image data. This class inherits fromtorch.utils.data.Dataset
and implements the__getitem__
and__len__
methods.utils.py
: This file defines some helper functions for image processing.train.py
: This file defines the training process of the model. During training, the dataset is loaded, the model and optimizer are created, and for each batch of data, forward propagation is performed, the loss is calculated, backpropagation is performed, and the model parameters are updated.
- PyTorch - The deep learning framework used
- torchvision - Used to handle image data
- tqdm - Used to display progress bars
- Yan Qian