Semantic segmentation

In [ ]:

import os
import torch
from matplotlib import pyplot as plt

from rastervision.core.data import ClassConfig, SemanticSegmentationLabels

import albumentations as A

from rastervision.pytorch_learner import (
    SemanticSegmentationRandomWindowGeoDataset,
    SemanticSegmentationSlidingWindowGeoDataset,
    SemanticSegmentationVisualizer,
    SemanticSegmentationGeoDataConfig,
    SemanticSegmentationLearnerConfig,
    SolverConfig,
    SemanticSegmentationLearner,
)

import os import torch from matplotlib import pyplot as plt from rastervision.core.data import ClassConfig, SemanticSegmentationLabels import albumentations as A from rastervision.pytorch_learner import ( SemanticSegmentationRandomWindowGeoDataset, SemanticSegmentationSlidingWindowGeoDataset, SemanticSegmentationVisualizer, SemanticSegmentationGeoDataConfig, SemanticSegmentationLearnerConfig, SolverConfig, SemanticSegmentationLearner, )

In [ ]:

os.environ["AWS_NO_SIGN_REQUEST"] = "YES"

os.environ["AWS_NO_SIGN_REQUEST"] = "YES"

In [ ]:

class_config = ClassConfig(
    names=["background", "building"],
    colors=["lightgray", "darkred"],
    null_class="background",
)

viz = SemanticSegmentationVisualizer(
    class_names=class_config.names, class_colors=class_config.colors
)

class_config = ClassConfig( names=["background", "building"], colors=["lightgray", "darkred"], null_class="background", ) viz = SemanticSegmentationVisualizer( class_names=class_config.names, class_colors=class_config.colors )

In [ ]:

train_image_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0331E-1257N_1327_3160_13/images/global_monthly_2018_01_mosaic_L15-0331E-1257N_1327_3160_13.tif"
train_label_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0331E-1257N_1327_3160_13/labels/global_monthly_2018_01_mosaic_L15-0331E-1257N_1327_3160_13_Buildings.geojson"

val_image_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/images/global_monthly_2018_01_mosaic_L15-0357E-1223N_1429_3296_13.tif"
val_label_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/labels/global_monthly_2018_01_mosaic_L15-0357E-1223N_1429_3296_13_Buildings.geojson"

train_image_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0331E-1257N_1327_3160_13/images/global_monthly_2018_01_mosaic_L15-0331E-1257N_1327_3160_13.tif" train_label_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0331E-1257N_1327_3160_13/labels/global_monthly_2018_01_mosaic_L15-0331E-1257N_1327_3160_13_Buildings.geojson" val_image_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/images/global_monthly_2018_01_mosaic_L15-0357E-1223N_1429_3296_13.tif" val_label_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/labels/global_monthly_2018_01_mosaic_L15-0357E-1223N_1429_3296_13_Buildings.geojson"

In [ ]:

pred_image_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/images/global_monthly_2020_01_mosaic_L15-0357E-1223N_1429_3296_13.tif"
pred_label_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/labels/global_monthly_2020_01_mosaic_L15-0357E-1223N_1429_3296_13_Buildings.geojson"

pred_image_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/images/global_monthly_2020_01_mosaic_L15-0357E-1223N_1429_3296_13.tif" pred_label_uri = "s3://spacenet-dataset/spacenet/SN7_buildings/train/L15-0357E-1223N_1429_3296_13/labels/global_monthly_2020_01_mosaic_L15-0357E-1223N_1429_3296_13_Buildings.geojson"

In [ ]:

data_augmentation_transform = A.Compose(
    [
        A.Flip(),
        A.ShiftScaleRotate(),
        A.OneOf(
            [
                A.HueSaturationValue(hue_shift_limit=10),
                A.RGBShift(),
                A.ToGray(),
                A.ToSepia(),
                A.RandomBrightness(),
                A.RandomGamma(),
            ]
        ),
        A.CoarseDropout(max_height=32, max_width=32, max_holes=5),
    ]
)

data_augmentation_transform = A.Compose( [ A.Flip(), A.ShiftScaleRotate(), A.OneOf( [ A.HueSaturationValue(hue_shift_limit=10), A.RGBShift(), A.ToGray(), A.ToSepia(), A.RandomBrightness(), A.RandomGamma(), ] ), A.CoarseDropout(max_height=32, max_width=32, max_holes=5), ] )

In [ ]:

train_ds = SemanticSegmentationRandomWindowGeoDataset.from_uris(
    class_config=class_config,
    image_uri=train_image_uri,
    label_vector_uri=train_label_uri,
    label_vector_default_class_id=class_config.get_class_id("building"),
    size_lims=(150, 200),
    out_size=256,
    max_windows=400,
    transform=data_augmentation_transform,
)

len(train_ds)

train_ds = SemanticSegmentationRandomWindowGeoDataset.from_uris( class_config=class_config, image_uri=train_image_uri, label_vector_uri=train_label_uri, label_vector_default_class_id=class_config.get_class_id("building"), size_lims=(150, 200), out_size=256, max_windows=400, transform=data_augmentation_transform, ) len(train_ds)

In [ ]:

x, y = viz.get_batch(train_ds, 4)
viz.plot_batch(x, y, show=True)

x, y = viz.get_batch(train_ds, 4) viz.plot_batch(x, y, show=True)

In [ ]:

val_ds = SemanticSegmentationSlidingWindowGeoDataset.from_uris(
    class_config=class_config,
    image_uri=val_image_uri,
    label_vector_uri=val_label_uri,
    label_vector_default_class_id=class_config.get_class_id("building"),
    size=200,
    stride=100,
    transform=A.Resize(256, 256),
)
len(val_ds)

val_ds = SemanticSegmentationSlidingWindowGeoDataset.from_uris( class_config=class_config, image_uri=val_image_uri, label_vector_uri=val_label_uri, label_vector_default_class_id=class_config.get_class_id("building"), size=200, stride=100, transform=A.Resize(256, 256), ) len(val_ds)

In [ ]:

x, y = viz.get_batch(val_ds, 4)
viz.plot_batch(x, y, show=True)

x, y = viz.get_batch(val_ds, 4) viz.plot_batch(x, y, show=True)

In [ ]:

pred_ds = SemanticSegmentationSlidingWindowGeoDataset.from_uris(
    class_config=class_config,
    image_uri=pred_image_uri,
    size=200,
    stride=100,
    transform=A.Resize(256, 256),
)
len(pred_ds)

pred_ds = SemanticSegmentationSlidingWindowGeoDataset.from_uris( class_config=class_config, image_uri=pred_image_uri, size=200, stride=100, transform=A.Resize(256, 256), ) len(pred_ds)

In [ ]:

model = torch.hub.load(
    "AdeelH/pytorch-fpn:0.3",
    "make_fpn_resnet",
    name="resnet18",
    fpn_type="panoptic",
    num_classes=len(class_config),
    fpn_channels=128,
    in_channels=3,
    out_size=(256, 256),
    pretrained=True,
)

model = torch.hub.load( "AdeelH/pytorch-fpn:0.3", "make_fpn_resnet", name="resnet18", fpn_type="panoptic", num_classes=len(class_config), fpn_channels=128, in_channels=3, out_size=(256, 256), pretrained=True, )

In [ ]:

data_cfg = SemanticSegmentationGeoDataConfig(
    class_names=class_config.names,
    class_colors=class_config.colors,
    num_workers=0,  # increase to use multi-processing
)

data_cfg = SemanticSegmentationGeoDataConfig( class_names=class_config.names, class_colors=class_config.colors, num_workers=0, # increase to use multi-processing )

In [ ]:

solver_cfg = SolverConfig(batch_sz=8, lr=3e-2, class_loss_weights=[1.0, 10.0])

solver_cfg = SolverConfig(batch_sz=8, lr=3e-2, class_loss_weights=[1.0, 10.0])

In [ ]:

learner_cfg = SemanticSegmentationLearnerConfig(data=data_cfg, solver=solver_cfg)

learner_cfg = SemanticSegmentationLearnerConfig(data=data_cfg, solver=solver_cfg)

In [ ]:

learner = SemanticSegmentationLearner(
    cfg=learner_cfg,
    output_dir="./train-demo/",
    model=model,
    train_ds=train_ds,
    valid_ds=val_ds,
)

learner = SemanticSegmentationLearner( cfg=learner_cfg, output_dir="./train-demo/", model=model, train_ds=train_ds, valid_ds=val_ds, )

In [ ]:

learner.log_data_stats()

learner.log_data_stats()

In [ ]:

%load_ext tensorboard

%load_ext tensorboard

In [ ]:

%tensorboard --bind_all --logdir "./train-demo/tb-logs" --reload_interval 10

%tensorboard --bind_all --logdir "./train-demo/tb-logs" --reload_interval 10

In [ ]:

learner.train(epochs=3)

learner.train(epochs=3)

In [ ]:

learner.train(epochs=1)

learner.train(epochs=1)

In [ ]:

learner.plot_predictions(split="valid", show=True)

learner.plot_predictions(split="valid", show=True)

In [ ]:

learner.save_model_bundle()

learner.save_model_bundle()

In [ ]:

learner = SemanticSegmentationLearner.from_model_bundle(
    model_bundle_uri="./train-demo/model-bundle.zip",
    output_dir="./train-demo/",
    model=model,
)

learner = SemanticSegmentationLearner.from_model_bundle( model_bundle_uri="./train-demo/model-bundle.zip", output_dir="./train-demo/", model=model, )

In [ ]:

learner = SemanticSegmentationLearner.from_model_bundle(
    model_bundle_uri="./train-demo/model-bundle.zip",
    output_dir="./train-demo/",
    model=model,
    train_ds=train_ds,
    valid_ds=val_ds,
    training=True,
)

learner = SemanticSegmentationLearner.from_model_bundle( model_bundle_uri="./train-demo/model-bundle.zip", output_dir="./train-demo/", model=model, train_ds=train_ds, valid_ds=val_ds, training=True, )

In [ ]:

learner.train(epochs=1)

learner.train(epochs=1)

In [ ]:

learner.plot_predictions(split="valid", show=True)

learner.plot_predictions(split="valid", show=True)

In [ ]:

predictions = learner.predict_dataset(
    pred_ds,
    raw_out=True,
    numpy_out=True,
    predict_kw=dict(out_shape=(325, 325)),
    progress_bar=True,
)

predictions = learner.predict_dataset( pred_ds, raw_out=True, numpy_out=True, predict_kw=dict(out_shape=(325, 325)), progress_bar=True, )

In [ ]:

pred_labels = SemanticSegmentationLabels.from_predictions(
    pred_ds.windows,
    predictions,
    smooth=True,
    extent=pred_ds.scene.extent,
    num_classes=len(class_config),
)

pred_labels = SemanticSegmentationLabels.from_predictions( pred_ds.windows, predictions, smooth=True, extent=pred_ds.scene.extent, num_classes=len(class_config), )

In [ ]:

scores = pred_labels.get_score_arr(pred_labels.extent)

scores = pred_labels.get_score_arr(pred_labels.extent)

In [ ]:

pred_labels.save(
    uri=f"predict",
    crs_transformer=pred_ds.scene.raster_source.crs_transformer,
    class_config=class_config,
)

pred_labels.save( uri=f"predict", crs_transformer=pred_ds.scene.raster_source.crs_transformer, class_config=class_config, )