iOS實(shí)現(xiàn)類(lèi)Prisma軟件

前言

Prisma在2016上線后就大火，該APP是利用神經(jīng)網(wǎng)絡(luò)和人工智能技術(shù)，為普通照片加入藝術(shù)效果的照片編輯軟件。

同年Google也發(fā)布了一篇《A LEARNED REPRESENTATION FOR ARTISTIC STYLE》論文，實(shí)現(xiàn)了前向運(yùn)算一次為照片整合多種藝術(shù)風(fēng)格的功能，并且優(yōu)化了內(nèi)存使用和運(yùn)算速度，可以在移動(dòng)設(shè)備上快速運(yùn)算。

最近在研究Tensorflow整合iOS過(guò)程中，發(fā)現(xiàn)google公開(kāi)了論文實(shí)現(xiàn)的源碼和訓(xùn)練數(shù)據(jù)，也就是說(shuō)我們可以通過(guò)自己寫(xiě)一個(gè)前向運(yùn)算圖，整合其訓(xùn)練參數(shù)就可以快速實(shí)現(xiàn)類(lèi)Prisma的應(yīng)用。

下面就介紹一下如何在iPhone上跑一個(gè)自己的"Prisma"。

招財(cái)和咕嚕

準(zhǔn)備工作

  1. 安裝Tensorflow，這個(gè)官網(wǎng)上有詳細(xì)教程這里就不多說(shuō)了。
  2. 搭建iOS+Tensorflow工程，這個(gè)可以根據(jù)Git上的步驟實(shí)現(xiàn)，也可以參考官方的Demo程序配置。（這個(gè)過(guò)程有很多坑，多次嘗試，應(yīng)該可以配置成功）
  3. 下載模型，本次使用的模型是image_stylization，google已開(kāi)源在GitHub上。
  4. 下載訓(xùn)練好的參數(shù)，Google提供了2個(gè)：
    Monet
    Varied
    Monet訓(xùn)練了10種藝術(shù)圖片，Varied訓(xùn)練了32種。
    當(dāng)然你也可以自己訓(xùn)練藝術(shù)圖片，但是得下載VGG的訓(xùn)練參數(shù)和ImageNet數(shù)據(jù)，然后自己訓(xùn)練，比較花時(shí)間。

構(gòu)建計(jì)算圖

雖然Google提供了模型的源碼，但是并沒(méi)有在源碼中輸出運(yùn)算圖以方便遷移到移動(dòng)設(shè)備中使用，Android的Demo中倒是提供了生成的pb，如果覺(jué)得自己寫(xiě)計(jì)算圖麻煩可以直接拷到自己iOS工程中使用。

我這里創(chuàng)建了一個(gè)python的工程，然后把Google源碼中model.py相關(guān)的文件都加入了工程。
我的建圖代碼如下：

import numpy as np
import tensorflow as tf
import ast
import os
from tensorflow.python import pywrap_tensorflow

from matplotlib import pyplot
from matplotlib.pyplot import imshow

import image_utils
import model
import ops
import argparse
import sys


num_styles = 32
imgWidth = 512
imgHeight = 512
channel = 3
checkpoint = "/Users/Jiao/Desktop/TFProject/style-image/checkpoint/multistyle-pastiche-generator-varied.ckpt"

inputImage = tf.placeholder(tf.float32,shape=[None,imgWidth,imgHeight,channel],name="input")
styles = tf.placeholder(tf.float32,shape=[num_styles],name="style")

with tf.name_scope(""):
    transform = model.transform(inputImage,
                            normalizer_fn=ops.weighted_instance_norm,
                            normalizer_params={
                                # 'weights': tf.constant(mixture),
                                'weights' : styles,
                                'num_categories': num_styles,
                                'center': True,
                                'scale': True})

model_saver = tf.train.Saver(tf.global_variables())

with tf.Session() as sess:
    tf.train.write_graph(sess.graph_def, "/Users/Jiao/Desktop/TFProject/style-image/protobuf", "input.pb")
    #checkpoint = os.path.expanduser(checkpoint)
    #if tf.gfile.IsDirectory(checkpoint):
    #    checkpoint = tf.train.latest_checkpoint(checkpoint)
    #    tf.logging.info('loading latest checkpoint file: {}'.format(checkpoint))
    #model_saver.restore(sess, checkpoint)

    #newstyle = np.zeros([num_styles], dtype=np.float32)
    #newstyle[18] = 0.5
    #newstyle[17] = 0.5
    #newImage = np.zeros((1,imgWidth,imgHeight,channel))
    #style_image = transform.eval(feed_dict={inputImage:newImage,styles:newstyle})
    #style_image = style_image[0]
    #imshow(style_image)
    #pyplot.show()

這里輸入節(jié)點(diǎn)是input和style，輸出節(jié)點(diǎn)是model中的transformer/expand/conv3/conv/Sigmoid。

到此就將模型的計(jì)算圖保存到了本地文件夾中。
接下來(lái)就是將圖和ckpt中的參數(shù)合并，并且生成移動(dòng)端的可以使用的pb文件，這一步可以參考我上一篇文章《iOS+Tensorflow實(shí)現(xiàn)圖像識(shí)別》，很容易就實(shí)現(xiàn)。

iOS工程

在上面準(zhǔn)備工作中，如果你已經(jīng)按步驟搭建好iOS+TF的工程，這里你只需要導(dǎo)入生成的最終pb文件就行了。工程結(jié)構(gòu)如圖：

XCode工程

然后在iOS使用pb文件，我這里直接導(dǎo)入了Google提供的tensorflow_utils，使用這個(gè)類(lèi)里面的LoadModel方法可以很快的生成含有計(jì)算圖的session。

- (void)viewDidLoad {
    [super viewDidLoad];
    tensorflow::Status load_status;
    load_status = LoadModel(@"rounded_graph", @"pb", &tf_session);
    if (!load_status.ok()) {
        LOG(FATAL) << "Couldn't load model: " << load_status;
    }
    currentStyle = 0;
    isDone = true;
    _styleImageView.layer.borderColor = [UIColor grayColor].CGColor;
    _styleImageView.layer.borderWidth = 0.5;
    _ogImageView.layer.borderColor = [UIColor grayColor].CGColor;
    _ogImageView.layer.borderWidth = 0.5;
}

最后就是獲取圖片，執(zhí)行運(yùn)算，生成藝術(shù)圖片展示。這里圖片需要轉(zhuǎn)換成bitmap然后獲取data值，展示圖片也是相識(shí)的過(guò)程。具體代碼如下：

- (void)runCnn:(UIImage *)compressedImg
{
    unsigned char *pixels = [self getImagePixel:compressedImg];
    int image_channels = 4;
    tensorflow::Tensor image_tensor(
                                    tensorflow::DT_FLOAT,
                                    tensorflow::TensorShape(
                                                            {1, wanted_input_height, wanted_input_width, wanted_input_channels}));
    auto image_tensor_mapped = image_tensor.tensor<float, 4>();
    tensorflow::uint8 *in = pixels;
    float *out = image_tensor_mapped.data();
    for (int y = 0; y < wanted_input_height; ++y) {
        float *out_row = out + (y * wanted_input_width * wanted_input_channels);
        for (int x = 0; x < wanted_input_width; ++x) {
            tensorflow::uint8 *in_pixel =
            in + (x * wanted_input_width * image_channels) + (y * image_channels);
            float *out_pixel = out_row + (x * wanted_input_channels);
            for (int c = 0; c < wanted_input_channels; ++c) {
                out_pixel[c] = in_pixel[c];
            }
        }
    }
    
    
    tensorflow::Tensor style(tensorflow::DT_FLOAT, tensorflow::TensorShape({32}));
    float *style_data = style.tensor<float, 1>().data();
    memset(style_data, 0, sizeof(float) * 32);
    style_data[currentStyle] = 1;
    
    if (tf_session.get()) {
        std::vector<tensorflow::Tensor> outputs;
        tensorflow::Status run_status = tf_session->Run(
                                                        {{contentNode, image_tensor},
                                                            {styleNode, style}},
                                                        {outputNode},
                                                        {},
                                                        &outputs);
        if (!run_status.ok()) {
            LOG(ERROR) << "Running model failed:" << run_status;
            isDone = true;
            free(pixels);
        } else {
            float *styledData = outputs[0].tensor<float,4>().data();
            UIImage *styledImg = [self createImage:styledData];
            dispatch_async(dispatch_get_main_queue(), ^{
                _styleImageView.image = styledImg;
                dispatch_after(dispatch_time(DISPATCH_TIME_NOW, (int64_t)(0.3 * NSEC_PER_SEC)), dispatch_get_main_queue(), ^{
                    isDone = true;
                    free(pixels);
                });
            });
        }
    }
}

- (unsigned char *)getImagePixel:(UIImage *)image
{
    int width = image.size.width;
    int height = image.size.height;
    CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
    unsigned char *rawData = (unsigned char*) calloc(height * width * 4, sizeof(unsigned char));
    NSUInteger bytesPerPixel = 4;
    NSUInteger bytesPerRow = bytesPerPixel * width;
    NSUInteger bitsPerComponent = 8;
    CGContextRef context = CGBitmapContextCreate(rawData, width, height,
                                                 
                                                 bitsPerComponent, bytesPerRow, colorSpace,
                                                 
                                                 kCGImageAlphaPremultipliedLast | kCGBitmapByteOrder32Big);
    
    CGColorSpaceRelease(colorSpace);
    CGContextDrawImage(context, CGRectMake(0, 0, width, height), image.CGImage);
    UIImage *ogImg = [UIImage imageWithCGImage:CGBitmapContextCreateImage(context)];
    dispatch_async(dispatch_get_main_queue(), ^{
        _ogImageView.image = ogImg;
    });
    CGContextRelease(context);
    return rawData;
}

- (UIImage *)createImage:(float *)pixels
{
    unsigned char *rawData = (unsigned char*) calloc(wanted_input_height * wanted_input_width * 4, sizeof(unsigned char));
    for (int y = 0; y < wanted_input_height; ++y) {
        unsigned char *out_row = rawData + (y * wanted_input_width * 4);
        for (int x = 0; x < wanted_input_width; ++x) {
            float *in_pixel =
            pixels + (x * wanted_input_width * 3) + (y * 3);
            unsigned char *out_pixel = out_row + (x * 4);
            for (int c = 0; c < wanted_input_channels; ++c) {
                out_pixel[c] = in_pixel[c] * 255;
            }
            out_pixel[3] = UINT8_MAX;
        }
    }
    CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
    NSUInteger bytesPerPixel = 4;
    NSUInteger bytesPerRow = bytesPerPixel * wanted_input_width;
    NSUInteger bitsPerComponent = 8;
    CGContextRef context = CGBitmapContextCreate(rawData, wanted_input_width, wanted_input_height,
                                                 
                                                 bitsPerComponent, bytesPerRow, colorSpace,
                                                 
                                                 kCGImageAlphaPremultipliedLast | kCGBitmapByteOrder32Big);
    
    CGColorSpaceRelease(colorSpace);
    UIImage *retImg = [UIImage imageWithCGImage:CGBitmapContextCreateImage(context)];
    CGContextRelease(context);
    free(rawData);
    return retImg;
}

這里說(shuō)明一下，前面python工程已經(jīng)定義了，我的輸入和輸出圖片的大小是512?512。

連接iPhone，運(yùn)行工程_

最后連上手機(jī)運(yùn)行，就可以自己創(chuàng)建自己的藝術(shù)類(lèi)圖片了。??

放幾張運(yùn)行效果圖：


截圖1
截圖2
截圖3

源碼已開(kāi)源：https://github.com/JiaoLiu/style-image/tree/master

最后編輯于 ：
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點(diǎn)，簡(jiǎn)書(shū)系信息發(fā)布平臺(tái)，僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

  • 前言 前面寫(xiě)了一篇利用TF在iOS實(shí)現(xiàn)類(lèi)Prisma軟件的文章后，收到很多網(wǎng)上朋友交流實(shí)現(xiàn)思想與求Demo回復(fù)，可...
    Jiao123閱讀 1,852評(píng)論 9贊 21
  • 推理過(guò)程中很容易出現(xiàn)各種各樣的謬誤，例如： 偷梁換柱謬誤：指在論證中關(guān)鍵詞語(yǔ)有兩種以上的含義，一旦不同含義之間的轉(zhuǎn)...
    飛鷹于凱閱讀 309評(píng)論 0贊 0
  • 轉(zhuǎn)眼間又到了端午節(jié) 各位客官可 三五成群，盡情吃吃粽子 但有一樣事要記住 千萬(wàn)不要 啪啪啪 （此處省略一萬(wàn)多個(gè)字…...
    葉姐分享閱讀 4,716評(píng)論 0贊 0

友情鏈接更多精彩內(nèi)容