論文復現丨基於ModelArts進行影像風格化繪畫

華為雲開發者聯盟發表於2022-12-22
摘要:這個 notebook 基於論文「Stylized Neural Painting, arXiv:2011.08114.」提供了最基本的「圖片生成繪畫」變換的可復現例子。

本文分享自華為雲社群《基於ModelArts進行影像風格化繪畫》,作者: HWCloudAI 。

專案首頁 | GitHub | 論文

ModelArts 專案地址:https://developer.huaweicloud.com/develop/aigallery/notebook/detail?id=b4e4c533-e0e7-4167-94d0-4d38b9bcfd63

下載程式碼和模型

import os
import moxing as mox
mox.file.copy('obs://obs-aigallery-zc/clf/code/stylized-neural-painting.zip','stylized-neural-painting.zip')
os.system('unzip stylized-neural-painting.zip')
cd stylized-neural-painting
import argparse

import torch
torch.cuda.current_device()
import torch.optim as optim

from painter import *
# 檢測執行裝置
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 配置
parser = argparse.ArgumentParser(description='STYLIZED NEURAL PAINTING')
args = parser.parse_args(args=[])
args.img_path = './test_images/sunflowers.jpg' # 輸入圖片路徑
args.renderer = 'oilpaintbrush' # 渲染器(水彩、馬克筆、油畫筆刷、矩形) [watercolor, markerpen, oilpaintbrush, rectangle]
args.canvas_color = 'black' # 畫布底色 [black, white]
args.canvas_size = 512 # 畫布渲染尺寸,單位畫素
args.max_m_strokes = 500 # 最大筆劃數量
args.m_grid = 5 # 將圖片分割為 m_grid x m_grid 的尺寸
args.beta_L1 = 1.0 # L1 loss 權重
args.with_ot_loss = False # 設為 True 以透過 optimal transportation loss 提高收斂。但會降低生成速度
args.beta_ot = 0.1 # optimal transportation loss 權重
args.net_G = 'zou-fusion-net' # 渲染器架構
args.renderer_checkpoint_dir = './checkpoints_G_oilpaintbrush' # 預訓練模型路徑
args.lr = 0.005 # 筆劃搜尋的學習率
args.output_dir = './output' # 輸出路徑

Download pretrained neural renderer.

Define a helper funtion to check the drawing status.

def _drawing_step_states(pt):
    acc = pt._compute_acc().item()
    print('iteration step %d, G_loss: %.5f, step_psnr: %.5f, strokes: %d / %d'
          % (pt.step_id, pt.G_loss.item(), acc,
              (pt.anchor_id+1)*pt.m_grid*pt.m_grid,
              pt.max_m_strokes))
    vis2 = utils.patches2img(pt.G_final_pred_canvas, pt.m_grid).clip(min=0, max=1)

定義最佳化迴圈

def optimize_x(pt):

    pt._load_checkpoint()
    pt.net_G.eval()

    pt.initialize_params()
    pt.x_ctt.requires_grad = True
    pt.x_color.requires_grad = True
    pt.x_alpha.requires_grad = True
    utils.set_requires_grad(pt.net_G, False)

    pt.optimizer_x = optim.RMSprop([pt.x_ctt, pt.x_color, pt.x_alpha], lr=pt.lr)

    print('begin to draw...')
    pt.step_id = 0
    for pt.anchor_id in range(0, pt.m_strokes_per_block):
        pt.stroke_sampler(pt.anchor_id)
        iters_per_stroke = 20
        if pt.anchor_id == pt.m_strokes_per_block - 1:
            iters_per_stroke = 40
        for i in range(iters_per_stroke):

            pt.optimizer_x.zero_grad()

            pt.x_ctt.data = torch.clamp(pt.x_ctt.data, 0.1, 1 - 0.1)
            pt.x_color.data = torch.clamp(pt.x_color.data, 0, 1)
            pt.x_alpha.data = torch.clamp(pt.x_alpha.data, 0, 1)

            if args.canvas_color == 'white':
                pt.G_pred_canvas = torch.ones([args.m_grid ** 2, 3, 128, 128]).to(device)
            else:
                pt.G_pred_canvas = torch.zeros(args.m_grid ** 2, 3, 128, 128).to(device)

            pt._forward_pass()
            _drawing_step_states(pt)
            pt._backward_x()
            pt.optimizer_x.step()

            pt.x_ctt.data = torch.clamp(pt.x_ctt.data, 0.1, 1 - 0.1)
            pt.x_color.data = torch.clamp(pt.x_color.data, 0, 1)
            pt.x_alpha.data = torch.clamp(pt.x_alpha.data, 0, 1)

            pt.step_id += 1

    v = pt.x.detach().cpu().numpy()
    pt._save_stroke_params(v)
    v_n = pt._normalize_strokes(pt.x)
    pt.final_rendered_images = pt._render_on_grids(v_n)
    pt._save_rendered_images()

處理圖片,可能需要一些時間,建議使用 32 GB+ 視訊記憶體

pt = Painter(args=args)
optimize_x(pt)

Check out your results at args.output_dir. Before you download that folder, let’s first have a look at what the generated painting looks like.

plt.subplot(1,2,1)
plt.imshow(pt.img_), plt.title('input')
plt.subplot(1,2,2)
plt.imshow(pt.final_rendered_images[-1]), plt.title('generated')
plt.show()

請下載 args.output_dir 目錄到本地檢視高解析度的生成結果/

# 將渲染進度用動互動畫形式展現

import matplotlib.animation as animation
from IPython.display import HTML

fig = plt.figure(figsize=(4,4))
plt.axis('off')
ims = [[plt.imshow(img, animated=True)] for img in pt.final_rendered_images[::10]]

ani = animation.ArtistAnimation(fig, ims, interval=50)

# HTML(ani.to_jshtml())
HTML(ani.to_html5_video())
論文復現丨基於ModelArts進行影像風格化繪畫

Next, let’s play style-transfer. Since we frame our stroke prediction under a parameter searching paradigm, our method naturally fits the neural style transfer framework.

接下來,讓我們嘗試風格遷移,由於我們是在引數搜尋正規化下構建的筆畫預測,因此我們的方法自然的適用於神經風格遷移框架

# 配置
args.content_img_path = './test_images/sunflowers.jpg' # 輸入圖片的路徑(原始的輸入圖片)
args.style_img_path = './style_images/fire.jpg' # 風格圖片路徑
args.vector_file = './output/sunflowers_strokes.npz' # 預生成筆劃向量檔案的路徑 
args.transfer_mode = 1 # 風格遷移模式,0:顏色遷移,1:遷移顏色和紋理
args.beta_L1 = 1.0 # L1 loss 權重
args.beta_sty = 0.5 # vgg style loss 權重
args.net_G = 'zou-fusion-net' # 渲染器架構
args.renderer_checkpoint_dir = './checkpoints_G_oilpaintbrush' # 預訓練模型路徑
args.lr = 0.005 # 筆劃搜尋的學習率
args.output_dir = './output' # 輸出路徑

Again, Let’s define a helper funtion to check the style transfer status.

def _style_transfer_step_states(pt):
      acc = pt._compute_acc().item()
      print('running style transfer... iteration step %d, G_loss: %.5f, step_psnr: %.5f'
            % (pt.step_id, pt.G_loss.item(), acc))
      vis2 = utils.patches2img(pt.G_final_pred_canvas, pt.m_grid).clip(min=0, max=1)

定義最佳化迴圈

def optimize_x(pt):

    pt._load_checkpoint()
    pt.net_G.eval()

    if args.transfer_mode == 0: # transfer color only
        pt.x_ctt.requires_grad = False
        pt.x_color.requires_grad = True
        pt.x_alpha.requires_grad = False
    else: # transfer both color and texture
        pt.x_ctt.requires_grad = True
        pt.x_color.requires_grad = True
        pt.x_alpha.requires_grad = True

    pt.optimizer_x_sty = optim.RMSprop([pt.x_ctt, pt.x_color, pt.x_alpha], lr=pt.lr)

    iters_per_stroke = 100
    for i in range(iters_per_stroke):
        pt.optimizer_x_sty.zero_grad()

        pt.x_ctt.data = torch.clamp(pt.x_ctt.data, 0.1, 1 - 0.1)
        pt.x_color.data = torch.clamp(pt.x_color.data, 0, 1)
        pt.x_alpha.data = torch.clamp(pt.x_alpha.data, 0, 1)

        if args.canvas_color == 'white':
            pt.G_pred_canvas = torch.ones([pt.m_grid*pt.m_grid, 3, 128, 128]).to(device)
        else:
            pt.G_pred_canvas = torch.zeros(pt.m_grid*pt.m_grid, 3, 128, 128).to(device)

        pt._forward_pass()
        _style_transfer_step_states(pt)
        pt._backward_x_sty()
        pt.optimizer_x_sty.step()

        pt.x_ctt.data = torch.clamp(pt.x_ctt.data, 0.1, 1 - 0.1)
        pt.x_color.data = torch.clamp(pt.x_color.data, 0, 1)
        pt.x_alpha.data = torch.clamp(pt.x_alpha.data, 0, 1)

        pt.step_id += 1

    print('saving style transfer result...')
    v_n = pt._normalize_strokes(pt.x)
    pt.final_rendered_images = pt._render_on_grids(v_n)

    file_dir = os.path.join(
        args.output_dir, args.content_img_path.split('/')[-1][:-4])
    plt.imsave(file_dir + '_style_img_' +
               args.style_img_path.split('/')[-1][:-4] + '.png', pt.style_img_)
    plt.imsave(file_dir + '_style_transfer_' +
               args.style_img_path.split('/')[-1][:-4] + '.png', pt.final_rendered_images[-1])

執行風格遷移

pt = NeuralStyleTransfer(args=args)
optimize_x(pt)

高解析度生成檔案儲存在 args.output_dir

讓我們預覽一下輸出結果:

plt.subplot(1,3,1)
plt.imshow(pt.img_), plt.title('input')
plt.subplot(1,3,2)
plt.imshow(pt.style_img_), plt.title('style')
plt.subplot(1,3,3)
plt.imshow(pt.final_rendered_images[-1]), plt.title('generated')
plt.show()
論文復現丨基於ModelArts進行影像風格化繪畫

 

點選關注,第一時間瞭解華為雲新鮮技術~

相關文章