Data Augmentation for Geospatial Training¶

This notebook demonstrates how to use data augmentation when preparing training tiles and training segmentation models. Data augmentation helps improve model generalization by creating variations of training data through transformations like flips, rotations, and photometric adjustments.

Key Features¶

• Tile Export with Augmentation: Generate augmented versions of tiles during export
• Default Augmentation Transforms: Use pre-configured transforms optimized for remote sensing
• Custom Augmentation: Define your own augmentation pipeline
• Enhanced Training Defaults: Improved default augmentations for segmentation model training

Install package¶

Uncomment the following line to install the geoai package if needed:

In [ ]:

# %pip install -U geoai

# %pip install -U geoai

Import libraries¶

In [ ]:

import os
import geoai
import numpy as np
import rasterio
import matplotlib.pyplot as plt
from pathlib import Path

import os import geoai import numpy as np import rasterio import matplotlib.pyplot as plt from pathlib import Path

Download sample data¶

We'll use the same NAIP imagery and building footprint dataset as used in the train_segmentation_model.ipynb example. This is real aerial imagery with building annotations.

In [ ]:

# Download NAIP imagery and building footprints
train_raster_url = (
    "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_rgb_train.tif"
)
train_vector_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_train_buildings.geojson"

sample_image = geoai.download_file(train_raster_url)
sample_vector = geoai.download_file(train_vector_url)

print(f"Downloaded sample image: {sample_image}")
print(f"Downloaded sample labels: {sample_vector}")

# Download NAIP imagery and building footprints train_raster_url = ( "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_rgb_train.tif" ) train_vector_url = "https://huggingface.co/datasets/giswqs/geospatial/resolve/main/naip_train_buildings.geojson" sample_image = geoai.download_file(train_raster_url) sample_vector = geoai.download_file(train_vector_url) print(f"Downloaded sample image: {sample_image}") print(f"Downloaded sample labels: {sample_vector}")

Visualize sample data¶

Let's visualize the NAIP imagery and building footprints.

In [ ]:

print("Image information:")
geoai.get_raster_info(sample_image)

print("Image information:") geoai.get_raster_info(sample_image)

In [ ]:

# Visualize on interactive map
geoai.view_vector_interactive(sample_vector, tiles=sample_image)

# Visualize on interactive map geoai.view_vector_interactive(sample_vector, tiles=sample_image)

Part 1: Export Tiles WITHOUT Augmentation (Baseline)¶

First, let's export tiles without augmentation to establish a baseline.

In [ ]:

# Create output directory
output_dir = "buildings_augmentation_demo"

# Export without augmentation
output_no_aug = f"{output_dir}/tiles_no_augmentation"

geoai.export_geotiff_tiles(
    sample_image,
    output_no_aug,
    in_class_data=sample_vector,
    tile_size=256,
    stride=128,
    apply_augmentation=False,  # No augmentation
)

# Count tiles
image_tiles = list(Path(output_no_aug, "images").glob("*.tif"))
label_tiles = list(Path(output_no_aug, "labels").glob("*.tif"))

print(f"\nWithout augmentation:")
print(f"  Image tiles: {len(image_tiles)}")
print(f"  Label tiles: {len(label_tiles)}")

# Create output directory output_dir = "buildings_augmentation_demo" # Export without augmentation output_no_aug = f"{output_dir}/tiles_no_augmentation" geoai.export_geotiff_tiles( sample_image, output_no_aug, in_class_data=sample_vector, tile_size=256, stride=128, apply_augmentation=False, # No augmentation ) # Count tiles image_tiles = list(Path(output_no_aug, "images").glob("*.tif")) label_tiles = list(Path(output_no_aug, "labels").glob("*.tif")) print(f"\nWithout augmentation:") print(f" Image tiles: {len(image_tiles)}") print(f" Label tiles: {len(label_tiles)}")

Part 2: Export Tiles WITH Default Augmentation¶

Now let's export tiles with default augmentation. The get_default_augmentation_transforms() function provides sensible defaults for remote sensing data:

Geometric Transforms:

• Horizontal/Vertical Flips (50% probability each)
• Random 90° Rotations (50% probability)
• Shift-Scale-Rotate (50% probability)

Photometric Transforms:

• Random Brightness/Contrast (50% probability)
• HSV Color Adjustments (30% probability)
• Gaussian Noise (20% probability)
• Gaussian Blur (20% probability)

In [ ]:

# Export with default augmentation - generate 3 augmented versions per tile
output_with_aug = f"{output_dir}/tiles_with_augmentation"

geoai.export_geotiff_tiles(
    sample_image,
    output_with_aug,
    in_class_data=sample_vector,
    tile_size=256,
    stride=128,
    apply_augmentation=True,  # Enable augmentation
    augmentation_count=3,  # Generate 3 augmented versions per tile
)

# Count tiles
aug_image_tiles = list(Path(output_with_aug, "images").glob("*.tif"))
aug_label_tiles = list(Path(output_with_aug, "labels").glob("*.tif"))

print(f"\nWith augmentation (3 per tile):")
print(f"  Image tiles: {len(aug_image_tiles)} (original + augmented)")
print(f"  Label tiles: {len(aug_label_tiles)} (original + augmented)")
print(f"  \nThis is {len(aug_image_tiles) / len(image_tiles):.1f}x more training data!")

# Export with default augmentation - generate 3 augmented versions per tile output_with_aug = f"{output_dir}/tiles_with_augmentation" geoai.export_geotiff_tiles( sample_image, output_with_aug, in_class_data=sample_vector, tile_size=256, stride=128, apply_augmentation=True, # Enable augmentation augmentation_count=3, # Generate 3 augmented versions per tile ) # Count tiles aug_image_tiles = list(Path(output_with_aug, "images").glob("*.tif")) aug_label_tiles = list(Path(output_with_aug, "labels").glob("*.tif")) print(f"\nWith augmentation (3 per tile):") print(f" Image tiles: {len(aug_image_tiles)} (original + augmented)") print(f" Label tiles: {len(aug_label_tiles)} (original + augmented)") print(f" \nThis is {len(aug_image_tiles) / len(image_tiles):.1f}x more training data!")

Visualize Original vs Augmented Tiles¶

Let's compare an original tile with its augmented versions to see the transformations.

In [ ]:

# Load original tile and some augmented versions
tile_files = sorted([f for f in Path(output_with_aug, "images").glob("*.tif")])
label_files = sorted([f for f in Path(output_with_aug, "labels").glob("*.tif")])

# Get first 4 tiles (1 original + 3 augmented)
n_display = min(4, len(tile_files))

fig, axes = plt.subplots(2, n_display, figsize=(15, 8))
if n_display == 1:
    axes = axes.reshape(2, 1)

for i in range(n_display):
    # Load and display image tile
    with rasterio.open(tile_files[i]) as src:
        img = src.read()  # Read all bands
        # If RGB, display as color
        if img.shape[0] >= 3:
            img_display = np.transpose(img[:3], (1, 2, 0))
            axes[0, i].imshow(img_display)
        else:
            axes[0, i].imshow(img[0], cmap="gray")

        title = "Original" if i == 0 else f"Augmented {i}"
        axes[0, i].set_title(f"{title}\n{tile_files[i].name}")
        axes[0, i].axis("off")

    # Load and display label tile
    with rasterio.open(label_files[i]) as src:
        label = src.read(1)
        axes[1, i].imshow(label, cmap="tab10", vmin=0, vmax=10)
        axes[1, i].set_title(f"Label\n{label_files[i].name}")
        axes[1, i].axis("off")

fig.suptitle("Original Tile vs Augmented Versions", fontsize=14, y=1.02)
plt.tight_layout()
plt.show()

print(
    "Notice how the augmented tiles have different orientations, colors, and brightness"
)
print("while the labels are transformed consistently with the images.")

# Load original tile and some augmented versions tile_files = sorted([f for f in Path(output_with_aug, "images").glob("*.tif")]) label_files = sorted([f for f in Path(output_with_aug, "labels").glob("*.tif")]) # Get first 4 tiles (1 original + 3 augmented) n_display = min(4, len(tile_files)) fig, axes = plt.subplots(2, n_display, figsize=(15, 8)) if n_display == 1: axes = axes.reshape(2, 1) for i in range(n_display): # Load and display image tile with rasterio.open(tile_files[i]) as src: img = src.read() # Read all bands # If RGB, display as color if img.shape[0] >= 3: img_display = np.transpose(img[:3], (1, 2, 0)) axes[0, i].imshow(img_display) else: axes[0, i].imshow(img[0], cmap="gray") title = "Original" if i == 0 else f"Augmented {i}" axes[0, i].set_title(f"{title}\n{tile_files[i].name}") axes[0, i].axis("off") # Load and display label tile with rasterio.open(label_files[i]) as src: label = src.read(1) axes[1, i].imshow(label, cmap="tab10", vmin=0, vmax=10) axes[1, i].set_title(f"Label\n{label_files[i].name}") axes[1, i].axis("off") fig.suptitle("Original Tile vs Augmented Versions", fontsize=14, y=1.02) plt.tight_layout() plt.show() print( "Notice how the augmented tiles have different orientations, colors, and brightness" ) print("while the labels are transformed consistently with the images.")

Part 3: Custom Augmentation Pipeline¶

You can also define your own custom augmentation transforms using albumentations. This is useful when you want specific augmentations for your use case.

In [ ]:

import albumentations as A

# Define custom augmentation pipeline
custom_transforms = A.Compose(
    [
        A.HorizontalFlip(p=1.0),  # Always flip horizontally
        A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.8),
        A.HueSaturationValue(
            hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5
        ),
    ]
)

# Export with custom augmentation
output_custom_aug = f"{output_dir}/tiles_custom_augmentation"

geoai.export_geotiff_tiles(
    sample_image,
    output_custom_aug,
    in_class_data=sample_vector,
    tile_size=256,
    stride=128,
    apply_augmentation=True,
    augmentation_count=2,
    augmentation_transforms=custom_transforms,  # Use custom transforms
)

custom_image_tiles = list(Path(output_custom_aug, "images").glob("*.tif"))
print(f"\nWith custom augmentation (2 per tile):")
print(f"  Image tiles: {len(custom_image_tiles)}")

import albumentations as A # Define custom augmentation pipeline custom_transforms = A.Compose( [ A.HorizontalFlip(p=1.0), # Always flip horizontally A.RandomBrightnessContrast(brightness_limit=0.3, contrast_limit=0.3, p=0.8), A.HueSaturationValue( hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.5 ), ] ) # Export with custom augmentation output_custom_aug = f"{output_dir}/tiles_custom_augmentation" geoai.export_geotiff_tiles( sample_image, output_custom_aug, in_class_data=sample_vector, tile_size=256, stride=128, apply_augmentation=True, augmentation_count=2, augmentation_transforms=custom_transforms, # Use custom transforms ) custom_image_tiles = list(Path(output_custom_aug, "images").glob("*.tif")) print(f"\nWith custom augmentation (2 per tile):") print(f" Image tiles: {len(custom_image_tiles)}")

Part 4: Using get_default_augmentation_transforms()¶

You can also access the default augmentation transforms directly to use in your own workflows.

In [ ]:

from geoai.utils import get_default_augmentation_transforms

# Get default transforms
default_transforms = get_default_augmentation_transforms(
    tile_size=256, include_normalize=False  # Don't normalize for visualization
)

print("Default augmentation pipeline:")
print(default_transforms)

# Apply to a sample image
with rasterio.open(tile_files[0]) as src:
    original_img = src.read()
    # Convert to HWC format for albumentations
    img_hwc = np.transpose(original_img, (1, 2, 0))

# Apply augmentation
augmented = default_transforms(image=img_hwc)
aug_img = augmented["image"]

# Visualize
fig, axes = plt.subplots(1, 2, figsize=(10, 5))
axes[0].imshow(img_hwc)
axes[0].set_title("Original")
axes[0].axis("off")

axes[1].imshow(aug_img)
axes[1].set_title("Augmented (random transform)")
axes[1].axis("off")

plt.tight_layout()
plt.show()

from geoai.utils import get_default_augmentation_transforms # Get default transforms default_transforms = get_default_augmentation_transforms( tile_size=256, include_normalize=False # Don't normalize for visualization ) print("Default augmentation pipeline:") print(default_transforms) # Apply to a sample image with rasterio.open(tile_files[0]) as src: original_img = src.read() # Convert to HWC format for albumentations img_hwc = np.transpose(original_img, (1, 2, 0)) # Apply augmentation augmented = default_transforms(image=img_hwc) aug_img = augmented["image"] # Visualize fig, axes = plt.subplots(1, 2, figsize=(10, 5)) axes[0].imshow(img_hwc) axes[0].set_title("Original") axes[0].axis("off") axes[1].imshow(aug_img) axes[1].set_title("Augmented (random transform)") axes[1].axis("off") plt.tight_layout() plt.show()

Part 5: Training with Enhanced Default Augmentation¶

The train_segmentation_model() function now uses improved default augmentations that include:

• Horizontal and vertical flips
• Random 90° rotations
• Brightness and contrast adjustments

Here's an example of how you would train with these defaults (we won't actually run training in this demo):

In [ ]:

from geoai.train import train_segmentation_model

# Train with enhanced default augmentation
model = train_segmentation_model(
    images_dir=f"{output_dir}/tiles_with_augmentation/images",
    labels_dir=f"{output_dir}/tiles_with_augmentation/labels",
    output_dir=f"{output_dir}/tiles_with_augmentation/training_output",
    architecture="unet",
    encoder_name="resnet34",
    num_classes=4,
    batch_size=8,
    num_epochs=20,
    # The following default augmentations are applied automatically:
    # - Horizontal flips (50%)
    # - Vertical flips (50%)
    # - Random 90° rotations (50%)
    # - Brightness adjustment (50%)
    # - Contrast adjustment (50%)
)

from geoai.train import train_segmentation_model # Train with enhanced default augmentation model = train_segmentation_model( images_dir=f"{output_dir}/tiles_with_augmentation/images", labels_dir=f"{output_dir}/tiles_with_augmentation/labels", output_dir=f"{output_dir}/tiles_with_augmentation/training_output", architecture="unet", encoder_name="resnet34", num_classes=4, batch_size=8, num_epochs=20, # The following default augmentations are applied automatically: # - Horizontal flips (50%) # - Vertical flips (50%) # - Random 90° rotations (50%) # - Brightness adjustment (50%) # - Contrast adjustment (50%) )

In [ ]:

geoai.plot_performance_metrics(
    history_path=f"{output_dir}/tiles_with_augmentation/training_output/training_history.pth",
    figsize=(15, 5),
    verbose=True,
)

geoai.plot_performance_metrics( history_path=f"{output_dir}/tiles_with_augmentation/training_output/training_history.pth", figsize=(15, 5), verbose=True, )

Part 6: Custom Training Transforms¶

If you want to use custom augmentations during training, you can define your own transform functions and pass them to train_segmentation_model().

In [ ]:

from geoai.train import (
    train_segmentation_model,
    SemanticTransforms,
    SemanticRandomHorizontalFlip,
    SemanticToTensor,
)

# Define custom training transforms
custom_train_transforms = SemanticTransforms(
    [
        SemanticToTensor(),
        SemanticRandomHorizontalFlip(0.5),
        # Add more custom transforms here...
    ]
)

# Train with custom augmentation
model = train_segmentation_model(
    images_dir=f"{output_dir}/tiles_no_augmentation/images",
    labels_dir=f"{output_dir}/tiles_no_augmentation/labels",
    output_dir=f"{output_dir}/tiles_no_augmentation/training_output",
    architecture="unet",
    num_classes=4,
    train_transforms=custom_train_transforms,  # Use custom transforms
    num_epochs=20,
)

from geoai.train import ( train_segmentation_model, SemanticTransforms, SemanticRandomHorizontalFlip, SemanticToTensor, ) # Define custom training transforms custom_train_transforms = SemanticTransforms( [ SemanticToTensor(), SemanticRandomHorizontalFlip(0.5), # Add more custom transforms here... ] ) # Train with custom augmentation model = train_segmentation_model( images_dir=f"{output_dir}/tiles_no_augmentation/images", labels_dir=f"{output_dir}/tiles_no_augmentation/labels", output_dir=f"{output_dir}/tiles_no_augmentation/training_output", architecture="unet", num_classes=4, train_transforms=custom_train_transforms, # Use custom transforms num_epochs=20, )

In [ ]:

geoai.plot_performance_metrics(
    history_path=f"{output_dir}/tiles_no_augmentation/training_output/training_history.pth",
    figsize=(15, 5),
    verbose=True,
)

geoai.plot_performance_metrics( history_path=f"{output_dir}/tiles_no_augmentation/training_output/training_history.pth", figsize=(15, 5), verbose=True, )

Summary¶

This notebook demonstrated:

1. Tile Export with Augmentation: Using export_geotiff_tiles() with apply_augmentation=True to generate augmented training data
2. Default Augmentation: Using get_default_augmentation_transforms() for sensible defaults optimized for remote sensing
3. Custom Augmentation: Defining custom augmentation pipelines with albumentations
4. Training with Augmentation: How the enhanced defaults work in train_segmentation_model()

Key Benefits of Data Augmentation:¶

• More Training Data: Generate 2-5x more training samples from existing data
• Better Generalization: Models learn to handle variations in orientation, lighting, and appearance
• Reduced Overfitting: More diverse training data helps prevent memorization
• Improved Accuracy: Typically results in 2-5% better validation accuracy

Best Practices:¶

1. Start with default augmentation - it works well for most remote sensing tasks
2. Use 2-5 augmented versions per tile (more isn't always better)
3. Ensure augmentations match your domain (e.g., avoid vertical flips if imagery has a consistent "up" direction)
4. Monitor validation performance to ensure augmentations help rather than hurt

For more information, see the geoai documentation.