apply_events.py

import argparse


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-b",
        "--batch_size",
        type=int,
        default=4,
        help="Batch size to process input images to events. Defaults to 4",
    )
    parser.add_argument(
        "-i",
        "--images_paths",
        type=str,
        required=True,
        help="Path to a directory with image files",
    )
    parser.add_argument(
        "-o",
        "--output_path",
        type=str,
        default=None,
        help="Path to a directory were events should be written. "
        + "Will NOT write anything to disk if this flag is not used.",
    )
    parser.add_argument(
        "-s",
        "--save_input",
        action="store_true",
        default=False,
        help="Binary flag to include the input image to the model (after crop and"
        + " resize) in the images written or uploaded (depending on saving options.)",
    )
    parser.add_argument(
        "-r",
        "--resume_path",
        type=str,
        default=None,
        help="Path to a directory containing the trainer to resume."
        + " In particular it must contain `opts.yam` and `checkpoints/`."
        + " Typically this points to a Masker, which holds the path to a"
        + " Painter in its opts",
    )
    parser.add_argument(
        "--no_time",
        action="store_true",
        default=False,
        help="Binary flag to prevent the timing of operations.",
    )
    parser.add_argument(
        "-f",
        "--flood_mask_binarization",
        type=float,
        default=0.5,
        help="Value to use to binarize masks (mask > value). "
        + "Set to -1 to use soft masks (not binarized). Defaults to 0.5.",
    )
    parser.add_argument(
        "-t",
        "--target_size",
        type=int,
        default=640,
        help="Output image size (when not using `keep_ratio_128`): images are resized"
        + " such that their smallest side is `target_size` then cropped in the middle"
        + " of the largest side such that the resulting input image (and output images)"
        + " has height and width `target_size x target_size`. **Must** be a multiple of"
        + " 2^7=128 (up/downscaling inside the models). Defaults to 640.",
    )
    parser.add_argument(
        "--half",
        action="store_true",
        default=False,
        help="Binary flag to use half precision (float16). Defaults to False.",
    )
    parser.add_argument(
        "-n",
        "--n_images",
        default=-1,
        type=int,
        help="Limit the number of images processed (if you have 100 images in "
        + "a directory but n is 10 then only the first 10 images will be loaded"
        + " for processing)",
    )
    parser.add_argument(
        "--no_conf",
        action="store_true",
        default=False,
        help="disable writing the apply_events hash and command in the output folder",
    )
    parser.add_argument(
        "--overwrite",
        action="store_true",
        default=False,
        help="Do not check for existing outdir, i.e. force overwrite"
        + " potentially existing files in the output path",
    )
    parser.add_argument(
        "--no_cloudy",
        action="store_true",
        default=False,
        help="Prevent the use of the cloudy intermediate"
        + " image to create the flood image. Rendering will"
        + " be more colorful but may seem less realistic",
    )
    parser.add_argument(
        "--keep_ratio_128",
        action="store_true",
        default=False,
        help="When loading the input images, resize and crop them in order for their "
        + "dimensions to match the closest multiples"
        + " of 128. Will force a batch size of 1 since images"
        + " now have different dimensions. "
        + "Use --max_im_width to cap the resulting dimensions.",
    )
    parser.add_argument(
        "--fuse",
        action="store_true",
        default=False,
        help="Use batch norm fusion to speed up inference",
    )
    parser.add_argument(
        "--save_masks",
        action="store_true",
        default=False,
        help="Save output masks along events",
    )
    parser.add_argument(
        "-m",
        "--max_im_width",
        type=int,
        default=-1,
        help="When using --keep_ratio_128, some images may still be too large. Use "
        + "--max_im_width to cap the resized image's width. Defaults to -1 (no cap).",
    )
    parser.add_argument(
        "--upload",
        action="store_true",
        help="Upload to comet.ml in a project called `climategan-apply`",
    )
    parser.add_argument(
        "--zip_outdir",
        "-z",
        action="store_true",
        help="Zip the output directory as '{outdir.parent}/{outdir.name}.zip'",
    )
    return parser.parse_args()


args = parse_args()


print("\n• Imports\n")
import time

import_time = time.time()
import sys
import shutil
from collections import OrderedDict
from pathlib import Path

import comet_ml  # noqa: F401
import torch
import numpy as np
import skimage.io as io
from skimage.color import rgba2rgb
from skimage.transform import resize
from tqdm import tqdm

from climategan.trainer import Trainer
from climategan.bn_fusion import bn_fuse
from climategan.tutils import print_num_parameters
from climategan.utils import Timer, find_images, get_git_revision_hash, to_128, resolve

import_time = time.time() - import_time


def to_m1_p1(img, i):
    """
    rescales a [0, 1] image to [-1, +1]

    Args:
        img (np.array): float32 numpy array of an image in [0, 1]
        i (int): Index of the image being rescaled

    Raises:
        ValueError: If the image is not in [0, 1]

    Returns:
        np.array(np.float32): array in [-1, +1]
    """
    if img.min() >= 0 and img.max() <= 1:
        return (img.astype(np.float32) - 0.5) * 2
    raise ValueError(f"Data range mismatch for image {i} : ({img.min()}, {img.max()})")


def uint8(array):
    """
    convert an array to np.uint8 (does not rescale or anything else than changing dtype)

    Args:
        array (np.array): array to modify

    Returns:
        np.array(np.uint8): converted array
    """
    return array.astype(np.uint8)


def resize_and_crop(img, to=640):
    """
    Resizes an image so that it keeps the aspect ratio and the smallest dimensions
    is `to`, then crops this resized image in its center so that the output is `to x to`
    without aspect ratio distortion

    Args:
        img (np.array): np.uint8 255 image

    Returns:
        np.array: [0, 1] np.float32 image
    """
    # resize keeping aspect ratio: smallest dim is 640
    h, w = img.shape[:2]
    if h < w:
        size = (to, int(to * w / h))
    else:
        size = (int(to * h / w), to)

    r_img = resize(img, size, preserve_range=True, anti_aliasing=True)
    r_img = uint8(r_img)

    # crop in the center
    H, W = r_img.shape[:2]

    top = (H - to) // 2
    left = (W - to) // 2

    rc_img = r_img[top : top + to, left : left + to, :]

    return rc_img / 255.0


def print_time(text, time_series, purge=-1):
    """
    Print a timeseries's mean and std with a label

    Args:
        text (str): label of the time series
        time_series (list): list of timings
        purge (int, optional): ignore first n values of time series. Defaults to -1.
    """
    if not time_series:
        return

    if purge > 0 and len(time_series) > purge:
        time_series = time_series[purge:]

    m = np.mean(time_series)
    s = np.std(time_series)

    print(
        f"{text.capitalize() + ' ':.<26}  {m:.5f}"
        + (f" +/- {s:.5f}" if len(time_series) > 1 else "")
    )


def print_store(store, purge=-1):
    """
    Pretty-print time series store

    Args:
        store (dict): maps string keys to lists of times
        purge (int, optional): ignore first n values of time series. Defaults to -1.
    """
    singles = OrderedDict({k: v for k, v in store.items() if len(v) == 1})
    multiples = OrderedDict({k: v for k, v in store.items() if len(v) > 1})
    empties = {k: v for k, v in store.items() if len(v) == 0}

    if empties:
        print("Ignoring empty stores ", ", ".join(empties.keys()))
        print()

    for k in singles:
        print_time(k, singles[k], purge)

    print()
    print("Unit: s/batch")
    for k in multiples:
        print_time(k, multiples[k], purge)
    print()


def write_apply_config(out):
    """
    Saves the args to `apply_events.py` in a text file for future reference
    """
    cwd = Path.cwd().expanduser().resolve()
    command = f"cd {str(cwd)}\n"
    command += " ".join(sys.argv)
    git_hash = get_git_revision_hash()
    with (out / "command.txt").open("w") as f:
        f.write(command)
    with (out / "hash.txt").open("w") as f:
        f.write(git_hash)


def get_outdir_name(half, keep_ratio, max_im_width, target_size, bin_value, cloudy):
    """
    Create the output directory's name based on uer-provided arguments
    """
    name_items = []
    if half:
        name_items.append("half")
    if keep_ratio:
        name_items.append("AR")
    if max_im_width and keep_ratio:
        name_items.append(f"{max_im_width}")
    if target_size and not keep_ratio:
        name_items.append("S")
        name_items.append(f"{target_size}")
    if bin_value != 0.5:
        name_items.append(f"bin{bin_value}")
    if not cloudy:
        name_items.append("no_cloudy")

    return "-".join(name_items)


def make_outdir(
    outdir, overwrite, half, keep_ratio, max_im_width, target_size, bin_value, cloudy
):
    """
    Creates the output directory if it does not exist. If it does exist,
    prompts the user for confirmation (except if `overwrite` is True).
    If the output directory's name is "_auto_" then it is created as:
        outdir.parent / get_outdir_name(...)
    """
    if outdir.name == "_auto_":
        outdir = outdir.parent / get_outdir_name(
            half, keep_ratio, max_im_width, target_size, bin_value, cloudy
        )
    if outdir.exists() and not overwrite:
        print(
            f"\nWARNING: outdir ({str(outdir)}) already exists."
            + " Files with existing names will be overwritten"
        )
        if "n" in input(">>> Continue anyway? [y / n] (default: y) : "):
            print("Interrupting execution from user input.")
            sys.exit()
        print()
    outdir.mkdir(exist_ok=True, parents=True)
    return outdir


def get_time_stores(import_time):
    return OrderedDict(
        {
            "imports": [import_time],
            "setup": [],
            "data pre-processing": [],
            "encode": [],
            "mask": [],
            "flood": [],
            "depth": [],
            "segmentation": [],
            "smog": [],
            "wildfire": [],
            "all events": [],
            "numpy": [],
            "inference on all images": [],
            "write": [],
        }
    )


if __name__ == "__main__":

    # -----------------------------------------
    # -----  Initialize script variables  -----
    # -----------------------------------------
    print(
        "• Using args\n\n"
        + "\n".join(["{:25}: {}".format(k, v) for k, v in vars(args).items()]),
    )

    batch_size = args.batch_size
    bin_value = args.flood_mask_binarization
    cloudy = not args.no_cloudy
    fuse = args.fuse
    half = args.half
    save_masks = args.save_masks
    images_paths = resolve(args.images_paths)
    keep_ratio = args.keep_ratio_128
    max_im_width = args.max_im_width
    n_images = args.n_images
    outdir = resolve(args.output_path) if args.output_path is not None else None
    resume_path = args.resume_path
    target_size = args.target_size
    time_inference = not args.no_time
    upload = args.upload
    zip_outdir = args.zip_outdir

    # -------------------------------------
    # -----  Validate size arguments  -----
    # -------------------------------------
    if keep_ratio:
        if target_size != 640:
            print(
                "\nWARNING: using --keep_ratio_128 overwrites target_size"
                + " which is ignored."
            )
        if batch_size != 1:
            print("\nWARNING: batch_size overwritten to 1 when using keep_ratio_128")
            batch_size = 1
        if max_im_width > 0 and max_im_width % 128 != 0:
            new_im_width = int(max_im_width / 128) * 128
            print("\nWARNING: max_im_width should be <0 or a multiple of 128.")
            print(
                "            Was {} but is now overwritten to {}".format(
                    max_im_width, new_im_width
                )
            )
            max_im_width = new_im_width
    else:
        if target_size % 128 != 0:
            print(f"\nWarning: target size {target_size} is not a multiple of 128.")
            target_size = target_size - (target_size % 128)
            print(f"Setting target_size to {target_size}.")

    # -------------------------------------
    # -----  Create output directory  -----
    # -------------------------------------
    if outdir is not None:
        outdir = make_outdir(
            outdir,
            args.overwrite,
            half,
            keep_ratio,
            max_im_width,
            target_size,
            bin_value,
            cloudy,
        )

    # -------------------------------
    # -----  Create time store  -----
    # -------------------------------
    stores = get_time_stores(import_time)

    # -----------------------------------
    # -----  Load Trainer instance  -----
    # -----------------------------------
    with Timer(store=stores.get("setup", []), ignore=time_inference):
        print("\n• Initializing trainer\n")
        torch.set_grad_enabled(False)
        trainer = Trainer.resume_from_path(
            resume_path,
            setup=True,
            inference=True,
            new_exp=None,
        )
        print()
        print_num_parameters(trainer, True)
        if fuse:
            trainer.G = bn_fuse(trainer.G)
        if half:
            trainer.G.half()

    # --------------------------------------------
    # -----  Read data from input directory  -----
    # --------------------------------------------
    print("\n• Reading & Pre-processing Data\n")

    # find all images
    data_paths = find_images(images_paths)
    base_data_paths = data_paths
    # filter images
    if 0 < n_images < len(data_paths):
        data_paths = data_paths[:n_images]
    # repeat data
    elif n_images > len(data_paths):
        repeats = n_images // len(data_paths) + 1
        data_paths = base_data_paths * repeats
        data_paths = data_paths[:n_images]

    with Timer(store=stores.get("data pre-processing", []), ignore=time_inference):
        # read images to numpy arrays
        data = [io.imread(str(d)) for d in data_paths]
        # rgba to rgb
        data = [im if im.shape[-1] == 3 else uint8(rgba2rgb(im) * 255) for im in data]
        # resize images to target_size or
        if keep_ratio:
            # to closest multiples of 128 <= max_im_width, keeping aspect ratio
            new_sizes = [to_128(d, max_im_width) for d in data]
            data = [resize(d, ns, anti_aliasing=True) for d, ns in zip(data, new_sizes)]
        else:
            # to args.target_size
            data = [resize_and_crop(d, target_size) for d in data]
            new_sizes = [(target_size, target_size) for _ in data]
        # resize() produces [0, 1] images, rescale to [-1, 1]
        data = [to_m1_p1(d, i) for i, d in enumerate(data)]

    n_batchs = len(data) // batch_size
    if len(data) % batch_size != 0:
        n_batchs += 1

    print("Found", len(base_data_paths), "images. Inferring on", len(data), "images.")

    # --------------------------------------------
    # -----  Batch-process images to events  -----
    # --------------------------------------------
    print(f"\n• Using device {str(trainer.device)}\n")

    all_events = []

    with Timer(store=stores.get("inference on all images", []), ignore=time_inference):
        for b in tqdm(range(n_batchs), desc="Infering events", unit="batch"):

            images = data[b * batch_size : (b + 1) * batch_size]
            if not images:
                continue

            # concatenate images in a batch batch_size x height x width x 3
            images = np.stack(images)
            # Retreive numpy events as a dict {event: array[BxHxWxC]}
            events = trainer.infer_all(
                images,
                numpy=True,
                stores=stores,
                bin_value=bin_value,
                half=half,
                cloudy=cloudy,
                return_masks=save_masks,
            )

            # save resized and cropped image
            if args.save_input:
                events["input"] = uint8((images + 1) / 2 * 255)

            # store events to write after inference loop
            all_events.append(events)

    # --------------------------------------------
    # -----  Save (write/upload) inferences  -----
    # --------------------------------------------
    if outdir is not None or upload:

        if upload:
            print("\n• Creating comet Experiment")
            exp = comet_ml.Experiment(project_name="climategan-apply")
            exp.log_parameters(vars(args))

        # --------------------------------------------------------------
        # -----  Change inferred data structure to a list of dicts  -----
        # --------------------------------------------------------------
        to_write = []
        events_names = list(all_events[0].keys())
        for events_data in all_events:
            n_ims = len(events_data[events_names[0]])
            for i in range(n_ims):
                item = {event: events_data[event][i] for event in events_names}
                to_write.append(item)

        progress_bar_desc = ""
        if outdir is not None:
            print("\n• Output directory:\n")
            print(str(outdir), "\n")
            if upload:
                progress_bar_desc = "Writing & Uploading events"
            else:
                progress_bar_desc = "Writing events"
        else:
            if upload:
                progress_bar_desc = "Uploading events"

        # ------------------------------------
        # -----  Save individual images  -----
        # ------------------------------------
        with Timer(store=stores.get("write", []), ignore=time_inference):

            # for each image
            for t, event_dict in tqdm(
                enumerate(to_write),
                desc=progress_bar_desc,
                unit="input image",
                total=len(to_write),
            ):

                idx = t % len(base_data_paths)
                stem = Path(data_paths[idx]).stem
                width = new_sizes[idx][1]

                if keep_ratio:
                    ar = "_AR"
                else:
                    ar = ""

                # for each event type
                event_bar = tqdm(
                    enumerate(event_dict.items()),
                    leave=False,
                    total=len(events_names),
                    unit="event",
                )
                for e, (event, im_data) in event_bar:
                    event_bar.set_description(
                        f"  {event.capitalize():<{len(progress_bar_desc) - 2}}"
                    )

                    if args.no_cloudy:
                        suffix = ar + "_no_cloudy"
                    else:
                        suffix = ar

                    im_path = Path(f"{stem}_{event}_{width}{suffix}.png")

                    if outdir is not None:
                        im_path = outdir / im_path
                        io.imsave(im_path, im_data)

                    if upload:
                        exp.log_image(im_data, name=im_path.name)
    if zip_outdir:
        print("\n• Zipping output directory... ", end="", flush=True)
        archive_path = Path(shutil.make_archive(outdir.name, "zip", root_dir=outdir))
        archive_path = archive_path.rename(outdir.parent / archive_path.name)
        print("Done:\n")
        print(str(archive_path))

    # ---------------------------
    # -----  Print timings  -----
    # ---------------------------
    if time_inference:
        print("\n• Timings\n")
        print_store(stores)

    # ---------------------------------------------
    # -----  Save apply_events.py run config  -----
    # ---------------------------------------------
    if not args.no_conf and outdir is not None:
        write_apply_config(outdir)