climategan/fire.py

import torch
import torch.nn.functional as F
import random
import kornia
from torchvision.transforms.functional import adjust_brightness, adjust_contrast

from climategan.tutils import normalize, retrieve_sky_mask

try:
    from kornia.filters import filter2d
except ImportError:
    from kornia.filters import filter2D as filter2d


def increase_sky_mask(mask, p_w=0, p_h=0):
    """
    Increases sky mask in width and height by a given pourcentage
    (Purpose: when applying Gaussian blur, there are no artifacts of blue sky behind)
    Args:
        sky_mask (torch.Tensor): Sky mask of shape (H,W)
        p_w (float): Percentage of mask width by which to increase
            the width of the sky region
        p_h (float): Percentage of mask height by which to increase
            the height of the sky region
    Returns:
        torch.Tensor: Sky mask increased given p_w and p_h
    """

    if p_h <= 0 and p_w <= 0:
        return mask

    n_lines = int(p_h * mask.shape[-2])
    n_cols = int(p_w * mask.shape[-1])

    temp_mask = mask.clone().detach()
    for i in range(1, n_cols):
        temp_mask[:, :, :, i::] += mask[:, :, :, 0:-i]
        temp_mask[:, :, :, 0:-i] += mask[:, :, :, i::]

    new_mask = temp_mask.clone().detach()
    for i in range(1, n_lines):
        new_mask[:, :, i::, :] += temp_mask[:, :, 0:-i, :]
        new_mask[:, :, 0:-i, :] += temp_mask[:, :, i::, :]

    new_mask[new_mask >= 1] = 1

    return new_mask


def paste_filter(x, filter_, mask):
    """
    Pastes a filter over an image given a mask
    Where the mask is 1, the filter is copied as is.
    Where the mask is 0, the current value is preserved.
    Intermediate values will mix the two images together.
    Args:
        x (torch.Tensor): Input tensor, range must be [0, 255]
        filer_ (torch.Tensor): Filter, range must be [0, 255]
        mask (torch.Tensor): Mask, range must be [0, 1]
    Returns:
        torch.Tensor: New tensor with filter pasted on it
    """
    assert len(x.shape) == len(filter_.shape) == len(mask.shape)
    x = filter_ * mask + x * (1 - mask)
    return x


def add_fire(x, seg_preds, fire_opts):
    """
    Transforms input tensor given wildfires event
    Args:
        x (torch.Tensor): Input tensor
        seg_preds (torch.Tensor): Semantic segmentation predictions for input tensor
        filter_color (tuple): (r,g,b) tuple for the color of the sky
        blur_radius (float): radius of the Gaussian blur that smooths
            the transition between sky and foreground
    Returns:
        torch.Tensor: Wildfire version of input tensor
    """
    wildfire_tens = normalize(x, 0, 255)

    # Warm the image
    wildfire_tens[:, 2, :, :] -= 20
    wildfire_tens[:, 1, :, :] -= 10
    wildfire_tens[:, 0, :, :] += 40
    wildfire_tens.clamp_(0, 255)
    wildfire_tens = wildfire_tens.to(torch.uint8)

    # Darken the picture and increase contrast
    wildfire_tens = adjust_contrast(wildfire_tens, contrast_factor=1.5)
    wildfire_tens = adjust_brightness(wildfire_tens, brightness_factor=0.73)

    sky_mask = retrieve_sky_mask(seg_preds).unsqueeze(1)

    if fire_opts.get("crop_bottom_sky_mask"):
        i = 2 * sky_mask.shape[-2] // 3
        sky_mask[..., i:, :] = 0

    sky_mask = F.interpolate(
        sky_mask.to(torch.float),
        (wildfire_tens.shape[-2], wildfire_tens.shape[-1]),
    )
    sky_mask = increase_sky_mask(sky_mask, 0.18, 0.18)

    kernel_size = (fire_opts.get("kernel_size", 301), fire_opts.get("kernel_size", 301))
    sigma = (fire_opts.get("kernel_sigma", 150.5), fire_opts.get("kernel_sigma", 150.5))
    border_type = "reflect"
    kernel = torch.unsqueeze(
        kornia.filters.kernels.get_gaussian_kernel2d(kernel_size, sigma), dim=0
    ).to(x.device)
    sky_mask = filter2d(sky_mask, kernel, border_type)

    filter_ = torch.ones(wildfire_tens.shape, device=x.device)
    filter_[:, 0, :, :] = 255
    filter_[:, 1, :, :] = random.randint(100, 150)
    filter_[:, 2, :, :] = 0

    wildfire_tens = paste_tensor(wildfire_tens, filter_, sky_mask, 200)

    wildfire_tens = adjust_brightness(wildfire_tens.to(torch.uint8), 0.8)
    wildfire_tens = wildfire_tens.to(torch.float)

    # dummy pixels to fool scaling and preserve range
    wildfire_tens[:, :, 0, 0] = 255.0
    wildfire_tens[:, :, -1, -1] = 0.0

    return wildfire_tens


def paste_tensor(source, filter_, mask, transparency):
    mask = transparency / 255.0 * mask
    new = mask * filter_ + (1.0 - mask) * source
    return new