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

没有做任何模型压缩,在我的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 图像风格迁移