This repository contains the code for the U-Net architecture, a popular convolutional neural network used for image segmentation. The code is implemented in Python and uses the PyTorch framework.
To install the dependencies required for this code, run the following command:
pip install -r requirements.txt
To use the U-Net architecture, follow these steps:
- Prepare your data by organizing it into two folders: one for images and one for their corresponding masks.
- Update the `data_dir` variable in `train.py` and `test.py` to point to the location of your data.
- Run `train.py` to train the model. You can adjust the hyperparameters in this file to optimize performance.
- Run `test.py` to generate predictions on a new set of images.
This code was developed by Yash Bhootda. If you find this repository useful, please consider giving it a star on GitHub!