Python CNN风格迁移

本文仅供学习交流使用，如侵立删！联系方式及demo下载见文末

---

环境：
python：3.6.7
tensorflow=2.2.0
torch=1.5.1+cpu
torchvision=0.6.0+cpu
Pillow=8.1.2

效果：

---

输入原图和风格图生成迁移图片

# -*- coding: utf-8 -*-
import cv2
import time


def style_transfer(pathIn='',
                   pathOut='',
                   model='',
                   width=None,
                   jpg_quality=80):
    '''
    pathIn: 原始图片的路径
    pathOut: 风格化图片的保存路径
    model: 预训练模型的路径
    width: 设置风格化图片的宽度，默认为None, 即原始图片尺寸
    jpg_quality: 0-100，设置输出图片的质量，默认80，越大图片质量越好
    '''

    # 读入原始图片，调整图片至所需尺寸，然后获取图片的宽度和高度
    img = cv2.imread(pathIn)
    (h, w) = img.shape[:2]
    if width is not None:
        img = cv2.resize(img, (width, round(width * h / w)), interpolation=cv2.INTER_CUBIC)
        (h, w) = img.shape[:2]

    # 从本地加载预训练模型
    print('加载预训练模型......')
    net = cv2.dnn.readNetFromTorch(model)

    # 将图片构建成一个blob：设置图片尺寸，将各通道像素值减去平均值（比如ImageNet所有训练样本各通道统计平均值）
    # 然后执行一次前馈网络计算，并输出计算所需的时间
    blob = cv2.dnn.blobFromImage(img, 1.0, (w, h), (103.939, 116.779, 123.680), swapRB=False, crop=False)
    net.setInput(blob)
    start = time.time()
    output = net.forward()
    end = time.time()
    print("风格迁移花费：{:.2f}秒".format(end - start))

    # reshape输出结果, 将减去的平均值加回来，并交换各颜色通道
    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)

    # 输出风格化后的图片
    cv2.imwrite(pathOut, output, [int(cv2.IMWRITE_JPEG_QUALITY), jpg_quality])

---

CNN模型训练

import __future__
import torch
import torchvision.models as models
import torchvision.transforms as transforms
import time
import os
from PIL import Image

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
'''
环境：
python：3.6.7
tensorflow==2.2.0
torch==1.5.1+cpu
torchvision==0.6.0+cpu
Pillow==8.1.2
'''

class GramMatrix(torch.nn.Module):
    def forward(self, input):
        b, n, h, w = input.size()
        features = input.view(b * n, h * w)
        G = torch.mm(features, features.t())
        return G.div(b * n * h * w)


class StyleLoss(torch.nn.Module):
    def __init__(self, style_feature, weight):
        super(StyleLoss, self).__init__()
        self.style_feature = style_feature.detach()
        self.weight = weight
        self.gram = GramMatrix()
        self.criterion = torch.nn.MSELoss()

    def forward(self, combination):
        # output = combination
        style_feature = self.gram(self.style_feature.clone() * self.weight)
        combination_features = self.gram(combination.clone() * self.weight)
        self.loss = self.criterion(combination_features, style_feature)
        return combination


class StyleTransfer:
    def __init__(self, content_image, style_image, style_weight=5, content_weight=0.025):
        # Weights of the different loss components
        self.vgg19 = models.vgg19()
        self.vgg19.load_state_dict(torch.load('vgg19-dcbb9e9d.pth'))
        self.img_ncols = 1280
        self.img_nrows = 720
        self.style_weight = style_weight
        self.content_weight = content_weight
        # 处理原图和风格图片
        self.content_tensor, self.content_name = self.process_img(content_image)
        self.style_tensor, self.style_name = self.process_img(style_image)
        self.conbination_tensor = self.content_tensor.clone()

    def process_img(self, img_path):
        img = Image.open(img_path)
        img_name = img_path.split('/')[-1][:-4]
        loader = transforms.Compose([transforms.Resize((self.img_nrows, self.img_ncols)),
                                     transforms.ToTensor()])
        img_tensor = loader(img)
        img_tensor = img_tensor.unsqueeze(0)
        return img_tensor.to(device, torch.float), img_name

    def deprocess_img(self, x, index):
        unloader = transforms.ToPILImage()
        x = x.cpu().clone()
        img_tensor = x.squeeze(0)
        img = unloader(img_tensor)
        result_folder = f'{self.content_name}_and_{self.style_name}'
        os.path.exists(result_folder) or os.mkdir(result_folder)
        filename = f'{result_folder}/rersult_{index}.png'
        img.save(filename)
        print(f'save {filename} successfully!')
        print()

    def get_loss_and_model(self, vgg_model, content_image, style_image):
        vgg_layers = vgg_model.features.to(device).eval()
        style_losses = []
        content_losses = []
        model = torch.nn.Sequential()
        style_layer_name_maping = {
     
            '0': "style_loss_1",
            '5': "style_loss_2",
            '10': "style_loss_3",
            '19': "style_loss_4",
            '28': "style_loss_5",
        }
        content_layer_name_maping = {
     '30': "content_loss"}
        for name, module in vgg_layers._modules.items():
            model.add_module(name, module)
            if name in content_layer_name_maping:
                content_feature = model(content_image).clone()
                content_loss = ContentLoss(content_feature, self.content_weight)
                model.add_module(f'{content_layer_name_maping[name]}', content_loss)
                content_losses.append(content_loss)
            if name in style_layer_name_maping:
                style_feature = model(style_image).clone()
                style_loss = StyleLoss(style_feature, self.style_weight)
                style_losses.append(style_loss)
                model.add_module(f'{style_layer_name_maping[name]}', style_loss)
        return content_losses, style_losses, model

    def get_input_param_optimizer(self, input_img):
        input_param = torch.nn.Parameter(input_img.data)
        optimizer = torch.optim.LBFGS([input_param])
        return input_param, optimizer

    def main_train(self, epoch=10):
        print('Load model preprocessing')
        combination_param, optimizer = self.get_input_param_optimizer(self.conbination_tensor)
        content_losses, style_losses, model = self.get_loss_and_model(self.vgg19, self.content_tensor,
                                                                      self.style_tensor)
        cur, pre = 10, 10
        for i in range(1, epoch + 1):
            start = time.time()

            def closure():
                combination_param.data.clamp_(0, 1)
                optimizer.zero_grad()
                model(combination_param)
                style_score = 0
                content_score = 0
                for cl in content_losses:
                    content_score += cl.loss
                for sl in style_losses:
                    style_score += sl.loss
                loss = content_score + style_score
                loss.backward()
                return style_score + content_score

            loss = optimizer.step(closure)
            cur, pre = loss, cur
            end = time.time()
            print(f'|using:{int(end - start):2d}s |epoch:{i:2d} |loss:{loss.data}')
            if pre <= cur:
                print('Early stopping!')
                break
            combination_param.data.clamp_(0, 1)

---

本文仅供学习交流使用，如侵立删！ 企鹅 、WX： 1033383881（备注来意）

---