ICML2022/resefa

1# python3.72"""Utility functions for image editing."""3 4import numpy as np5import cv26import torch7 8 9__all__ = ['to_numpy', 'linear_interpolate', 'make_transform',10           'get_ind', 'mask2image']11 12 13def to_numpy(data):14    """Converts the input data to `numpy.ndarray`."""15    if isinstance(data, (int, float)):16        return np.array(data)17    if isinstance(data, np.ndarray):18        return data19    if isinstance(data, torch.Tensor):20        return data.detach().cpu().numpy()21    raise TypeError(f'Not supported data type `{type(data)}` for '22                    f'converting to `numpy.ndarray`!')23 24 25def linear_interpolate(latent_code,26                       boundary,27                       layer_index=None,28                       start_distance=-10.0,29                       end_distance=10.0,30                       steps=21):31    """Interpolate between the latent code and boundary."""32    assert (len(latent_code.shape) == 3 and len(boundary.shape) == 3 and33            latent_code.shape[0] == 1 and boundary.shape[0] == 1 and34            latent_code.shape[1] == boundary.shape[1])35    linspace = np.linspace(start_distance, end_distance, steps)36    linspace = linspace.reshape([-1, 1, 1]).astype(np.float32)37    inter_code = linspace * boundary38    is_manipulatable = np.zeros(inter_code.shape, dtype=bool)39    is_manipulatable[:, layer_index, :] = True40    mani_code = np.where(is_manipulatable, latent_code+inter_code, latent_code)41    return mani_code42 43 44def make_transform(tx, ty, angle):45    """Transform the input feature maps with given46    coordinates and rotation angle.47 48    cos(theta) -sin(theta) tx49    sin(theta)  cos(theta) ty50        0          0        151 52    """53    m = np.eye(3)54    s = np.sin(angle/360.0*np.pi*2)55    c = np.cos(angle/360.0*np.pi*2)56    m[0][0] = c57    m[0][1] = s58    m[0][2] = tx59    m[1][0] = -s60    m[1][1] = c61    m[1][2] = ty62    return m63 64 65def get_ind(seg_mask, label):66    """Get the index of the masked and unmasked region."""67    mask = np.where(seg_mask == label,68                    np.ones_like(seg_mask),69                    np.zeros_like(seg_mask))70    f_ind = np.where(mask == 1)71    b_ind = np.where((1 - mask) == 1)72    return f_ind, b_ind, mask73 74 75def mask2image(image, mask, r=3, g=255, b=118):76    """Show the mask on the given image."""77    assert image.shape[0] == image.shape[1]78    r_c = np.ones([256, 256, 1]) * r79    g_c = np.ones([256, 256, 1]) * g80    b_c = np.ones([256, 256, 1]) * b81    img1 = np.concatenate([r_c, g_c, b_c], axis=2).astype(np.uint8)82    mask = np.expand_dims(mask, axis=2).astype(np.uint8)83    img1 = img1 * mask84    image = cv2.addWeighted(image, 0.4, img1, 0.6, 0)85    mask_i = np.tile(mask, [1, 1, 3]) * 25586    return image, mask_i87