基于openCV、PyTroch视频图像实时风格迁移实现

前言

利用openCV和预训练的PyTorch风格迁移模型结合摄像头实现视频图像实时风格迁移。

 

效果展示

视频图像实时风格迁移效果.gif

没有做任何模型压缩，在我的Mac笔记本上可以跑到10FPS，基本是准实时的效果。

运行环境

openCV-4.0 以上

imutils

time

算法实现

核心算法是风格迁移，关于风格迁移的原理总结起来就是，生成一张新的图像，使其同时兼具内容图像的内容和风格图像的风格。通过提取预训练好的分类网络（VGG、ResNEt等）不同的层作为内容层和风格层，冻结网络参数调整生成的图像，在训练一定的轮次后得到结果。细节部分不在本文展开，本文重点放在生产部署上。后续抽空写文章详细补充一下原理部分。

实现过程：

首先，你要有一个USB连接的摄像头。

其次，下载预训练好的PyTorch模型。

之后，实时处理每一帧。

最后，逐帧或者跳帧保存。

 

代码部分：

# 用cv提供的dnn模块加载风格迁移模型

net = cv2.dnn.readNetFromTorch(weights)

# 设置网络的输入格式

blob = cv2.dnn.blobFromImage(current_image,1.0,(224,224),(103.939, 116.779, 123.680),swapRB=False, crop=False)

net.setInput(blob)

# 前向推理

output = net.forward()


# torch的tensor shape是（N，C，H，W）需要转换成正常图像格式（H，W，C）
 
output = output.reshape((3, output.shape[2], output.shape[3]))

# 注意前面对图像进行了逐像素去均值，所以在结果图像上需要加回来

output[0] += 103.939
output[1] += 116.779
output[2] += 123.680

output = output.transpose(1,2,0) # （H，W，C）
# 需要注意的是，得到的结果数据类型是float，cv的imshow方法不支持float类型数据，但是
# 把数据转换成uint8之后，发现图像中出现了色块，不进行转换，直接保存结果则没有这种问题。

image_name = 'images/' + str(count) + '.jpg'
cv2.imwrite(image_name, output, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
count += 1

torch模型下载

#!/bin/bash

BASE_URL="http://cs.stanford.edu/people/jcjohns/fast-neural-style/models/"

cd models
wget -N --no-check-certificate "$BASE_URL/instance_norm/candy.t7"
wget -N --no-check-certificate "$BASE_URL/instance_norm/la_muse.t7"
wget -N --no-check-certificate "$BASE_URL/instance_norm/mosaic.t7"
wget -N --no-check-certificate "$BASE_URL/instance_norm/feathers.t7"
wget -N --no-check-certificate "$BASE_URL/instance_norm/the_scream.t7"
wget -N --no-check-certificate "$BASE_URL/instance_norm/udnie.t7"

GitHub 图像风格迁移