inference_large.py

"""
Copyright (c) Microsoft Corporation. All rights reserved.
Licensed under the MIT License.

Version of the inference script that writes outputs to a similar directory structure
as the inputs.
"""
import argparse
import datetime
import os
import time

import numpy as np
import pandas as pd
import rasterio
import rasterio.mask
import torch
import torch.nn.functional as F

from cafo import models, utils
from cafo.data.TileDatasets import TileInferenceDataset

os.environ.update(utils.RASTERIO_BEST_PRACTICES)

NUM_WORKERS = 4
CHIP_SIZE = 256
PADDING = 32
assert PADDING % 2 == 0
HALF_PADDING = PADDING // 2
CHIP_STRIDE = CHIP_SIZE - PADDING

parser = argparse.ArgumentParser(description="CAFO detection model inference script")

parser.add_argument(
    "--input_fn",
    type=str,
    required=True,
    help="Path to a text file containing a list of files to run the model on.",
)
parser.add_argument(
    "--model_fn", type=str, required=True, help="Path to the model file to use."
)
parser.add_argument(
    "--base_output_dir",
    type=str,
    required=True,
    help="Path to a directory where outputs will be saved. This directory will be"
    + " created if it does not exist.",
)
parser.add_argument("--gpu", type=int, default=0, help="ID of the GPU to run on.")
parser.add_argument(
    "--batch_size", type=int, default=64, help="Batch size to use during inference."
)

parser.add_argument(
    "--model", default="unet", choices=("unet", "fcn"), help="Model to use"
)

parser.add_argument(
    "--save_soft", action="store_true", help="Whether to save soft versions of output."
)

args = parser.parse_args()


def main():
    print(
        "Starting CAFO detection model inference script at %s"
        % (str(datetime.datetime.now()))
    )

    # Load files
    assert os.path.exists(args.input_fn)
    assert os.path.exists(args.model_fn)

    # Ensure output directory exists
    assert os.path.exists(args.base_output_dir)

    input_dataframe = pd.read_csv(args.input_fn)
    image_fns = input_dataframe["image_fn"].values
    print("Running on %d files" % (len(image_fns)))

    # Load model
    if torch.cuda.is_available():
        device = torch.device("cuda:%d" % args.gpu)
    else:
        print("WARNING! Torch is reporting that CUDA isn't available, exiting...")
        return
    print("Using device:", device)

    if args.model == "unet":
        model = models.get_unet()
    elif args.model == "fcn":
        model = models.get_fcn()
    else:
        raise ValueError("Invalid model")
    model.load_state_dict(
        torch.load(args.model_fn, map_location="cpu")["model_checkpoint"]
    )
    model = model.eval().to(device)

    # Make sure all output directories exist
    output_dir_set = set()
    for image_idx, image_fn in enumerate(image_fns):
        image_fn = image_fn.replace("http://naipblobs.blob.core.windows.net/naip/", "")
        output_dir = os.path.dirname(image_fn)
        output_dir = os.path.join(args.base_output_dir, output_dir)
        output_dir_set.add(output_dir)

    for output_dir in output_dir_set:
        os.makedirs(output_dir, exist_ok=True)

    # Run model on all files and save output
    for image_idx, image_fn in enumerate(image_fns):
        tic = time.time()

        with rasterio.open(image_fn) as f:
            input_width, input_height = f.width, f.height
            input_profile = f.profile.copy()

        dataset = TileInferenceDataset(
            image_fn,
            chip_size=CHIP_SIZE,
            stride=CHIP_STRIDE,
            transform=utils.chip_transformer,
            verbose=False,
        )
        dataloader = torch.utils.data.DataLoader(
            dataset,
            batch_size=args.batch_size,
            num_workers=NUM_WORKERS,
            pin_memory=True,
        )

        # Run model and organize output
        output = np.zeros((2, input_height, input_width), dtype=np.float32)
        kernel = np.ones((CHIP_SIZE, CHIP_SIZE), dtype=np.float32)
        kernel[HALF_PADDING:-HALF_PADDING, HALF_PADDING:-HALF_PADDING] = 5
        counts = np.zeros((input_height, input_width), dtype=np.float32)

        for i, (data, coords) in enumerate(dataloader):
            data = data.to(device)
            with torch.no_grad():
                t_output = model(data)
                t_output = F.softmax(t_output, dim=1).cpu().numpy()

            for j in range(t_output.shape[0]):
                y, x = coords[j]

                output[:, y : y + CHIP_SIZE, x : x + CHIP_SIZE] += t_output[j] * kernel
                counts[y : y + CHIP_SIZE, x : x + CHIP_SIZE] += kernel

        output = output / counts

        # Save output
        output_profile = input_profile.copy()
        output_profile["driver"] = "GTiff"
        output_profile["dtype"] = "uint8"
        output_profile["compress"] = "lzw"
        output_profile["predictor"] = 2
        output_profile["count"] = 1
        output_profile["nodata"] = 0
        output_profile["tiled"] = True
        output_profile["blockxsize"] = 512
        output_profile["blockysize"] = 512

        if args.save_soft:
            output = output / output.sum(axis=0, keepdims=True)
            output = (output * 255).astype(np.uint8)

            output_fn = image_fn.replace(
                "http://naipblobs.blob.core.windows.net/naip/", ""
            ).replace(".tif", "_predictions-soft.tif")
            output_fn = os.path.join(args.base_output_dir, output_fn)

            with rasterio.open(output_fn, "w", **output_profile) as f:
                f.write(output[1], 1)

        else:
            output_hard = output.argmax(axis=0).astype(np.uint8)

            output_fn = image_fn.replace(
                "http://naipblobs.blob.core.windows.net/naip/", ""
            ).replace(".tif", "_predictions.tif")
            output_fn = os.path.join(args.base_output_dir, output_fn)

            with rasterio.open(output_fn, "w", **output_profile) as f:
                f.write(output_hard, 1)
                f.write_colormap(
                    1,
                    {
                        0: (0, 0, 0, 0),
                        1: (255, 0, 0, 255),
                    },
                )

        print(
            "(%d/%d) Processed %s in %0.4f seconds"
            % (image_idx, len(image_fns), image_fn, time.time() - tic)
        )


if __name__ == "__main__":
    main()