demo-painttransformer / render_utils.py
jaekookang
remove cv
48c86e5
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
# import cv2
import numpy as np
from PIL import Image
import math
class Erosion2d(nn.Layer):
"""
Erosion2d
"""
def __init__(self, m=1):
super(Erosion2d, self).__init__()
self.m = m
self.pad = [m, m, m, m]
def forward(self, x):
batch_size, c, h, w = x.shape
x_pad = F.pad(x, pad=self.pad, mode='constant', value=1e9)
channel = nn.functional.unfold(x_pad, 2 * self.m + 1, strides=1, paddings=0).reshape([batch_size, c, -1, h, w])
result = paddle.min(channel, axis=2)
return result
class Dilation2d(nn.Layer):
"""
Dilation2d
"""
def __init__(self, m=1):
super(Dilation2d, self).__init__()
self.m = m
self.pad = [m, m, m, m]
def forward(self, x):
batch_size, c, h, w = x.shape
x_pad = F.pad(x, pad=self.pad, mode='constant', value=-1e9)
channel = nn.functional.unfold(x_pad, 2 * self.m + 1, strides=1, paddings=0).reshape([batch_size, c, -1, h, w])
result = paddle.max(channel, axis=2)
return result
def param2stroke(param, H, W, meta_brushes):
"""
param2stroke
"""
b = param.shape[0]
param_list = paddle.split(param, 8, axis=1)
x0, y0, w, h, theta = [item.squeeze(-1) for item in param_list[:5]]
sin_theta = paddle.sin(math.pi * theta)
cos_theta = paddle.cos(math.pi * theta)
index = paddle.full((b,), -1, dtype='int64').numpy()
index[(h > w).numpy()] = 0
index[(h <= w).numpy()] = 1
meta_brushes_resize = F.interpolate(meta_brushes, (H, W)).numpy()
brush = paddle.to_tensor(meta_brushes_resize[index])
warp_00 = cos_theta / w
warp_01 = sin_theta * H / (W * w)
warp_02 = (1 - 2 * x0) * cos_theta / w + (1 - 2 * y0) * sin_theta * H / (W * w)
warp_10 = -sin_theta * W / (H * h)
warp_11 = cos_theta / h
warp_12 = (1 - 2 * y0) * cos_theta / h - (1 - 2 * x0) * sin_theta * W / (H * h)
warp_0 = paddle.stack([warp_00, warp_01, warp_02], axis=1)
warp_1 = paddle.stack([warp_10, warp_11, warp_12], axis=1)
warp = paddle.stack([warp_0, warp_1], axis=1)
grid = nn.functional.affine_grid(warp, [b, 3, H, W]) # paddle和torch默认值是反过来的
brush = nn.functional.grid_sample(brush, grid)
return brush
def read_img(img_path, img_type='RGB', h=None, w=None):
"""
read img
"""
img = Image.open(img_path).convert(img_type)
if h is not None and w is not None:
img = img.resize((w, h), resample=Image.NEAREST)
img = np.array(img)
if img.ndim == 2:
img = np.expand_dims(img, axis=-1)
img = img.transpose((2, 0, 1))
img = paddle.to_tensor(img).unsqueeze(0).astype('float32') / 255.
return img
# def preprocess(img, w=512, h=512):
# image = cv2.resize(img, (w, h), cv2.INTER_NEAREST)
# image = image.transpose((2, 0, 1))
# image = paddle.to_tensor(image).unsqueeze(0).astype('float32') / 255.
# return image
def pad(img, H, W):
b, c, h, w = img.shape
pad_h = (H - h) // 2
pad_w = (W - w) // 2
remainder_h = (H - h) % 2
remainder_w = (W - w) % 2
expand_img = nn.functional.pad(img, [pad_w, pad_w + remainder_w,
pad_h, pad_h + remainder_h])
return expand_img