基于百度PaddleHub实现视频/图片风格转换

百度PaddleHub官网
stylepro_artistic模型介绍
项目地址

项目说明

该项目基于百度PaddleHub提供的stylepro_artistic模型开发，实现将视频或图片转换为多种风格的功能。

模型概述

艺术风格迁移模型可以将给定的图像转换为任意的艺术风格。本模型StyleProNet整体采用全卷积神经网络架构(FCNs)，通过encoder-decoder重建艺术风格图片。StyleProNet的核心是无参数化的内容-风格融合算法Style Projection，模型规模小，响应速度快。模型训练的损失函数包含style loss、content perceptual loss以及content KL loss，确保模型高保真还原内容图片的语义细节信息与风格图片的风格信息。预训练数据集采用MS-COCO数据集作为内容端图像，WikiArt数据集作为风格端图像。

风格预览

• 原图
  

• 梵高抽象派风格
  

• 水墨风格
  

• 动画风格
  

• 梵高风格视频

  使用百度PaddleHub生成梵高风格视频

• 动画风格视频

  使用百度PaddleHub生成动画风格视频

• 水墨风格视频
  

  使用百度PaddleHub生成水墨画风格视频

开发过程

百度PaddleHub提供了完整的模型已经极其简单的调用接口，基本使用方法见stylepro_artistic模型介绍

1. 安装stylepro_artistic模型

   $ hub install stylepro_artistic==1.0.0
   

2. 定义风格转换类

   import paddlehub as hub
   import cv2
   
   class PicTransfer:
       def __init__(self):
           self.fangao_style = './pic/fangao.jpg'
           self.donghua_style = './pic/donghua2.jpg'
           self.shuimo_style = './pic/shuimo.jpg'
           self.stylepro_artistic = hub.Module(name="stylepro_artistic")
       
       def transfer(self, picture, style='Fangao'):
           if style == 'Fangao':
               transfer_style = self.fangao_style
           elif style == 'Donghua':
               transfer_style = self.donghua_style
           elif style == 'Shuimo':
               transfer_style = self.shuimo_style
           else:
               return null
           result = self.stylepro_artistic.style_transfer(
               images=[{
                   'content': cv2.imread(picture),
                   'styles': [cv2.imread(transfer_style)]
               }])
           return result
   

3. 执行风格转换

   pic_transfer = PicTransfer()
   style_list = ['Fangao', 'Donghua', 'Shuimo']
   picture = './pic/fengjing_big.jpg'
   result_list = []
   for style in style_list:
       result = pic_transfer.transfer(picture, style)
       result_list.append(result)
   

4. 可视化

   import matplotlib.pyplot as plt
   from matplotlib.image import imread
   plt.figure(figsize=(120,30))
   plt.subplot(411)
   plt.title(picture)
   plt.imshow(imread(picture))
   plt.axis('off')
   plt.subplot(412)
   plt.title(style_list[0]+'-'+picture)
   plt.imshow(result_list[0][0]['data'])
   plt.axis('off')
   plt.subplot(413)
   plt.title(style_list[1]+'-'+picture)
   plt.imshow(result_list[1][0]['data'])
   plt.axis('off')
   plt.subplot(414)
   plt.title(style_list[2]+'-'+picture)
   plt.imshow(result_list[2][0]['data'])
   plt.axis('off')
   plt.show()
   

5. 视频处理：提取视频帧为图片

   ## 将视频转换为图片
   import os
   import cv2
   
   # 原视频读取
   video = './video/山水视频.mp4'
   frame_path = './video/山水视频/'
   if not os.path.exists(frame_path):
       os.mkdir(frame_path)
   vc = cv2.VideoCapture(video)
   
   if vc.isOpened():
       rval, frame = vc.read()
   else:
       rval = False
   
   i = 0
   while rval:
       rval, frame = vc.read()
       cv2.imwrite(frame_path+'{}.jpg'.format(i), frame)
       i += 1
   print('视频帧数：{}'.format(i))
   FPS = vc.get(cv2.CAP_PROP_FPS)
   SIZE = (int(vc.get(cv2.CAP_PROP_FRAME_WIDTH)), 
           int(vc.get(cv2.CAP_PROP_FRAME_HEIGHT)))
   print('FPS:{}, SIZE:{}'
       .format(FPS,SIZE))
   

6. 转换提取的图片并合成为新的视频

   from tqdm import tqdm
   pic_transfer = PicTransfer()
   style = 'Shuimo'
   save_video = './video/山水视频_'+style+'.mp4'
   fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
   vw = cv2.VideoWriter(save_video, fourcc, FPS, SIZE, True)
   files = os.listdir(frame_path)
   for i in  tqdm(range(len(files)-1)):
       picture = os.path.join(frame_path,str(i)+'.jpg')
       result = pic_transfer.transfer(picture, style)
       data = result[0]['data']
       vw.write(data)
   vw.release()
   cv2.destroyAllWindows()
   

心得

百度PaddleHub提供了很多预训练模型，基于这些模型可以实现很多有趣又实用的功能
此项目的核心部分仅用了不足10行代码就能实现强大的功能，简单方便
希望百度Paddle Paddle越来越强大，自己能力也越来越强