This repository contains a Jupyter Notebook (U_net_model_main.ipynb) that implements a U-Net model for image segmentation. The U-Net is a convolutional neural network architecture designed for biomedical image segmentation, known for its U-shaped architecture that allows it to capture context and localize precisely.
The Jupyter Notebook demonstrates the following:
- Data Loading and Preprocessing: Loads images and their corresponding masks, resizes them to a uniform size (128x128 pixels), and performs necessary preprocessing steps.
- U-Net Model Architecture: Defines and builds the U-Net model using TensorFlow/Keras, including the contracting path (encoder) and the expansive path (decoder) with skip connections.
- Model Training: Compiles the model with an Adam optimizer and binary cross-entropy loss, and trains it on the prepared dataset.
- Prediction and Visualization: Makes predictions on unseen data and visualizes the original images, ground truth masks, and predicted masks to evaluate the model's performance.
To run this notebook, you will need the following Python libraries:
tensorflowosrandomnumpytqdmskimage(specificallyskimage.ioforimread,imshowandskimage.transformforresize)matplotlib(specificallymatplotlib.pyplotfor plotting)
You can install these dependencies using pip:
pip install tensorflow os random numpy tqdm scikit-image matplotlibThe model is trained and evaluated on a dataset of images and their corresponding masks. Based on the notebook, the dataset structure is expected to be as follows:
/path/to/your/dataset/
├── stage1_train/
│ ├── <image_id_1>/
│ │ ├── images/
│ │ │ └── <image_id_1>.png
│ │ └── masks/
│ │ ├── <mask_id_1>.png
│ │ └── <mask_id_2>.png
│ │ └── ...
│ ├── <image_id_2>/
│ │ ├── images/
│ │ │ └── <image_id_2>.png
│ │ └── masks/
│ │ ├── <mask_id_1>.png
│ │ └── ...
│ └── ...
└── stage1_test/
├── <image_id_A>/
│ └── images/
│ └── <image_id_A>.png
└── ...
The notebook specifies the following paths for the dataset:
TRAIN_PATH = '/Users/rudra/Desktop/Projects/U-net/stage1_train/'TEST_PATH = '/Users/rudra/Desktop/Projects/U-net/stage1_test/'
Please ensure your dataset is organized in this manner and update the TRAIN_PATH and TEST_PATH variables in the notebook to point to the correct location of your data.