|
|
|
"""Utility functions for image editing.""" |
|
|
|
import numpy as np |
|
import cv2 |
|
import torch |
|
|
|
|
|
__all__ = ['to_numpy', 'linear_interpolate', 'make_transform', |
|
'get_ind', 'mask2image'] |
|
|
|
|
|
def to_numpy(data): |
|
"""Converts the input data to `numpy.ndarray`.""" |
|
if isinstance(data, (int, float)): |
|
return np.array(data) |
|
if isinstance(data, np.ndarray): |
|
return data |
|
if isinstance(data, torch.Tensor): |
|
return data.detach().cpu().numpy() |
|
raise TypeError(f'Not supported data type `{type(data)}` for ' |
|
f'converting to `numpy.ndarray`!') |
|
|
|
|
|
def linear_interpolate(latent_code, |
|
boundary, |
|
layer_index=None, |
|
start_distance=-10.0, |
|
end_distance=10.0, |
|
steps=21): |
|
"""Interpolate between the latent code and boundary.""" |
|
assert (len(latent_code.shape) == 3 and len(boundary.shape) == 3 and |
|
latent_code.shape[0] == 1 and boundary.shape[0] == 1 and |
|
latent_code.shape[1] == boundary.shape[1]) |
|
linspace = np.linspace(start_distance, end_distance, steps) |
|
linspace = linspace.reshape([-1, 1, 1]).astype(np.float32) |
|
inter_code = linspace * boundary |
|
is_manipulatable = np.zeros(inter_code.shape, dtype=bool) |
|
is_manipulatable[:, layer_index, :] = True |
|
mani_code = np.where(is_manipulatable, latent_code+inter_code, latent_code) |
|
return mani_code |
|
|
|
|
|
def make_transform(tx, ty, angle): |
|
"""Transform the input feature maps with given |
|
coordinates and rotation angle. |
|
|
|
cos(theta) -sin(theta) tx |
|
sin(theta) cos(theta) ty |
|
0 0 1 |
|
|
|
""" |
|
m = np.eye(3) |
|
s = np.sin(angle/360.0*np.pi*2) |
|
c = np.cos(angle/360.0*np.pi*2) |
|
m[0][0] = c |
|
m[0][1] = s |
|
m[0][2] = tx |
|
m[1][0] = -s |
|
m[1][1] = c |
|
m[1][2] = ty |
|
return m |
|
|
|
|
|
def get_ind(seg_mask, label): |
|
"""Get the index of the masked and unmasked region.""" |
|
mask = np.where(seg_mask == label, |
|
np.ones_like(seg_mask), |
|
np.zeros_like(seg_mask)) |
|
f_ind = np.where(mask == 1) |
|
b_ind = np.where((1 - mask) == 1) |
|
return f_ind, b_ind, mask |
|
|
|
|
|
def mask2image(image, mask, r=3, g=255, b=118): |
|
"""Show the mask on the given image.""" |
|
assert image.shape[0] == image.shape[1] |
|
r_c = np.ones([256, 256, 1]) * r |
|
g_c = np.ones([256, 256, 1]) * g |
|
b_c = np.ones([256, 256, 1]) * b |
|
img1 = np.concatenate([r_c, g_c, b_c], axis=2).astype(np.uint8) |
|
mask = np.expand_dims(mask, axis=2).astype(np.uint8) |
|
img1 = img1 * mask |
|
image = cv2.addWeighted(image, 0.4, img1, 0.6, 0) |
|
mask_i = np.tile(mask, [1, 1, 3]) * 255 |
|
return image, mask_i |
|
|
