This tutorial provides a step-by-step guide on how to implement and train a UNet Tensorflow model for Melanoma detection using TensorFlow and Keras.
🔍 What You’ll Learn 🔍:
Building U-net model : Learn how to construct the model using TensorFlow and U-net Keras.
Unet Tensorflow Model Training: We’ll guide you through the training process, optimizing your model to generate masks in the lungs position
Testing and Evaluation: Run the pre-trained model on a new fresh images , and visual the test image next to the predicted mask .
Check out our tutorial here : https://youtu.be/-AejMcdeOOM?si=JibmFeIoXPaiVVZC
Link for the full code : https://ko-fi.com/s/3f1fdb4316
You can find more tutorials, and join my newsletter here : https://eranfeit.net/
This tutorial is based on the U-net Architecture :
Here is the code for segment X-Ray lungs using U-net and Tensorflow :
Dataset name – Academic Montgomery County X-ray Set
You can download the dataset here : https://academictorrents.com/details/ac786f74878a5775c81d490b23842fd4736bfe33
Part 1 : Dataset loading and Preprocessing :
This code processes the Montgomery County X-ray Dataset, which contains chest X-rays and their corresponding lung segmentation masks (left and right lungs separately).
The code loads these images, preprocesses them, and prepares them for training a segmentation Unet Tensorflow model. It performs tasks like resizing images, normalizing pixel values, combining left and right lung masks, and splitting the data into training and validation sets.
# Import required libraries import cv2 # OpenCV library for image processing import numpy as np # NumPy for numerical operations import glob # For finding files matching a pattern from tqdm import tqdm # For displaying progress bars # Define image dimensions for resizing Height = 256 Width = 256 # Define file paths for images and masks path = "E:/Data-sets/Lung Segmentation/MontgomerySet/" imagesPath = path + "CXR_png/*.png" # Path to X-ray images leftMaskPath = path + "ManualMask/leftMask/*.png" # Path to left lung masks rightMaskPath = path + "ManualMask/rightMask/*.png" # Path to right lung masks # Get lists of all image files using glob print ("Images in folder , left mask images , right mask images :") listOfImages = glob.glob(imagesPath) # Get list of X-ray images listOfLeftMaskImages = glob.glob(leftMaskPath) # Get list of left mask images listOfRightMaskImages = glob.glob(rightMaskPath) # Get list of right mask images # Print number of images found in each category print(len(listOfImages)) print(len(listOfLeftMaskImages)) print(len(listOfRightMaskImages)) # Load and display a sample image and its masks img = cv2.imread(listOfImages[0], cv2.IMREAD_COLOR) # Read first X-ray image print(img.shape) # Print original image dimensions # Resize X-ray image to target dimensions img = cv2.resize(img , (Width, Height)) # Load left and right lung masks for the first image left_mask = cv2.imread(listOfLeftMaskImages[0], cv2.IMREAD_GRAYSCALE) right_mask = cv2.imread(listOfRightMaskImages[0], cv2.IMREAD_GRAYSCALE) # Resize both masks to target dimensions left_mask = cv2.resize(left_mask , (Width, Height)) right_mask = cv2.resize(right_mask , (Width, Height)) # Combine left and right masks into a single mask finalMask = left_mask + right_mask # Display the images and masks cv2.imshow("img", img) cv2.imshow("left_mask", left_mask) cv2.imshow("right_mask", right_mask) cv2.imshow("finalMask", finalMask) cv2.waitKey(0) # Wait for key press before continuing # Demonstrate mask value processing mask16 = cv2.resize(left_mask, (16,16)) # Resize mask to smaller dimensions for demonstration print(mask16) # Print original mask values # Convert mask values to binary (0 and 1) mask16[mask16 > 0 ] = 1 print("======================================") print(mask16) # Print binary mask values # Initialize empty lists for storing processed images and masks allImages = [] maskImages = [] # Process all images and masks print("start loading the train images and masks + images augmentation X3") for imgFile , leftMask, rightMask in tqdm(zip(listOfImages, listOfLeftMaskImages, listOfRightMaskImages), total = len(listOfImages)): # Load and preprocess X-ray image img = cv2.imread(imgFile, cv2.IMREAD_COLOR) img = cv2.resize(img, (Width, Height)) # Resize image img = img / 255.0 # Normalize pixel values to range [0,1] img = img.astype(np.float32) # Convert to float32 allImages.append(img) # Load and process mask images leftMask = cv2.imread(leftMask, cv2.IMREAD_GRAYSCALE) rightMask = cv2.imread(rightMask, cv2.IMREAD_GRAYSCALE) mask = leftMask + rightMask # Combine masks mask = cv2.resize(mask, (Width, Height)) # Resize mask mask[mask>0] = 1 # Convert to binary mask maskImages.append(mask) # Convert lists to NumPy arrays allImagesNP = np.array(allImages) maskImagesNP = np.array(maskImages) maskImagesNP = maskImagesNP.astype(int) # Convert masks to integer type # Print shapes of processed data print("Shapes of train images and masks :") print(allImagesNP.shape) print(maskImagesNP.shape) # Split data into training and validation sets from sklearn.model_selection import train_test_split split = 0.1 # 10% validation split # Perform train-test split train_imgs, valid_imgs = train_test_split(allImagesNP, test_size=split , random_state=42) train_masks, valid_masks = train_test_split(maskImagesNP, test_size=split, random_state=42) # Print shapes of training and validation sets print("Shapes of train images and masks : ") print(train_imgs.shape) print(train_masks.shape) print("Shapes of Validat images and masks : ") print(valid_imgs.shape) print(valid_masks.shape) # Save processed data to NumPy files print("Save the data :") np.save("e:/temp/Unet-Train-Lung-Images.npy", train_imgs) np.save("e:/temp/Unet-Train-Lung-Masks.npy", train_masks) np.save("e:/temp/Unet-Validate-Lung-Images.npy", valid_imgs) np.save("e:/temp/Unet-Validate-Lung-Masks.npy", valid_masks) print("Finish save the data")
Link for the full code : https://ko-fi.com/s/3f1fdb4316
Part 2 : The Unet Tensoflow model (Save it as “Step02Model.py”) :
This code implements a U-Net Keras model using TensorFlow and Keras. U-Net is a popular architecture for image segmentation tasks, with an encoder-decoder structure and skip connections. The model takes an input image and produces a segmentation mask. The architecture consists of a convolutional block helper function and the main model-building function.
This code defines a U-Net architecture with:
- An encoder path that reduces spatial dimensions while increasing feature channels
- A bridge section at the bottom with the highest number of filters
- A decoder path that increases spatial dimensions while decreasing feature channels
- Skip connections that connect corresponding encoder and decoder levels
- A final output layer that produces a binary segmentation mask
The Unet Tensorflow model is particularly effective for medical image segmentation tasks, like the lung segmentation task it’s being used for here.
Key features:
- Uses repeated blocks of Conv2D, BatchNormalization, and ReLU activation
- Implements skip connections to preserve spatial information
- Uses max pooling for downsampling and upsampling for decoder path
- Produces a binary segmentation mask through sigmoid activation
- Follows the classic U-Net architecture pattern of contracting and expanding paths
# Import required TensorFlow and Keras libraries import tensorflow as tf from tensorflow.keras.layers import * # Import all Keras layers from tensorflow.keras.models import Model # Import Keras Model class # Define a convolutional block function that will be used throughout the network def conv_block(x, num_filters): # First convolutional layer x = Conv2D(num_filters, (3,3), padding="same")(x) # Apply 3x3 convolution x = BatchNormalization()(x) # Normalize the activations x = Activation("relu")(x) # Apply ReLU activation function # Second convolutional layer x = Conv2D(num_filters, (3,3), padding="same")(x) # Apply another 3x3 convolution x = BatchNormalization()(x) # Normalize the activations x = Activation("relu")(x) # Apply ReLU activation function return x # Define the main U-Net model building function def build_model(shape): # Define the number of filters for each level of the network num_filters = [64,128,256,512] # Filters double at each level # Create the input layer with specified shape inputs = Input((shape)) # List to store skip connections skip_x = [] x = inputs # Encoder path (Contracting path) for f in num_filters: x = conv_block(x,f) # Apply convolution block skip_x.append(x) # Store the output for skip connection x = MaxPool2D((2,2))(x) # Apply max pooling to reduce dimensions # Bridge (Bottom of the U) x = conv_block(x, 1024) # Apply convolution block with 1024 filters # Reverse the filter list and skip connections for the decoder path num_filters.reverse() # Reverse filter order for decoder skip_x.reverse() # Reverse skip connections order # Decoder path (Expanding path) for i, f in enumerate(num_filters): x = UpSampling2D((2,2))(x) # Upsample the feature map xs = skip_x[i] # Get corresponding skip connection x = Concatenate()([x,xs]) # Concatenate with skip connection x = conv_block(x,f) # Apply convolution block # Output layer x = Conv2D(1, (1,1), padding="same")(x) # 1x1 convolution to get final output x = Activation("sigmoid")(x) # Sigmoid activation for binary segmentation # Create and return the model return Model(inputs, x)
Link for the full code : https://ko-fi.com/s/3f1fdb4316
Part 3 : Train the Unet Tensorflow model :
This script loads preprocessed image data for training and validating a U-Net Keras model for lung segmentation.
It sets up a Unet TensorFlow based deep learning model, compiles it with an Adam optimizer and binary cross-entropy loss, and trains it using predefined hyperparameters while applying model checkpointing, learning rate reduction, and early stopping callbacks.
import numpy as np ### Import NumPy, a library for numerical computing, which is used for handling arrays. # load the data ### Print a message indicating the start of data loading. print("start lodaing ") ### Load the training images from a `.npy` file. allImagesNp = np.load("e:/temp/Unet-Train-Lung-Images.npy") ### Load the corresponding masks (segmentation labels) for training images. maskImagesNp = np.load("e:/temp/Unet-Train-Lung-Masks.npy") ### Load the validation images. allValidateImagesNP = np.load("e:/temp/Unet-Validate-Lung-Images.npy") ### Load the validation masks. maskValidateImagesNP = np.load("e:/temp/Unet-Validate-Lung-Masks.npy") ### Print the shapes of the loaded datasets to verify correct dimensions. print(allImagesNp.shape) print(maskImagesNp.shape) print(allValidateImagesNP.shape) print(maskValidateImagesNP.shape) ### Define the height and width of the input images. Height = 256 Width = 256 # build the model ### Import TensorFlow and necessary Keras utilities for defining and training the model. import tensorflow as tf ### Import `build_model` function from `Step02Model` to construct the U-Net model. from Step02Model import build_model ### Import Keras callbacks for saving the best model, adjusting learning rate, and stopping early. from keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping ### Define the input shape (256x256 with 3 color channels), learning rate, batch size, and epochs. shape = (256, 256, 3) lr = 1e-4 # 0.0001 batchSize = 4 epochs = 50 ### Build the U-Net model. model = build_model(shape) ### Print the model architecture summary. print(model.summary()) ### Use the Adam optimizer with the defined learning rate. opt = tf.keras.optimizers.Adam(lr) ### Compile the model with binary cross-entropy loss and accuracy as a metric. model.compile(loss="binary_crossentropy", optimizer=opt, metrics=['accuracy']) ### Compute the number of steps per epoch for training and validation. stepsPerEpoch = np.ceil(len(allImagesNp) / batchSize) validationSteps = np.ceil(len(allValidateImagesNP) / batchSize) ### Define the file path where the best model will be saved. best_model_file = "e:/temp/lung-Unet.h5" ### Set up callbacks: ### - Save the best model based on validation loss. ### - Reduce the learning rate if validation loss does not improve for 5 epochs. ### - Stop training if validation accuracy does not improve for 20 epochs. callbacks = [ ModelCheckpoint(best_model_file, verbose=1, save_best_only=True), ReduceLROnPlateau(monitor="val_loss", patience=5, factor=0.1, verbose=1, min_lr=1e-7), EarlyStopping(monitor="val_accuracy", patience=20, verbose=1) ] ### Train the model using the loaded data. ### Use mini-batches, shuffle data, validate after each epoch, and apply callbacks. history = model.fit( allImagesNp, maskImagesNp, batch_size=batchSize, epochs=epochs, verbose=1, validation_data=(allValidateImagesNP, maskValidateImagesNP), validation_steps=validationSteps, steps_per_epoch=stepsPerEpoch, shuffle=True, callbacks=callbacks )
Link for the full code : https://ko-fi.com/s/3f1fdb4316
Part 4 : Test the Unet Tensorflow model (inference) :
This script loads a pre-trained U-Net Keras model for lung segmentation, processes a test image, predicts the segmented lung region, and displays both the original image and the predicted mask.
It resizes the image, normalizes pixel values, applies a binary threshold to the mask, and visualizes the results using OpenCV.
import numpy as np import tensorflow as tf import cv2 ### Import necessary libraries: ### - NumPy for handling arrays. ### - TensorFlow for loading the trained model and making predictions. ### - OpenCV for image processing. # load the saved model ### Define the path to the best saved model. best_model_file = "e:/temp/lung-Unet.h5" ### Load the pre-trained U-Net model from the specified file. model = tf.keras.models.load_model(best_model_file) ### Print the model architecture summary. print(model.summary()) ### Define the width and height of the input image for resizing. Width = 256 Height = 256 # load test image ### Define the path to the test image. testImagePath = "TensorFlowProjects/Unet-Projects/Lung Segmentation/Lung-test-Image-From-Google.jpeg" ### Read the test image using OpenCV. img = cv2.imread(testImagePath) ### Resize the image to match the input size expected by the model. img2 = cv2.resize(img, (Width, Height)) ### Normalize the pixel values by dividing by 255 (scaling to the range [0,1]). img2 = img2 / 255.0 ### Expand dimensions to match the model's input shape (batch size of 1). imgForModel = np.expand_dims(img2, axis=0) ### Perform prediction using the trained model. p = model.predict(imgForModel) ### Extract the predicted segmentation mask. resultMask = p[0] ### Print the shape of the predicted mask to verify dimensions. print(resultMask.shape) # Process the prediction output ### Since this is a binary segmentation task: ### - Any value below or equal to 0.5 is set to 0 (black, no object detected). ### - Any value above 0.5 is set to 255 (white, object detected). resultMask[resultMask <= 0.5] = 0 resultMask[resultMask > 0.5] = 255 ### Define the scaling percentage for displaying images. scale_precent = 60 # Resize factor for display purposes. ### Compute the new width and height based on the scale percentage. w = int(img.shape[1] * scale_precent / 100) h = int(img.shape[0] * scale_precent / 100) ### Define the new image dimensions. dim = (w, h) ### Resize the original image for display. img = cv2.resize(img, dim, interpolation=cv2.INTER_AREA) ### Resize the predicted mask for display. mask = cv2.resize(resultMask, dim, interpolation=cv2.INTER_AREA) ### Display the original test image. cv2.imshow("first image", img) ### Display the predicted segmentation mask. cv2.imshow("Predicted mask", mask) ### Wait for a key press to close the displayed images. cv2.waitKey(0)
Link for the full code : https://ko-fi.com/s/3f1fdb4316
The Unet Tensorflow result :
Connect :
☕ Buy me a coffee — https://ko-fi.com/eranfeit
🖥️ Email : feitgemel@gmail.com
🤝 Fiverr : https://www.fiverr.com/s/mB3Pbb
Enjoy,
Eran