This project implements a U-Net model to perform lung segmentation on medical images (e.g., CT scans or X-rays). With convolutional layers and skip connections, the U-Net model accurately segments lung regions, supporting various medical imaging tasks and disease analysis.
- Custom Data Loader: Loads and preprocesses images and corresponding masks.
- Data Augmentation: Simple augmentations, such as horizontal flips, to increase data diversity.
- U-Net Architecture: Encoder-decoder structure with skip connections for detailed segmentation.
- Custom IoU Metric: Uses Intersection over Union (IoU) to evaluate segmentation accuracy.
- Training Callbacks: Includes model checkpointing, early stopping, and learning rate reduction.
- Visualization: Plots training history and displays sample predictions.
DataLoader
class: Loads images and masks, resizes to target size, and normalizes pixel values.SegmentationGenerator
class: Generates batches of images and masks with optional augmentation.simple_unet
function: Defines a U-Net model with an adjustable input size.iou_metric
function: Calculates the IoU score as a custom evaluation metric.plot_training_history
function: Plots training and validation loss and IoU metrics.predict_and_visualize
function: Predicts masks on the validation set and visualizes results.
- Python Libraries: Install the required libraries by running:
pip install tensorflow numpy opencv-python scikit-learn matplotlib
Store images and corresponding mask files in separate directories:
dataset/ ├── images/ └── masks/
Update image_dir
and mask_dir
in the main script to point to your dataset directories.
To execute the project, run the main script. It will:
- Load and preprocess data
- Train the U-Net model with defined callbacks
- Save the best model
- Plot training metrics and visualize sample predictions
python lung_segmentation.py
## Results
-The script will output:
-Saved model files (best_model.keras and final_model.keras)
-Plots of training and validation metrics
-Sample predictions showing the original image, true mask, and predicted mask