inference update

inference.py

@@ -0,0 +1,496 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+import network
+import morphology
+import os
+import math
+
+idx = 0
+
+
+def save_img(img, output_path):
+    result = Image.fromarray((img.data.cpu().numpy().transpose((1, 2, 0)) * 255).astype(np.uint8))
+    result.save(output_path)
+
+
+def param2stroke(param, H, W, meta_brushes):
+    """
+    Input a set of stroke parameters and output its corresponding foregrounds and alpha maps.
+    Args:
+        param: a tensor with shape n_strokes x n_param_per_stroke. Here, param_per_stroke is 8:
+        x_center, y_center, width, height, theta, R, G, and B.
+        H: output height.
+        W: output width.
+        meta_brushes: a tensor with shape 2 x 3 x meta_brush_height x meta_brush_width.
+         The first slice on the batch dimension denotes vertical brush and the second one denotes horizontal brush.
+
+    Returns:
+        foregrounds: a tensor with shape n_strokes x 3 x H x W, containing color information.
+        alphas: a tensor with shape n_strokes x 3 x H x W,
+         containing binary information of whether a pixel is belonging to the stroke (alpha mat), for painting process.
+    """
+    # Firstly, resize the meta brushes to the required shape,
+    # in order to decrease GPU memory especially when the required shape is small.
+    meta_brushes_resize = F.interpolate(meta_brushes, (H, W))
+    b = param.shape[0]
+    # Extract shape parameters and color parameters.
+    param_list = torch.split(param, 1, dim=1)
+    x0, y0, w, h, theta = [item.squeeze(-1) for item in param_list[:5]]
+    R, G, B = param_list[5:]
+    # Pre-compute sin theta and cos theta
+    sin_theta = torch.sin(torch.acos(torch.tensor(-1., device=param.device)) * theta)
+    cos_theta = torch.cos(torch.acos(torch.tensor(-1., device=param.device)) * theta)
+    # index means each stroke should use which meta stroke? Vertical meta stroke or horizontal meta stroke.
+    # When h > w, vertical stroke should be used. When h <= w, horizontal stroke should be used.
+    index = torch.full((b,), -1, device=param.device, dtype=torch.long)
+    index[h > w] = 0
+    index[h <= w] = 1
+    brush = meta_brushes_resize[index.long()]
+
+    # Calculate warp matrix according to the rules defined by pytorch, in order for warping.
+    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 = torch.stack([warp_00, warp_01, warp_02], dim=1)
+    warp_1 = torch.stack([warp_10, warp_11, warp_12], dim=1)
+    warp = torch.stack([warp_0, warp_1], dim=1)
+    # Conduct warping.
+    grid = F.affine_grid(warp, [b, 3, H, W], align_corners=False)
+    brush = F.grid_sample(brush, grid, align_corners=False)
+    # alphas is the binary information suggesting whether a pixel is belonging to the stroke.
+    alphas = (brush > 0).float()
+    brush = brush.repeat(1, 3, 1, 1)
+    alphas = alphas.repeat(1, 3, 1, 1)
+    # Give color to foreground strokes.
+    color_map = torch.cat([R, G, B], dim=1)
+    color_map = color_map.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, H, W)
+    foreground = brush * color_map
+    # Dilation and erosion are used for foregrounds and alphas respectively to prevent artifacts on stroke borders.
+    foreground = morphology.dilation(foreground)
+    alphas = morphology.erosion(alphas)
+    return foreground, alphas
+
+
+def param2img_serial(
+        param, decision, meta_brushes, cur_canvas, frame_dir, has_border=False, original_h=None, original_w=None):
+    """
+    Input stroke parameters and decisions for each patch, meta brushes, current canvas, frame directory,
+    and whether there is a border (if intermediate painting results are required).
+    Output the painting results of adding the corresponding strokes on the current canvas.
+    Args:
+        param: a tensor with shape batch size x patch along height dimension x patch along width dimension
+         x n_stroke_per_patch x n_param_per_stroke
+        decision: a 01 tensor with shape batch size x patch along height dimension x patch along width dimension
+         x n_stroke_per_patch
+        meta_brushes: a tensor with shape 2 x 3 x meta_brush_height x meta_brush_width.
+        The first slice on the batch dimension denotes vertical brush and the second one denotes horizontal brush.
+        cur_canvas: a tensor with shape batch size x 3 x H x W,
+         where H and W denote height and width of padded results of original images.
+        frame_dir: directory to save intermediate painting results. None means intermediate results are not required.
+        has_border: on the last painting layer, in order to make sure that the painting results do not miss
+         any important detail, we choose to paint again on this layer but shift patch_size // 2 pixels when
+         cutting patches. In this case, if intermediate results are required, we need to cut the shifted length
+         on the border before saving, or there would be a black border.
+        original_h: to indicate the original height for cropping when saving intermediate results.
+        original_w: to indicate the original width for cropping when saving intermediate results.
+
+    Returns:
+        cur_canvas: a tensor with shape batch size x 3 x H x W, denoting painting results.
+    """
+    # param: b, h, w, stroke_per_patch, param_per_stroke
+    # decision: b, h, w, stroke_per_patch
+    b, h, w, s, p = param.shape
+    H, W = cur_canvas.shape[-2:]
+    is_odd_y = h % 2 == 1
+    is_odd_x = w % 2 == 1
+    patch_size_y = 2 * H // h
+    patch_size_x = 2 * W // w
+    even_idx_y = torch.arange(0, h, 2, device=cur_canvas.device)
+    even_idx_x = torch.arange(0, w, 2, device=cur_canvas.device)
+    odd_idx_y = torch.arange(1, h, 2, device=cur_canvas.device)
+    odd_idx_x = torch.arange(1, w, 2, device=cur_canvas.device)
+    even_y_even_x_coord_y, even_y_even_x_coord_x = torch.meshgrid([even_idx_y, even_idx_x])
+    odd_y_odd_x_coord_y, odd_y_odd_x_coord_x = torch.meshgrid([odd_idx_y, odd_idx_x])
+    even_y_odd_x_coord_y, even_y_odd_x_coord_x = torch.meshgrid([even_idx_y, odd_idx_x])
+    odd_y_even_x_coord_y, odd_y_even_x_coord_x = torch.meshgrid([odd_idx_y, even_idx_x])
+    cur_canvas = F.pad(cur_canvas, [patch_size_x // 4, patch_size_x // 4,
+                                    patch_size_y // 4, patch_size_y // 4, 0, 0, 0, 0])
+
+    def partial_render(this_canvas, patch_coord_y, patch_coord_x, stroke_id):
+        canvas_patch = F.unfold(this_canvas, (patch_size_y, patch_size_x),
+                                stride=(patch_size_y // 2, patch_size_x // 2))
+        # canvas_patch: b, 3 * py * px, h * w
+        canvas_patch = canvas_patch.view(b, 3, patch_size_y, patch_size_x, h, w).contiguous()
+        canvas_patch = canvas_patch.permute(0, 4, 5, 1, 2, 3).contiguous()
+        # canvas_patch: b, h, w, 3, py, px
+        selected_canvas_patch = canvas_patch[:, patch_coord_y, patch_coord_x, :, :, :]
+        selected_h, selected_w = selected_canvas_patch.shape[1:3]
+        selected_param = param[:, patch_coord_y, patch_coord_x, stroke_id, :].view(-1, p).contiguous()
+        selected_decision = decision[:, patch_coord_y, patch_coord_x, stroke_id].view(-1).contiguous()
+        selected_foregrounds = torch.zeros(selected_param.shape[0], 3, patch_size_y, patch_size_x,
+                                           device=this_canvas.device)
+        selected_alphas = torch.zeros(selected_param.shape[0], 3, patch_size_y, patch_size_x, device=this_canvas.device)
+        if selected_param[selected_decision, :].shape[0] > 0:
+            selected_foregrounds[selected_decision, :, :, :], selected_alphas[selected_decision, :, :, :] = \
+                param2stroke(selected_param[selected_decision, :], patch_size_y, patch_size_x, meta_brushes)
+        selected_foregrounds = selected_foregrounds.view(
+            b, selected_h, selected_w, 3, patch_size_y, patch_size_x).contiguous()
+        selected_alphas = selected_alphas.view(b, selected_h, selected_w, 3, patch_size_y, patch_size_x).contiguous()
+        selected_decision = selected_decision.view(b, selected_h, selected_w, 1, 1, 1).contiguous()
+        selected_canvas_patch = selected_foregrounds * selected_alphas * selected_decision + selected_canvas_patch * (
+                1 - selected_alphas * selected_decision)
+        this_canvas = selected_canvas_patch.permute(0, 3, 1, 4, 2, 5).contiguous()
+        # this_canvas: b, 3, selected_h, py, selected_w, px
+        this_canvas = this_canvas.view(b, 3, selected_h * patch_size_y, selected_w * patch_size_x).contiguous()
+        # this_canvas: b, 3, selected_h * py, selected_w * px
+        return this_canvas
+
+    global idx
+    if has_border:
+        factor = 2
+    else:
+        factor = 4
+    if even_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
+        for i in range(s):
+            canvas = partial_render(cur_canvas, even_y_even_x_coord_y, even_y_even_x_coord_x, i)
+            if not is_odd_y:
+                canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, :canvas.shape[3]]], dim=2)
+            if not is_odd_x:
+                canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+            cur_canvas = canvas
+            idx += 1
+            if frame_dir is not None:
+                frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
+                             patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
+                save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+
+    if odd_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
+        for i in range(s):
+            canvas = partial_render(cur_canvas, odd_y_odd_x_coord_y, odd_y_odd_x_coord_x, i)
+            canvas = torch.cat([cur_canvas[:, :, :patch_size_y // 2, -canvas.shape[3]:], canvas], dim=2)
+            canvas = torch.cat([cur_canvas[:, :, -canvas.shape[2]:, :patch_size_x // 2], canvas], dim=3)
+            if is_odd_y:
+                canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, :canvas.shape[3]]], dim=2)
+            if is_odd_x:
+                canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+            cur_canvas = canvas
+            idx += 1
+            if frame_dir is not None:
+                frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
+                             patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
+                save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+
+    if odd_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
+        for i in range(s):
+            canvas = partial_render(cur_canvas, odd_y_even_x_coord_y, odd_y_even_x_coord_x, i)
+            canvas = torch.cat([cur_canvas[:, :, :patch_size_y // 2, :canvas.shape[3]], canvas], dim=2)
+            if is_odd_y:
+                canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, :canvas.shape[3]]], dim=2)
+            if not is_odd_x:
+                canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+            cur_canvas = canvas
+            idx += 1
+            if frame_dir is not None:
+                frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
+                             patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
+                save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+
+    if even_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
+        for i in range(s):
+            canvas = partial_render(cur_canvas, even_y_odd_x_coord_y, even_y_odd_x_coord_x, i)
+            canvas = torch.cat([cur_canvas[:, :, :canvas.shape[2], :patch_size_x // 2], canvas], dim=3)
+            if not is_odd_y:
+                canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, -canvas.shape[3]:]], dim=2)
+            if is_odd_x:
+                canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+            cur_canvas = canvas
+            idx += 1
+            if frame_dir is not None:
+                frame = crop(cur_canvas[:, :, patch_size_y // factor:-patch_size_y // factor,
+                             patch_size_x // factor:-patch_size_x // factor], original_h, original_w)
+                save_img(frame[0], os.path.join(frame_dir, '%03d.jpg' % idx))
+
+    cur_canvas = cur_canvas[:, :, patch_size_y // 4:-patch_size_y // 4, patch_size_x // 4:-patch_size_x // 4]
+
+    return cur_canvas
+
+
+def param2img_parallel(param, decision, meta_brushes, cur_canvas):
+    """
+        Input stroke parameters and decisions for each patch, meta brushes, current canvas, frame directory,
+        and whether there is a border (if intermediate painting results are required).
+        Output the painting results of adding the corresponding strokes on the current canvas.
+        Args:
+            param: a tensor with shape batch size x patch along height dimension x patch along width dimension
+             x n_stroke_per_patch x n_param_per_stroke
+            decision: a 01 tensor with shape batch size x patch along height dimension x patch along width dimension
+             x n_stroke_per_patch
+            meta_brushes: a tensor with shape 2 x 3 x meta_brush_height x meta_brush_width.
+            The first slice on the batch dimension denotes vertical brush and the second one denotes horizontal brush.
+            cur_canvas: a tensor with shape batch size x 3 x H x W,
+             where H and W denote height and width of padded results of original images.
+
+        Returns:
+            cur_canvas: a tensor with shape batch size x 3 x H x W, denoting painting results.
+        """
+    # param: b, h, w, stroke_per_patch, param_per_stroke
+    # decision: b, h, w, stroke_per_patch
+    b, h, w, s, p = param.shape
+    param = param.view(-1, 8).contiguous()
+    decision = decision.view(-1).contiguous().bool()
+    H, W = cur_canvas.shape[-2:]
+    is_odd_y = h % 2 == 1
+    is_odd_x = w % 2 == 1
+    patch_size_y = 2 * H // h
+    patch_size_x = 2 * W // w
+    even_idx_y = torch.arange(0, h, 2, device=cur_canvas.device)
+    even_idx_x = torch.arange(0, w, 2, device=cur_canvas.device)
+    odd_idx_y = torch.arange(1, h, 2, device=cur_canvas.device)
+    odd_idx_x = torch.arange(1, w, 2, device=cur_canvas.device)
+    even_y_even_x_coord_y, even_y_even_x_coord_x = torch.meshgrid([even_idx_y, even_idx_x])
+    odd_y_odd_x_coord_y, odd_y_odd_x_coord_x = torch.meshgrid([odd_idx_y, odd_idx_x])
+    even_y_odd_x_coord_y, even_y_odd_x_coord_x = torch.meshgrid([even_idx_y, odd_idx_x])
+    odd_y_even_x_coord_y, odd_y_even_x_coord_x = torch.meshgrid([odd_idx_y, even_idx_x])
+    cur_canvas = F.pad(cur_canvas, [patch_size_x // 4, patch_size_x // 4,
+                                    patch_size_y // 4, patch_size_y // 4, 0, 0, 0, 0])
+    foregrounds = torch.zeros(param.shape[0], 3, patch_size_y, patch_size_x, device=cur_canvas.device)
+    alphas = torch.zeros(param.shape[0], 3, patch_size_y, patch_size_x, device=cur_canvas.device)
+    valid_foregrounds, valid_alphas = param2stroke(param[decision, :], patch_size_y, patch_size_x, meta_brushes)
+    foregrounds[decision, :, :, :] = valid_foregrounds
+    alphas[decision, :, :, :] = valid_alphas
+    # foreground, alpha: b * h * w * stroke_per_patch, 3, patch_size_y, patch_size_x
+    foregrounds = foregrounds.view(-1, h, w, s, 3, patch_size_y, patch_size_x).contiguous()
+    alphas = alphas.view(-1, h, w, s, 3, patch_size_y, patch_size_x).contiguous()
+    # foreground, alpha: b, h, w, stroke_per_patch, 3, render_size_y, render_size_x
+    decision = decision.view(-1, h, w, s, 1, 1, 1).contiguous()
+
+    # decision: b, h, w, stroke_per_patch, 1, 1, 1
+
+    def partial_render(this_canvas, patch_coord_y, patch_coord_x):
+
+        canvas_patch = F.unfold(this_canvas, (patch_size_y, patch_size_x),
+                                stride=(patch_size_y // 2, patch_size_x // 2))
+        # canvas_patch: b, 3 * py * px, h * w
+        canvas_patch = canvas_patch.view(b, 3, patch_size_y, patch_size_x, h, w).contiguous()
+        canvas_patch = canvas_patch.permute(0, 4, 5, 1, 2, 3).contiguous()
+        # canvas_patch: b, h, w, 3, py, px
+        selected_canvas_patch = canvas_patch[:, patch_coord_y, patch_coord_x, :, :, :]
+        selected_foregrounds = foregrounds[:, patch_coord_y, patch_coord_x, :, :, :, :]
+        selected_alphas = alphas[:, patch_coord_y, patch_coord_x, :, :, :, :]
+        selected_decisions = decision[:, patch_coord_y, patch_coord_x, :, :, :, :]
+        for i in range(s):
+            cur_foreground = selected_foregrounds[:, :, :, i, :, :, :]
+            cur_alpha = selected_alphas[:, :, :, i, :, :, :]
+            cur_decision = selected_decisions[:, :, :, i, :, :, :]
+            selected_canvas_patch = cur_foreground * cur_alpha * cur_decision + selected_canvas_patch * (
+                    1 - cur_alpha * cur_decision)
+        this_canvas = selected_canvas_patch.permute(0, 3, 1, 4, 2, 5).contiguous()
+        # this_canvas: b, 3, h_half, py, w_half, px
+        h_half = this_canvas.shape[2]
+        w_half = this_canvas.shape[4]
+        this_canvas = this_canvas.view(b, 3, h_half * patch_size_y, w_half * patch_size_x).contiguous()
+        # this_canvas: b, 3, h_half * py, w_half * px
+        return this_canvas
+
+    if even_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
+        canvas = partial_render(cur_canvas, even_y_even_x_coord_y, even_y_even_x_coord_x)
+        if not is_odd_y:
+            canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, :canvas.shape[3]]], dim=2)
+        if not is_odd_x:
+            canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+        cur_canvas = canvas
+
+    if odd_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
+        canvas = partial_render(cur_canvas, odd_y_odd_x_coord_y, odd_y_odd_x_coord_x)
+        canvas = torch.cat([cur_canvas[:, :, :patch_size_y // 2, -canvas.shape[3]:], canvas], dim=2)
+        canvas = torch.cat([cur_canvas[:, :, -canvas.shape[2]:, :patch_size_x // 2], canvas], dim=3)
+        if is_odd_y:
+            canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, :canvas.shape[3]]], dim=2)
+        if is_odd_x:
+            canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+        cur_canvas = canvas
+
+    if odd_idx_y.shape[0] > 0 and even_idx_x.shape[0] > 0:
+        canvas = partial_render(cur_canvas, odd_y_even_x_coord_y, odd_y_even_x_coord_x)
+        canvas = torch.cat([cur_canvas[:, :, :patch_size_y // 2, :canvas.shape[3]], canvas], dim=2)
+        if is_odd_y:
+            canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, :canvas.shape[3]]], dim=2)
+        if not is_odd_x:
+            canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+        cur_canvas = canvas
+
+    if even_idx_y.shape[0] > 0 and odd_idx_x.shape[0] > 0:
+        canvas = partial_render(cur_canvas, even_y_odd_x_coord_y, even_y_odd_x_coord_x)
+        canvas = torch.cat([cur_canvas[:, :, :canvas.shape[2], :patch_size_x // 2], canvas], dim=3)
+        if not is_odd_y:
+            canvas = torch.cat([canvas, cur_canvas[:, :, -patch_size_y // 2:, -canvas.shape[3]:]], dim=2)
+        if is_odd_x:
+            canvas = torch.cat([canvas, cur_canvas[:, :, :canvas.shape[2], -patch_size_x // 2:]], dim=3)
+        cur_canvas = canvas
+
+    cur_canvas = cur_canvas[:, :, patch_size_y // 4:-patch_size_y // 4, patch_size_x // 4:-patch_size_x // 4]
+
+    return cur_canvas
+
+
+def read_img(img_path, img_type='RGB', h=None, w=None):
+    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 = torch.from_numpy(img).unsqueeze(0).float() / 255.
+    return img
+
+
+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
+    img = torch.cat([torch.zeros((b, c, pad_h, w), device=img.device), img,
+                     torch.zeros((b, c, pad_h + remainder_h, w), device=img.device)], dim=-2)
+    img = torch.cat([torch.zeros((b, c, H, pad_w), device=img.device), img,
+                     torch.zeros((b, c, H, pad_w + remainder_w), device=img.device)], dim=-1)
+    return img
+
+
+def crop(img, h, w):
+    H, W = img.shape[-2:]
+    pad_h = (H - h) // 2
+    pad_w = (W - w) // 2
+    remainder_h = (H - h) % 2
+    remainder_w = (W - w) % 2
+    img = img[:, :, pad_h:H - pad_h - remainder_h, pad_w:W - pad_w - remainder_w]
+    return img
+
+
+def main(input_path, model_path, output_dir, need_animation=False, resize_h=None, resize_w=None, serial=False):
+    if not os.path.exists(output_dir):
+        os.mkdir(output_dir)
+    input_name = os.path.basename(input_path)
+    output_path = os.path.join(output_dir, input_name)
+    frame_dir = None
+    if need_animation:
+        if not serial:
+            print('It must be under serial mode if animation results are required, so serial flag is set to True!')
+            serial = True
+        frame_dir = os.path.join(output_dir, input_name[:input_name.find('.')])
+        if not os.path.exists(frame_dir):
+            os.mkdir(frame_dir)
+    patch_size = 32
+    stroke_num = 8
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    net_g = network.Painter(5, stroke_num, 256, 8, 3, 3).to(device)
+    net_g.load_state_dict(torch.load(model_path))
+    net_g.eval()
+    for param in net_g.parameters():
+        param.requires_grad = False
+
+    brush_large_vertical = read_img('brush/brush_large_vertical.png', 'L').to(device)
+    brush_large_horizontal = read_img('brush/brush_large_horizontal.png', 'L').to(device)
+    meta_brushes = torch.cat(
+        [brush_large_vertical, brush_large_horizontal], dim=0)
+
+    with torch.no_grad():
+        original_img = read_img(input_path, 'RGB', resize_h, resize_w).to(device)
+        original_h, original_w = original_img.shape[-2:]
+        K = max(math.ceil(math.log2(max(original_h, original_w) / patch_size)), 0)
+        original_img_pad_size = patch_size * (2 ** K)
+        original_img_pad = pad(original_img, original_img_pad_size, original_img_pad_size)
+        final_result = torch.zeros_like(original_img_pad).to(device)
+        for layer in range(0, K + 1):
+            layer_size = patch_size * (2 ** layer)
+            img = F.interpolate(original_img_pad, (layer_size, layer_size))
+            result = F.interpolate(final_result, (patch_size * (2 ** layer), patch_size * (2 ** layer)))
+            img_patch = F.unfold(img, (patch_size, patch_size), stride=(patch_size, patch_size))
+            result_patch = F.unfold(result, (patch_size, patch_size),
+                                    stride=(patch_size, patch_size))
+            # There are patch_num * patch_num patches in total
+            patch_num = (layer_size - patch_size) // patch_size + 1
+
+            # img_patch, result_patch: b, 3 * output_size * output_size, h * w
+            img_patch = img_patch.permute(0, 2, 1).contiguous().view(-1, 3, patch_size, patch_size).contiguous()
+            result_patch = result_patch.permute(0, 2, 1).contiguous().view(
+                -1, 3, patch_size, patch_size).contiguous()
+            shape_param, stroke_decision = net_g(img_patch, result_patch)
+            stroke_decision = network.SignWithSigmoidGrad.apply(stroke_decision)
+
+            grid = shape_param[:, :, :2].view(img_patch.shape[0] * stroke_num, 1, 1, 2).contiguous()
+            img_temp = img_patch.unsqueeze(1).contiguous().repeat(1, stroke_num, 1, 1, 1).view(
+                img_patch.shape[0] * stroke_num, 3, patch_size, patch_size).contiguous()
+            color = F.grid_sample(img_temp, 2 * grid - 1, align_corners=False).view(
+                img_patch.shape[0], stroke_num, 3).contiguous()
+            stroke_param = torch.cat([shape_param, color], dim=-1)
+            # stroke_param: b * h * w, stroke_per_patch, param_per_stroke
+            # stroke_decision: b * h * w, stroke_per_patch, 1
+            param = stroke_param.view(1, patch_num, patch_num, stroke_num, 8).contiguous()
+            decision = stroke_decision.view(1, patch_num, patch_num, stroke_num).contiguous().bool()
+            # param: b, h, w, stroke_per_patch, 8
+            # decision: b, h, w, stroke_per_patch
+            param[..., :2] = param[..., :2] / 2 + 0.25
+            param[..., 2:4] = param[..., 2:4] / 2
+            if serial:
+                final_result = param2img_serial(param, decision, meta_brushes, final_result,
+                                                frame_dir, False, original_h, original_w)
+            else:
+                final_result = param2img_parallel(param, decision, meta_brushes, final_result)
+
+        border_size = original_img_pad_size // (2 * patch_num)
+        img = F.interpolate(original_img_pad, (patch_size * (2 ** layer), patch_size * (2 ** layer)))
+        result = F.interpolate(final_result, (patch_size * (2 ** layer), patch_size * (2 ** layer)))
+        img = F.pad(img, [patch_size // 2, patch_size // 2, patch_size // 2, patch_size // 2,
+                          0, 0, 0, 0])
+        result = F.pad(result, [patch_size // 2, patch_size // 2, patch_size // 2, patch_size // 2,
+                                0, 0, 0, 0])
+        img_patch = F.unfold(img, (patch_size, patch_size), stride=(patch_size, patch_size))
+        result_patch = F.unfold(result, (patch_size, patch_size), stride=(patch_size, patch_size))
+        final_result = F.pad(final_result, [border_size, border_size, border_size, border_size, 0, 0, 0, 0])
+        h = (img.shape[2] - patch_size) // patch_size + 1
+        w = (img.shape[3] - patch_size) // patch_size + 1
+        # img_patch, result_patch: b, 3 * output_size * output_size, h * w
+        img_patch = img_patch.permute(0, 2, 1).contiguous().view(-1, 3, patch_size, patch_size).contiguous()
+        result_patch = result_patch.permute(0, 2, 1).contiguous().view(-1, 3, patch_size, patch_size).contiguous()
+        shape_param, stroke_decision = net_g(img_patch, result_patch)
+
+        grid = shape_param[:, :, :2].view(img_patch.shape[0] * stroke_num, 1, 1, 2).contiguous()
+        img_temp = img_patch.unsqueeze(1).contiguous().repeat(1, stroke_num, 1, 1, 1).view(
+            img_patch.shape[0] * stroke_num, 3, patch_size, patch_size).contiguous()
+        color = F.grid_sample(img_temp, 2 * grid - 1, align_corners=False).view(
+            img_patch.shape[0], stroke_num, 3).contiguous()
+        stroke_param = torch.cat([shape_param, color], dim=-1)
+        # stroke_param: b * h * w, stroke_per_patch, param_per_stroke
+        # stroke_decision: b * h * w, stroke_per_patch, 1
+        param = stroke_param.view(1, h, w, stroke_num, 8).contiguous()
+        decision = stroke_decision.view(1, h, w, stroke_num).contiguous().bool()
+        # param: b, h, w, stroke_per_patch, 8
+        # decision: b, h, w, stroke_per_patch
+        param[..., :2] = param[..., :2] / 2 + 0.25
+        param[..., 2:4] = param[..., 2:4] / 2
+        if serial:
+            final_result = param2img_serial(param, decision, meta_brushes, final_result,
+                                            frame_dir, True, original_h, original_w)
+        else:
+            final_result = param2img_parallel(param, decision, meta_brushes, final_result)
+        final_result = final_result[:, :, border_size:-border_size, border_size:-border_size]
+
+        final_result = crop(final_result, original_h, original_w)
+        save_img(final_result[0], output_path)
+
+
+if __name__ == '__main__':
+    main(input_path='input/chicago.jpg',
+         model_path='model.pth',
+         output_dir='output/',
+         need_animation=False,  # whether need intermediate results for animation.
+         resize_h=None,         # resize original input to this size. None means do not resize.
+         resize_w=None,         # resize original input to this size. None means do not resize.
+         serial=False)          # if need animation, serial must be True.

morphology.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Erosion2d(nn.Module):
+
+    def __init__(self, m=1):
+        super(Erosion2d, self).__init__()
+        self.m = m
+        self.pad = [m, m, m, m]
+        self.unfold = nn.Unfold(2 * m + 1, padding=0, stride=1)
+
+    def forward(self, x):
+        batch_size, c, h, w = x.shape
+        x_pad = F.pad(x, pad=self.pad, mode='constant', value=1e9)
+        channel = self.unfold(x_pad).view(batch_size, c, -1, h, w)
+        result = torch.min(channel, dim=2)[0]
+        return result
+
+
+def erosion(x, m=1):
+    b, c, h, w = x.shape
+    x_pad = F.pad(x, pad=[m, m, m, m], mode='constant', value=1e9)
+    channel = nn.functional.unfold(x_pad, 2 * m + 1, padding=0, stride=1).view(b, c, -1, h, w)
+    result = torch.min(channel, dim=2)[0]
+    return result
+
+
+class Dilation2d(nn.Module):
+
+    def __init__(self, m=1):
+        super(Dilation2d, self).__init__()
+        self.m = m
+        self.pad = [m, m, m, m]
+        self.unfold = nn.Unfold(2 * m + 1, padding=0, stride=1)
+
+    def forward(self, x):
+        batch_size, c, h, w = x.shape
+        x_pad = F.pad(x, pad=self.pad, mode='constant', value=-1e9)
+        channel = self.unfold(x_pad).view(batch_size, c, -1, h, w)
+        result = torch.max(channel, dim=2)[0]
+        return result
+
+
+def dilation(x, m=1):
+    b, c, h, w = x.shape
+    x_pad = F.pad(x, pad=[m, m, m, m], mode='constant', value=-1e9)
+    channel = nn.functional.unfold(x_pad, 2 * m + 1, padding=0, stride=1).view(b, c, -1, h, w)
+    result = torch.max(channel, dim=2)[0]
+    return result

network.py

@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+
+
+class SignWithSigmoidGrad(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, x):
+        result = (x > 0).float()
+        sigmoid_result = torch.sigmoid(x)
+        ctx.save_for_backward(sigmoid_result)
+        return result
+
+    @staticmethod
+    def backward(ctx, grad_result):
+        (sigmoid_result,) = ctx.saved_tensors
+        if ctx.needs_input_grad[0]:
+            grad_input = grad_result * sigmoid_result * (1 - sigmoid_result)
+        else:
+            grad_input = None
+        return grad_input
+
+
+class Painter(nn.Module):
+
+    def __init__(self, param_per_stroke, total_strokes, hidden_dim, n_heads=8, n_enc_layers=3, n_dec_layers=3):
+        super().__init__()
+        self.enc_img = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(3, 32, 3, 1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(32, 64, 3, 2),
+            nn.BatchNorm2d(64),
+            nn.ReLU(True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(64, 128, 3, 2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(True))
+        self.enc_canvas = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(3, 32, 3, 1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(32, 64, 3, 2),
+            nn.BatchNorm2d(64),
+            nn.ReLU(True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(64, 128, 3, 2),
+            nn.BatchNorm2d(128),
+            nn.ReLU(True))
+        self.conv = nn.Conv2d(128 * 2, hidden_dim, 1)
+        self.transformer = nn.Transformer(hidden_dim, n_heads, n_enc_layers, n_dec_layers)
+        self.linear_param = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(True),
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(True),
+            nn.Linear(hidden_dim, param_per_stroke))
+        self.linear_decider = nn.Linear(hidden_dim, 1)
+        self.query_pos = nn.Parameter(torch.rand(total_strokes, hidden_dim))
+        self.row_embed = nn.Parameter(torch.rand(8, hidden_dim // 2))
+        self.col_embed = nn.Parameter(torch.rand(8, hidden_dim // 2))
+
+    def forward(self, img, canvas):
+        b, _, H, W = img.shape
+        img_feat = self.enc_img(img)
+        canvas_feat = self.enc_canvas(canvas)
+        h, w = img_feat.shape[-2:]
+        feat = torch.cat([img_feat, canvas_feat], dim=1)
+        feat_conv = self.conv(feat)
+
+        pos_embed = torch.cat([
+            self.col_embed[:w].unsqueeze(0).contiguous().repeat(h, 1, 1),
+            self.row_embed[:h].unsqueeze(1).contiguous().repeat(1, w, 1),
+        ], dim=-1).flatten(0, 1).unsqueeze(1)
+        hidden_state = self.transformer(pos_embed + feat_conv.flatten(2).permute(2, 0, 1).contiguous(),
+                                        self.query_pos.unsqueeze(1).contiguous().repeat(1, b, 1))
+        hidden_state = hidden_state.permute(1, 0, 2).contiguous()
+        param = self.linear_param(hidden_state)
+        decision = self.linear_decider(hidden_state)
+        return param, decision