Last Updated on 08/10/2025 by Eran Feit
Deep Learning Image Segmentation with U-Net
This tutorial demonstrates a complete U-Net image segmentation workflow. It is designed as a practical image segmentation tutorial, showing how deep learning image segmentation can be applied to
Check out our tutorial here : https://youtu.be/ZiGMTFle7bw
The tutorial is divided into four parts:
Part 1: Data Preprocessing and Preparation
In this part, you load and preprocess the persons dataset, including resizing images and masks, converting masks to binary format, and splitting the data into training, validation, and testing sets.
Part 2: U-Net Model Architecture
This part defines the U-Net model architecture using Keras. It includes building blocks for convolutional layers, constructing the encoder and decoder parts of the U-Net, and defining the final output layer.
Part 3: Model Training
Here, you load the preprocessed data and train the U-Net model. You compile the model, define training parameters like learning rate and batch size, and use callbacks for model checkpointing, learning rate reduction, and early stopping.
Part 4: Model Evaluation and Inference
The final part demonstrates how to load the trained model, perform inference on test data, and visualize the predicted segmentation masks
You can find more tutorials, and join my newsletter here : https://eranfeit.net/
Link for the full code here : https://ko-fi.com/s/372080130c
Link for my blog : https://eranfeit.net/blog/
This tutorial is based on the U-net Architecture :
Here is the code for U-net Image Segmentation tutorial :
How to segment persons in images
Link for the dataset : https://www.kaggle.com/datasets/balraj98/cvcclinicdb
Part 1: Person Image Segmentation Data Preprocessing
This Python script prepares training and validation datasets for a person segmentation task using U-Net.
It loads images and corresponding segmentation masks, resizes them to a specified size, normalizes image pixel values,
and saves the processed data as .npy files for efficient use in deep learning models.
The dataset used is sourced from Kaggle and contains images of people with labeled segmentation masks.
# Import required libraries for image processing and data manipulation
import cv2
import numpy as np
import pandas as pd
# Define image dimensions for resizing
Height = 256
Width = 256
# Initialize empty lists to store images and masks for both training and validation sets
allImages = []
maskImages = []
allValidateImages = []
maskValidatImages = []
# Define file paths for the dataset
path = "e:/Data-sets/people_segmentation/"
imagesPath = path + "images"
maskPath = path + "masks"
TrainFile = path + "segmentation/train.txt"
validateFile = path + "segmentation/val.txt"
# Load training file list using pandas
df = pd.read_csv(TrainFile, sep=" ", header=None)
filesList = df[0].values
# Test loading process with a single image and mask
# Load and resize a sample image
img = cv2.imread(imagesPath+"/ache-adult-depression-expression-41253.jpg", cv2.IMREAD_COLOR)
img = cv2.resize(img, (Width, Height))
cv2.imshow("img", img)
# Load and process a sample mask
# Masks are binary: 0 for background, 1 for human
mask = cv2.imread(maskPath+"/ache-adult-depression-expression-41253.png", cv2.IMREAD_GRAYSCALE)
MASK16 = cv2.resize(mask , (16,16))
print(MASK16)
# Scale mask values to 255 for visualization
mask = cv2.resize(mask , (Width, Height))
mask = mask * 255
cv2.imshow("mask", mask)
cv2.waitKey(0)
# Process all training images and masks
print("Start loading the train images and masks ..............................")
for file in filesList:
# Construct file paths for each image and mask
filePathForImage = imagesPath + "/" +file + ".jpg"
filePathForMask = maskPath + "/" + file + ".png"
print(file)
# Load and preprocess image: resize and normalize to [0,1]
img = cv2.imread(filePathForImage , cv2.IMREAD_COLOR)
img = cv2.resize(img , (Width, Height))
img = img / 255.0
img = img.astype(np.float32)
allImages.append(img)
# Load and resize mask
mask = cv2.imread(filePathForMask, cv2.IMREAD_GRAYSCALE)
mask = cv2.resize(mask , (Width, Height))
maskImages.append(mask)
# Convert lists to numpy arrays and ensure proper data types
allImagesNP = np.array(allImages)
maskImagesNP = np.array(maskImages)
maskImagesNP = maskImagesNP.astype(int)
# Print shapes for verification
print ("Shapes of train images and masks :")
print(allImagesNP.shape)
print(maskImagesNP.shape)
print(maskImagesNP.dtype)
# Process validation images and masks
df = pd.read_csv(validateFile, sep=" ", header=None)
filesList = df[0].values
print("Start loading the Validate images and masks ..............................")
for file in filesList:
# Similar process as training data but for validation set
filePathForImage = imagesPath + "/" +file + ".jpg"
filePathForMask = maskPath + "/" + file + ".png"
print(file)
# Load and preprocess validation images
img = cv2.imread(filePathForImage , cv2.IMREAD_COLOR)
img = cv2.resize(img , (Width, Height))
img = img / 255.0
img = img.astype(np.float32)
allValidateImages.append(img)
# Load and process validation masks
mask = cv2.imread(filePathForMask, cv2.IMREAD_GRAYSCALE)
mask = cv2.resize(mask , (Width, Height))
maskValidatImages.append(mask)
# Convert validation data to numpy arrays
allValidateImagesNP = np.array(allValidateImages)
maskValidateImagesNP = np.array(maskValidatImages)
maskValidateImagesNP = maskValidateImagesNP.astype(int)
# Print validation set shapes
print ("Shapes of train images and masks :")
print(allValidateImagesNP.shape)
print(maskValidateImagesNP.shape)
print(maskValidateImagesNP.dtype)
# Save processed arrays to disk
print("Save the Data ......")
np.save("e:/temp/Unet-Human-Train-Images.npy", allImagesNP)
np.save("e:/temp/Unet-Human-Train-masks.npy", maskImagesNP)
np.save("e:/temp/Unet-Human-Validate-Images.npy", allValidateImagesNP)
np.save("e:/temp/Unet-Human-Validate-Masks.npy", maskValidateImagesNP)
print("Finish save the data .............")Link for the full code here : https://ko-fi.com/s/372080130c
Part 2: U-Net Architecture Implementation for Image Segmentation
This code implements a modified U-Net architecture, a specialized convolutional neural network designed for image segmentation tasks.
The implementation includes a compact version of the original U-Net with reduced filter sizes for efficiency while maintaining the characteristic encoder-decoder structure with skip connections.
The network is particularly optimized for binary segmentation tasks (like separating person from background) using sigmoid activation in the output layer.
The architecture consists of an encoder path that captures context, a bridge that processes the most compressed representation, and a decoder path that enables precise localization, with skip connections linking the encoder and decoder paths to preserve spatial information.
# Import required TensorFlow libraries and specific Keras layers
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model
# Define a convolutional block function that serves as a basic building block
# Each block contains two sets of Conv2D -> BatchNorm -> ReLU
def conv_block(x, num_filters):
# First convolutional layer with batch normalization and ReLU activation
x = Conv2D(num_filters, (3,3), padding="same")(x) # Apply 3x3 convolution
x = BatchNormalization()(x) # Normalize the outputs
x = Activation("relu")(x) # Apply ReLU activation
# Second convolutional layer with batch normalization and ReLU activation
x = Conv2D(num_filters, (3,3), padding="same")(x) # Repeat conv operation
x = BatchNormalization()(x) # Normalize again
x = Activation("relu")(x) # Apply ReLU activation
return x
# Main function to build the U-Net model
def build_model(shape):
# Define number of filters for each level (reduced from original U-Net for efficiency)
num_filters = [16,32,48,64] # Original U-Net used [64,128,256,512]
# Create input layer with specified shape
inputs = Input((shape))
# Initialize list to store skip connections
skip_x = []
x = inputs
# Encoder path: progressively reduce spatial dimensions while increasing filters
for f in num_filters:
x = conv_block(x, f) # Apply convolutional block
skip_x.append(x) # Store the output for skip connection
x = MaxPool2D((2,2))(x) # Reduce spatial dimensions by half
# Bridge: the bottommost layer that connects encoder to decoder
x = conv_block(x, 128) # Original U-Net used 1024 filters here
# Decoder path: progressively increase spatial dimensions while decreasing filters
num_filters.reverse() # Reverse filter list for decoder path
skip_x.reverse() # Reverse skip connections to match decoder levels
for i, f in enumerate(num_filters):
x = UpSampling2D((2,2))(x) # Double the spatial dimensions
xs = skip_x[i] # Get corresponding skip connection
x = Concatenate()([x,xs]) # Combine upsampled and skip features
x = conv_block(x,f) # Apply convolutional block
# Output layer
x = Conv2D(1, (1,1), padding="same")(x) # 1x1 convolution to get single channel
x = Activation("sigmoid")(x) # Sigmoid for binary segmentation
# Create and return the complete model
return Model(inputs, x)Link for the full code here : https://ko-fi.com/s/372080130c
The code implements a U-Net architecture with several key features:
- Modified filter sizes (16,32,48,64 instead of 64,128,256,512) for efficiency
- Consistent use of double convolutional blocks with batch normalization
- Skip connections to preserve spatial information
- Symmetric encoder-decoder structure
- Binary segmentation output using sigmoid activation
- Four levels of encoding/decoding with corresponding skip connections
- A bridge layer with 128 filters connecting encoder and decoder paths
This implementation is particularly suitable for tasks like person segmentation where the goal is to separate a subject from the background, offering a good balance between model capacity and computational efficiency.
Part 3: U-Net Model Training Pipeline for Person Segmentation
# Import numpy for array operations
import numpy as np
# Load the preprocessed training and validation data from saved numpy files
print("start loading the Train data ....... ")
allImagesNP = np.load("e:/temp/Unet-Human-Train-Images.npy")
maskImagesNP= np.load("e:/temp/Unet-Human-Train-masks.npy")
print("start loading the validate data ....... ")
allValidateImagesNP = np.load("e:/temp/Unet-Human-Validate-Images.npy")
maskValidateImagesNP = np.load("e:/temp/Unet-Human-Validate-Masks.npy")
print("Finish save the data .............")
# Print shapes of loaded arrays to verify dimensions
print(allImagesNP.shape)
print(maskImagesNP.shape)
print(allValidateImagesNP.shape)
print(maskValidateImagesNP.shape)
# Define image dimensions
Height = 256
Width = 256
# Import required libraries for model building and training
import tensorflow as tf
from Step02Model import build_model
from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
# Set up model parameters
shape = (256, 256, 3) # Input shape: 256x256 RGB images
lr = 1e-4 # Initial learning rate (0.0001)
batch_size = 8 # Number of images per batch
epochs = 50 # Maximum number of training epochs
# Build the U-Net model using imported function
model = build_model(shape)
print(model.summary()) # Display model architecture
# Configure the optimizer (Adam with specified learning rate)
opt = tf.keras.optimizers.Adam(lr)
# Compile model with binary cross-entropy loss (suitable for binary segmentation)
model.compile(loss="binary_crossentropy", optimizer=opt, metrics=['accuracy'])
# Calculate steps per epoch and validation steps based on batch size
stepsPerEpoch = np.ceil(len(allImagesNP) / batch_size)
validationSteps = np.ceil(len(allValidateImagesNP) / batch_size)
# Define path for saving the best model
best_model_file = "e:/temp/Human-Unet.h5"
# Set up training callbacks for optimization:
callbacks = [
# Save the best model based on validation performance
ModelCheckpoint(best_model_file, verbose=1, save_best_only=True),
# Reduce learning rate when validation loss plateaus
# Reduces by factor of 0.1 after 3 epochs without improvement
ReduceLROnPlateau(monitor="val_loss",
patience=3,
factor=0.1,
verbose=1,
min_lr=1e-6),
# Stop training if validation loss doesn't improve after 5 epochs
EarlyStopping(monitor="val_loss",
patience=5,
verbose=1)
]
# Start model training
history = model.fit(
allImagesNP, # Training images
maskImagesNP, # Training masks
batch_size=batch_size, # Batch size
epochs=epochs, # Maximum epochs
verbose=1, # Show progress
validation_data=(allValidateImagesNP, # Validation images
maskValidateImagesNP), # Validation masks
validation_steps=validationSteps, # Steps per validation
steps_per_epoch=stepsPerEpoch, # Steps per training epoch
shuffle=True, # Shuffle training data
callbacks=callbacks # Training callbacks
)Link for the full code here : https://ko-fi.com/s/372080130c
This code represents a sophisticated training pipeline with several key features:
- Efficient data loading from preprocessed numpy arrays
- Carefully tuned hyperparameters (learning rate, batch size, epochs)
- Advanced training optimization through callbacks:
- Model checkpointing to save best performing model
- Adaptive learning rate reduction to overcome plateaus
- Early stopping to prevent overfitting
- Binary cross-entropy loss suitable for segmentation
- Progress monitoring through accuracy metrics
- Data shuffling for better generalization
- Proper batch size handling with step calculations
This summary shows the architecture of the U-Net model, widely used in deep learning image
Part 4: Test the model – Person Segmentation Model Inference Pipeline
This script implements a complete inference pipeline for a trained U-Net segmentation model.
It demonstrates the practical application of the model by loading a pre-trained model and performing person segmentation on a single image.
The pipeline includes image preprocessing (resizing and normalization), model prediction, post-processing of the prediction mask (thresholding), and visualization of both the input image and resulting segmentation mask.
The script is particularly useful for testing the model’s performance on new images and provides visual feedback through OpenCV’s display capabilities.
# Import required libraries
import numpy as np
import tensorflow as tf
import cv2
# Load the trained model from saved file
best_model_file = "e:/temp/Human-Unet.h5"
model = tf.keras.models.load_model(best_model_file)
print(model.summary()) # Display model architecture
# Define image dimensions matching training configuration
Height = 256
Width = 256
# Define path and load test image
imgPath = "TensorFlowProjects/Unet-Projects/Human Image Segmentation/One-Human.jpg"
img = cv2.imread(imgPath, cv2.IMREAD_COLOR) # Load original image
# Preprocess image for model input
img2 = cv2.resize(img, (Width, Height)) # Resize to model's input dimensions
img2 = img2 / 255.0 # Normalize pixel values to [0,1]
# Add batch dimension as model expects 4D input (batch_size, height, width, channels)
imgForModel = np.expand_dims(img2, axis=0)
# Generate prediction using the model
p = model.predict(imgForModel)
resultMask = p[0] # Extract first (and only) mask from batch
print(resultMask.shape)
# Post-process the prediction mask
# Convert probabilistic output to binary mask using 0.5 threshold
# Values <= 0.5 become 0 (background)
# Values > 0.5 become 255 (person)
resultMask[resultMask <= 0.5] = 0
resultMask[resultMask > 0.5] = 255
# Calculate dimensions for display
# Reduce image size to 25% for better visualization
scale_precent = 25
width = int(img.shape[1] * scale_precent / 100)
height = int(img.shape[0] * scale_precent / 100)
dim = (width, height)
# Resize both original image and mask for display
img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA)
mask = cv2.resize(resultMask, dim, interpolation=cv2.INTER_AREA)
# Display results
cv2.imshow("image ", img) # Show original image
cv2.imshow("Mask", mask) # Show segmentation mask
cv2.waitKey(0) # Wait for key press
# Save the resulting mask
cv2.imwrite("e:/temp/testMask.png", mask)Link for the full code here : https://ko-fi.com/s/372080130c
This inference pipeline includes several key components:
- Model loading from saved weights
- Image preprocessing:
- Resizing to match model input requirements
- Normalization to [0,1] range
- Batch dimension addition
- Prediction generation and post-processing:
- Model inference
- Binary thresholding at 0.5
- Scaling to 8-bit range (0–255)
- Visualization:
- Image resizing for display
- Side-by-side comparison of input and output
- Mask saving functionality
- Professional error handling through proper OpenCV usage
- Efficient memory usage through appropriate array operations
The script serves as a practical tool for testing the segmentation model on new images and visualizing its performance in real-world applications.
“We successfully trained a U-Net to segment persons from the background
U-Net Image Segmentation Tutorial – Result :
Here is our test image and the predicted mask :
Connect :
☕ Buy me a coffee — https://ko-fi.com/eranfeit
🖥️ Email : feitgemel@gmail.com
🤝 Fiverr : https://www.fiverr.com/s/mB3Pbb
Planning a trip and want ideas you can copy fast?
Here are three detailed guides from our travels:
• 5-Day Ireland Itinerary: Cliffs, Castles, Pubs & Wild Atlantic Views
https://eranfeit.net/unforgettable-trip-to-ireland-full-itinerary/
• My Kraków Travel Guide: Best Places to Eat, Stay & Explore
https://eranfeit.net/my-krakow-travel-guide-best-places-to-eat-stay-explore/
• Northern Greece: Athens, Meteora, Tzoumerka, Ioannina & Nafpaktos (7 Days)
https://eranfeit.net/my-amazing-trip-to-greece/
Each guide includes maps, practical tips, and family-friendly stops—so you can plan in minutes, not hours.
Enjoy,
Eran