Semantic Segmentation with Pretrained U-Net in Keras + Albumentations 🔗

pip install tensorflow-metal 

# Install required packages %pip install -q --upgrade albumentationsx tensorflow tensorflow-datasets keras matplotlib tqdm jupytext 

import os import platform import warnings from typing import Tuple from dataclasses import dataclass  import numpy as np import tensorflow as tf import tensorflow_datasets as tfds import keras from keras import layers, models, optimizers, callbacks import albumentations as A import matplotlib.pyplot as plt from tqdm.keras import TqdmCallback from tqdm import tqdm  # Display plots inline %matplotlib inline  # Suppress warnings for cleaner output warnings.filterwarnings('ignore') os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Set random seeds for reproducibility np.random.seed(42) tf.random.set_seed(42)  print("Versions:") print(f" TensorFlow: {tf.__version__}") print(f" Keras: {keras.__version__}") print(f" Albumentations: {A.__version__}") print(f" Platform: {platform.system()} {platform.processor()}") 

def setup_device():  """Setup and detect available compute device."""  import platform  system = platform.system()  processor = platform.processor()    # Check for GPU availability  gpus = tf.config.list_physical_devices('GPU')    if gpus:  try:  # Enable memory growth for GPUs  for gpu in gpus:  tf.config.experimental.set_memory_growth(gpu, True)    # Determine if it's CUDA or MPS  if system == 'Darwin':  print("✓ Apple Silicon MPS GPU detected")  print(" TensorFlow will use Metal Performance Shaders")  return 'MPS'  else:  device_name = f"CUDA GPU ({len(gpus)} device(s))"  print(f"✓ Using {device_name}")  try:  gpu_details = tf.config.experimental.get_device_details(gpus[0])  print(f" GPU Name: {gpu_details.get('device_name', 'Unknown')}")  except:  pass  return 'GPU'  except RuntimeError as e:  print(f"GPU setup failed: {e}")    # Check if on Apple Silicon without GPU detected  if system == 'Darwin' and processor == 'arm':  print("⚠️ Apple Silicon detected but MPS not available!")  print(" To enable GPU acceleration, install tensorflow-metal:")  print(" pip install tensorflow-metal")  print(" Currently using CPU - training will be slower")  return 'CPU'    # Regular CPU fallback  print("ℹ Using CPU (no GPU detected)")  print(" Training will be slower. Consider using Google Colab with GPU.")  return 'CPU'  # Detect and setup device device_type = setup_device() 

@dataclass class Config:  """Configuration for training."""  # Model  backbone: str = 'resnet50' # Options: 'resnet34', 'resnet50', 'mobilenetv2', 'vgg16'  input_shape: Tuple[int, int, int] = (256, 256, 3)  num_classes: int = 3 # Background, pet, border    # Training  batch_size: int = 8 # Adjust based on your GPU memory  epochs: int = 10 # Pretrained models converge faster  learning_rate: float = 5e-4    # Dataset splits  train_split: str = "train[:80%]"  val_split: str = "train[80%:90%]"  test_split: str = "train[90%:]"    # Paths  checkpoint_dir: str = 'checkpoints'  # Create configuration config = Config()  print("Training Configuration:") print(f" Backbone: {config.backbone.upper()} (pretrained on ImageNet)") print(f" Input shape: {config.input_shape}") print(f" Batch size: {config.batch_size}") print(f" Epochs: {config.epochs}") print(f" Learning rate: {config.learning_rate}") print(f" Classes: {config.num_classes} (background, pet, border)")  # Adjust batch size based on device if device_type == 'CPU':  config.batch_size = 4  print(f" → Reduced batch size to {config.batch_size} for CPU") 

def build_unet_model(backbone_name: str, input_shape: Tuple, num_classes: int) -> keras.Model:  """  Build U-Net with pretrained backbone from ImageNet.    The U-Net architecture consists of:  1. Encoder: Pretrained backbone (ResNet, MobileNet, or VGG)  2. Decoder: Transposed convolutions with skip connections  3. Output: Softmax activation for multi-class segmentation  """  inputs = layers.Input(shape=input_shape)    # Select pretrained encoder based on backbone choice  if backbone_name == 'resnet34':  print("Loading ResNet34 pretrained on ImageNet...")  # Keras doesn't have ResNet34, so we use ResNet50V2  encoder = tf.keras.applications.ResNet50V2(  input_tensor=inputs,  weights='imagenet',  include_top=False  )  skip_layer_names = ['conv1_conv', 'conv2_block3_1_relu',   'conv3_block4_1_relu', 'conv4_block6_1_relu']  print(" (Using ResNet50V2 as ResNet34 equivalent)")    elif backbone_name == 'resnet50':  print("Loading ResNet50 pretrained on ImageNet...")  encoder = tf.keras.applications.ResNet50(  input_tensor=inputs,  weights='imagenet',  include_top=False  )  skip_layer_names = ['conv1_relu', 'conv2_block3_out',   'conv3_block4_out', 'conv4_block6_out']    elif backbone_name == 'mobilenetv2':  print("Loading MobileNetV2 pretrained on ImageNet...")  encoder = tf.keras.applications.MobileNetV2(  input_tensor=inputs,  weights='imagenet',  include_top=False  )  skip_layer_names = ['block_1_expand_relu', 'block_3_expand_relu',  'block_6_expand_relu', 'block_13_expand_relu']    elif backbone_name == 'vgg16':  print("Loading VGG16 pretrained on ImageNet...")  encoder = tf.keras.applications.VGG16(  input_tensor=inputs,  weights='imagenet',  include_top=False  )  skip_layer_names = ['block1_pool', 'block2_pool',   'block3_pool', 'block4_pool']  else:  raise ValueError(f"Unknown backbone: {backbone_name}")    # Make encoder trainable for fine-tuning  encoder.trainable = True    # Get skip connections from encoder layers  skip_layers = [encoder.get_layer(name).output for name in skip_layer_names]    # Build U-Net decoder  x = encoder.output  decoder_filters = [256, 128, 64, 32]    # Decoder blocks with skip connections  for i, (skip, filters) in enumerate(zip(reversed(skip_layers), decoder_filters)):  # Upsampling  x = layers.Conv2DTranspose(filters, 3, strides=2, padding='same')(x)  x = layers.BatchNormalization()(x)  x = layers.Activation('relu')(x)    # Skip connection (concatenate)  x = layers.Concatenate()([x, skip])    # Double convolution block  x = layers.Conv2D(filters, 3, padding='same')(x)  x = layers.BatchNormalization()(x)  x = layers.Activation('relu')(x)    x = layers.Conv2D(filters, 3, padding='same')(x)  x = layers.BatchNormalization()(x)  x = layers.Activation('relu')(x)    # Dropout for regularization (except last block)  if filters > 32:  x = layers.Dropout(0.3)(x)    # Final upsampling to match input resolution  x = layers.Conv2DTranspose(16, 3, strides=2, padding='same')(x)  x = layers.BatchNormalization()(x)  x = layers.Activation('relu')(x)    # Output layer with softmax for multi-class segmentation  outputs = layers.Conv2D(num_classes, 1, activation='softmax')(x)    model = models.Model(inputs=inputs, outputs=outputs, name=f'{backbone_name}_unet')  print(f"Model created: {model.count_params():,} parameters")    return model 

def create_augmentations(training: bool = True, backbone_name: str = 'resnet50'):  """Create Albumentations pipeline optimized for segmentation.    Key principles for segmentation augmentations:  - Avoid aggressive geometric distortions (no ElasticTransform, GridDistortion)  - Use conservative spatial transforms (mild rotation, scaling)  - Apply color augmentations for robustness  - Include normalization as the last step  """    # Get normalization parameters for each backbone  if backbone_name in ['resnet34', 'resnet50', 'vgg16']:  # Standard ImageNet normalization  mean = (0.485, 0.456, 0.406)  std = (0.229, 0.224, 0.225)  elif backbone_name == 'mobilenetv2':  # MobileNetV2 normalizes to [-1, 1]  mean = (0.5, 0.5, 0.5)  std = (0.5, 0.5, 0.5)  else:  # Default ImageNet normalization  mean = (0.485, 0.456, 0.406)  std = (0.229, 0.224, 0.225)    if training:  return A.Compose([  # Spatial augmentations (conservative for segmentation)  A.RandomCrop(height=256, width=256, pad_if_needed=True),  A.HorizontalFlip(p=0.5),  A.Affine(  scale=(0.5, 2), # Moderate scaling  rotate=(-15, 15), # Mild rotation  balanced_scale=True,  keep_ratio=True,  p=0.5  ),    # Color augmentations  A.ColorJitter(  brightness=[0.8, 1.2],  contrast=[0.8, 1.2],  saturation=[0.8, 1.2],  hue=[-0.5, 0.5],  p=1  ), # Random brightness, contrast, saturation, hue    # Noise and blur (mild)  A.OneOf([  A.GaussNoise(std_range=(0.1, 0.2), p=1),  A.GaussianBlur(blur_limit=(3, 5), p=1),  ], p=0.2),    # Normalization (must be last!)  A.Normalize(mean=mean, std=std),  ], seed=137, strict=True) # seed for reproducibility  else:  return A.Compose([  # For validation/test: only resize and normalize  A.CenterCrop(height=256, width=256, pad_if_needed=True, p=1),  A.Normalize(mean=mean, std=std),  ])  # Preview augmentation settings print("Augmentation Pipeline:") print(" Training:") print(" - RandomCrop with padding") print(" - HorizontalFlip (50%)") print(" - Affine transforms (scale, rotate)") print(" - ColorJitter") print(" - Noise/Blur (20%)") print(" - Normalization") print(" Validation:") print(" - CenterCrop with padding") print(" - Normalization") 

def dice_coefficient(y_true, y_pred, smooth=1.0):  """Dice coefficient metric for segmentation.    Dice = 2 * |A ∩ B| / (|A| + |B|)  Range: [0, 1] where 1 is perfect overlap  """  y_true_f = tf.reshape(y_true, [-1, tf.shape(y_true)[-1]])  y_pred_f = tf.reshape(y_pred, [-1, tf.shape(y_pred)[-1]])  intersection = tf.reduce_sum(y_true_f * y_pred_f, axis=0)  union = tf.reduce_sum(y_true_f, axis=0) + tf.reduce_sum(y_pred_f, axis=0)  dice = tf.reduce_mean((2. * intersection + smooth) / (union + smooth))  return dice  def dice_loss(y_true, y_pred):  """Dice loss for training."""  return 1.0 - dice_coefficient(y_true, y_pred)  def combined_loss(y_true, y_pred):  """Combined dice + categorical crossentropy loss.    This combination helps with:  - Dice: Handles class imbalance  - CrossEntropy: Provides stable gradients  """  cce = tf.keras.losses.categorical_crossentropy(y_true, y_pred)  return dice_loss(y_true, y_pred) + tf.reduce_mean(cce)  class MeanIoU(tf.keras.metrics.MeanIoU):  """Mean IoU metric for one-hot encoded masks."""  def update_state(self, y_true, y_pred, sample_weight=None):  # Convert from one-hot to class indices  y_true = tf.argmax(y_true, axis=-1)  y_pred = tf.argmax(y_pred, axis=-1)  return super().update_state(y_true, y_pred, sample_weight)  print("Loss and Metrics:") print(" Loss: Combined (Dice + Categorical Crossentropy)") print(" Metrics: Dice Coefficient, Mean IoU, Categorical Accuracy") 

def load_dataset(config: Config):  """Load Oxford-IIIT Pet dataset.    The dataset contains:  - Images of pets (cats and dogs)  - Segmentation masks with 3 classes:  0: Background  1: Pet  2: Border  """  print("Loading Oxford-IIIT Pet dataset...")  print("This may take a few minutes on first run...")    with tqdm(total=3, desc="Loading splits", unit="split") as pbar:  (ds_train, ds_val, ds_test), ds_info = tfds.load(  'oxford_iiit_pet',  split=[config.train_split, config.val_split, config.test_split],  with_info=True  )  pbar.update(3)    print(f"\n✓ Dataset loaded successfully")  print(f" Train samples: {len(ds_train)}")  print(f" Val samples: {len(ds_val)}")  print(f" Test samples: {len(ds_test)}")    return ds_train, ds_val, ds_test, ds_info  # Load the dataset ds_train, ds_val, ds_test, ds_info = load_dataset(config) 

def preprocess_data(data_dict):  """Preprocess raw data from TensorFlow datasets."""  image = tf.cast(data_dict['image'], tf.float32)  mask = data_dict['segmentation_mask']    # Convert mask to int32 and ensure 2D  mask = tf.cast(mask, tf.int32)  mask = tf.squeeze(mask) # Remove any extra dimensions    # Remap mask values for cleaner classes  # Original: 1=pet, 2=border, 3=background+border  # New: 0=background, 1=pet, 2=border  mask = tf.where(mask == 2, 0, mask) # border -> background  mask = tf.where(mask == 3, 2, mask) # background+border -> border    return image, mask  def create_data_pipeline(dataset, config: Config, augmentations, training: bool = True):  """Create tf.data pipeline with augmentations."""    def augment_fn(image, mask):  """Apply Albumentations augmentations."""  def aug(img, msk):  img = img.numpy().astype(np.uint8)  msk = msk.numpy().astype(np.uint8)    # Apply augmentations (including normalization)  augmented = augmentations(image=img, mask=msk)    return augmented['image'], augmented['mask'].astype(np.int32)    # Use tf.py_function to apply numpy-based augmentations  aug_img, aug_mask = tf.py_function(  aug, [image, mask], [tf.float32, tf.int32]  )    # Set shapes (required after py_function)  aug_img.set_shape([config.input_shape[0], config.input_shape[1], 3])  aug_mask.set_shape([config.input_shape[0], config.input_shape[1]])    # One-hot encode the mask for multi-class segmentation  aug_mask = tf.one_hot(aug_mask, config.num_classes)    return aug_img, aug_mask    # Build pipeline  dataset = dataset.map(preprocess_data, num_parallel_calls=tf.data.AUTOTUNE)  dataset = dataset.map(augment_fn, num_parallel_calls=tf.data.AUTOTUNE)    if training:  dataset = dataset.shuffle(buffer_size=1000)    dataset = dataset.batch(config.batch_size)  dataset = dataset.prefetch(tf.data.AUTOTUNE)    return dataset  # Create augmentation pipelines train_aug = create_augmentations(training=True, backbone_name=config.backbone) val_aug = create_augmentations(training=False, backbone_name=config.backbone)  # Create data pipelines print("\nCreating data pipelines...") train_dataset = create_data_pipeline(ds_train, config, train_aug, training=True) val_dataset = create_data_pipeline(ds_val, config, val_aug, training=False) test_dataset = create_data_pipeline(ds_test, config, val_aug, training=False) print("✓ Data pipelines ready") 

def denormalize_image(img, backbone_name='resnet50'):  """Properly denormalize image for visualization."""  # Get normalization parameters  if backbone_name in ['resnet34', 'resnet50', 'vgg16']:  # Standard ImageNet normalization  mean = np.array([0.485, 0.456, 0.406])  std = np.array([0.229, 0.224, 0.225])  elif backbone_name == 'mobilenetv2':  # MobileNetV2 normalizes to [-1, 1]  mean = np.array([0.5, 0.5, 0.5])  std = np.array([0.5, 0.5, 0.5])  else:  # Default ImageNet normalization  mean = np.array([0.485, 0.456, 0.406])  std = np.array([0.229, 0.224, 0.225])    # Denormalize: x = (x_norm * std) + mean  img_denorm = (img * std) + mean  # Clip to [0, 1] range for valid image  img_denorm = np.clip(img_denorm, 0, 1)  return img_denorm  def visualize_samples(dataset, num_samples=3):  """Visualize images and masks from dataset."""  fig, axes = plt.subplots(num_samples, 3, figsize=(12, num_samples * 3))    if num_samples == 1:  axes = axes.reshape(1, -1)    # Get a batch of data  for images, masks in dataset.take(1):  for i in range(min(num_samples, images.shape[0])):  # Properly denormalize image for visualization  img = images[i].numpy()  img = denormalize_image(img, config.backbone)    # Convert one-hot mask to class indices  mask = tf.argmax(masks[i], axis=-1).numpy()    # Create colored mask for better visualization  colored_mask = np.zeros((*mask.shape, 3))  colored_mask[mask == 0] = [0, 0, 0] # Background: black  colored_mask[mask == 1] = [0, 1, 0] # Pet: green  colored_mask[mask == 2] = [1, 0, 0] # Border: red    # Plot image  axes[i, 0].imshow(img)  axes[i, 0].set_title("Input Image")  axes[i, 0].axis('off')    # Plot mask  axes[i, 1].imshow(colored_mask)  axes[i, 1].set_title("Segmentation Mask")  axes[i, 1].axis('off')    # Plot overlay  axes[i, 2].imshow(img)  axes[i, 2].imshow(colored_mask, alpha=0.5)  axes[i, 2].set_title("Overlay")  axes[i, 2].axis('off')    plt.suptitle("Training Samples with Albumentations Augmentations", fontsize=14)  plt.tight_layout()  plt.show()  # Visualize training samples print("Training samples with augmentations:") visualize_samples(train_dataset, num_samples=3) 

# Build the model print(f"Building {config.backbone.upper()} U-Net...\n") model = build_unet_model(config.backbone, config.input_shape, config.num_classes)  # Compile the model model.compile(  optimizer=optimizers.Adam(learning_rate=config.learning_rate),  loss=combined_loss,  metrics=[  dice_coefficient,  MeanIoU(num_classes=config.num_classes, name='mean_iou'),  'categorical_accuracy'  ] )  print("\n📊 Model Summary:") print(f" Architecture: U-Net with {config.backbone.upper()} encoder") print(f" Total parameters: {model.count_params():,}") print(f" Loss: Combined (Dice + Categorical Crossentropy)") print(f" Optimizer: Adam (lr={config.learning_rate})") 

# Setup callbacks os.makedirs(config.checkpoint_dir, exist_ok=True)  checkpoint = callbacks.ModelCheckpoint(  filepath=os.path.join(config.checkpoint_dir, f'best_{config.backbone}_unet.keras'),  monitor='val_dice_coefficient',  mode='max',  save_best_only=True,  verbose=0 )  early_stop = callbacks.EarlyStopping(  monitor='val_dice_coefficient',  mode='max',  patience=7,  restore_best_weights=True,  verbose=0 )  reduce_lr = callbacks.ReduceLROnPlateau(  monitor='val_dice_coefficient',  mode='max',  factor=0.5,  patience=3,  min_lr=1e-7,  verbose=0 )  # Create tqdm callback for progress bars tqdm_callback = TqdmCallback(verbose=2, leave=True)  # Train the model print(f"\n🚂 Starting training for {config.epochs} epochs...") print(f" Backbone: {config.backbone.upper()} (pretrained)") print(f" Batch size: {config.batch_size}") print(f" Learning rate: {config.learning_rate}\n")  history = model.fit(  train_dataset,  validation_data=val_dataset,  epochs=config.epochs,  callbacks=[checkpoint, early_stop, reduce_lr, tqdm_callback],  verbose=0 # Disable default progress bar (using tqdm instead) )  print("\n✅ Training complete!")  # Print best metrics best_dice = max(history.history['val_dice_coefficient']) best_iou = max(history.history['val_mean_iou']) print(f"\n📊 Best validation metrics:") print(f" Dice coefficient: {best_dice:.4f}") print(f" Mean IoU: {best_iou:.4f}") 

# Plot training history fig, axes = plt.subplots(1, 3, figsize=(15, 4))  # Loss axes[0].plot(history.history['loss'], label='Train', linewidth=2) axes[0].plot(history.history['val_loss'], label='Val', linewidth=2) axes[0].set_title('Loss') axes[0].set_xlabel('Epoch') axes[0].set_ylabel('Loss') axes[0].legend() axes[0].grid(True, alpha=0.3)  # Dice Coefficient axes[1].plot(history.history['dice_coefficient'], label='Train', linewidth=2) axes[1].plot(history.history['val_dice_coefficient'], label='Val', linewidth=2) axes[1].set_title('Dice Coefficient') axes[1].set_xlabel('Epoch') axes[1].set_ylabel('Dice') axes[1].legend() axes[1].grid(True, alpha=0.3)  # Mean IoU axes[2].plot(history.history['mean_iou'], label='Train', linewidth=2) axes[2].plot(history.history['val_mean_iou'], label='Val', linewidth=2) axes[2].set_title('Mean IoU') axes[2].set_xlabel('Epoch') axes[2].set_ylabel('IoU') axes[2].legend() axes[2].grid(True, alpha=0.3)  plt.suptitle(f'Training History - {config.backbone.upper()} U-Net', fontsize=14) plt.tight_layout() plt.show()  # Print final metrics print(f"\nFinal Training Metrics (Epoch {len(history.history['loss'])}):") print(f" Train Dice: {history.history['dice_coefficient'][-1]:.4f}") print(f" Val Dice: {history.history['val_dice_coefficient'][-1]:.4f}") print(f" Train IoU: {history.history['mean_iou'][-1]:.4f}") print(f" Val IoU: {history.history['val_mean_iou'][-1]:.4f}") 

# Evaluate on test set print("🔍 Evaluating on test set...")  test_results = model.evaluate(  test_dataset,  verbose=0,  callbacks=[TqdmCallback(verbose=1, leave=True)] )  print("\n📊 Test Set Results:") print(f" Loss: {test_results[0]:.4f}") print(f" Dice Coefficient: {test_results[1]:.4f}") print(f" Mean IoU: {test_results[2]:.4f}") print(f" Categorical Accuracy: {test_results[3]:.4f}")  # Performance interpretation if test_results[1] > 0.8:  print("\n✨ Excellent performance! The pretrained backbone works great!") elif test_results[1] > 0.7:  print("\n✅ Good performance! Transfer learning is effective!") elif test_results[1] > 0.6:  print("\n📈 Decent results. Consider training for more epochs.") else:  print("\n⚠️ Room for improvement. Try adjusting hyperparameters or augmentations.") 

def visualize_predictions(model, dataset, num_samples=3):  """Visualize model predictions vs ground truth."""  fig, axes = plt.subplots(num_samples, 4, figsize=(15, num_samples * 3.5))    if num_samples == 1:  axes = axes.reshape(1, -1)    # Get predictions  for images, masks_true in dataset.take(1):  masks_pred = model.predict(images, verbose=0)    for i in range(min(num_samples, images.shape[0])):  # Denormalize image  img = images[i].numpy()  img = denormalize_image(img, config.backbone)    # Convert masks to class indices  mask_true = tf.argmax(masks_true[i], axis=-1).numpy()  mask_pred = tf.argmax(masks_pred[i], axis=-1).numpy()    # Create colored masks  def create_colored_mask(mask):  colored = np.zeros((*mask.shape, 3))  colored[mask == 0] = [0, 0, 0] # Background: black  colored[mask == 1] = [0, 1, 0] # Pet: green  colored[mask == 2] = [1, 0, 0] # Border: red  return colored    colored_true = create_colored_mask(mask_true)  colored_pred = create_colored_mask(mask_pred)    # Calculate accuracy for this sample  accuracy = np.mean(mask_true == mask_pred)    # Plot  axes[i, 0].imshow(img)  axes[i, 0].set_title("Input Image")  axes[i, 0].axis('off')    axes[i, 1].imshow(colored_true)  axes[i, 1].set_title("Ground Truth")  axes[i, 1].axis('off')    axes[i, 2].imshow(colored_pred)  axes[i, 2].set_title(f"Prediction (Acc: {accuracy:.2%})")  axes[i, 2].axis('off')    axes[i, 3].imshow(img)  axes[i, 3].imshow(colored_pred, alpha=0.5)  axes[i, 3].set_title("Overlay")  axes[i, 3].axis('off')    plt.suptitle(f'Model Predictions - {config.backbone.upper()} U-Net', fontsize=14)  plt.tight_layout()  plt.show()  # Visualize predictions print("Test set predictions:") visualize_predictions(model, test_dataset, num_samples=3) 

# Save the model model_path = os.path.join(config.checkpoint_dir, f'final_{config.backbone}_unet.keras') model.save(model_path) print(f"✅ Model saved to: {model_path}") print(f" Size: {os.path.getsize(model_path) / (1024*1024):.2f} MB")  # To load the model later: print("\nTo load this model later, use:") print(f"model = keras.models.load_model('{model_path}', custom_objects={{") print(" 'dice_coefficient': dice_coefficient,") print(" 'combined_loss': combined_loss,") print(" 'MeanIoU': MeanIoU") print("})")