Spaces:
Runtime error
Runtime error
| 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 | |