教程帖：用TensorFlow自制Taylor Swift识别器

作者：Sara Robinson

编译：Bing

编者按：本文作者Sara Robinson在Medium上发布了一个有趣的项目，她自制了一款APP，能自动识别歌手Taylor Swift。这与我们之前介绍的寻找威利项目很像。该教程非常详细，有兴趣的同学可以学习一下，动手做一个自己的图像识别工具哦~本文已获作者授权，以下是论智对原文的编译。

注：由于写作本文时TensorFlow没有Swift库，我用Swift构建了针对我的模型的预测请求的APP。

以下就是我们创建的APP：

TensorFlow物体检测API能让你识别出一张图片中特定物体的位置，这可以应用到许多有趣的程序上。不过我平常拍人比较多，所以就想把这一技术应用到人脸识别上。结果发现模型表现得非常好！也就是上图我创建的Taylor Swift检测器。

本文将列出模型的构建步骤，从收集Taylor Swift的照片到模型的训练：

• 对图像进行预处理，改变大小、贴标签、将它们分成训练和测试两部分，并修改成Pascal VOC格式；

• 将图片转化成TFRecords文件以符合物体检测API；

• 利用MobileNet在谷歌Cloud ML Engine上训练模型；

• 导出训练好的模型并将其部署到ML Engine上进行服务；

• 构建一个iOS前端，根据训练好的模型做出预测请求（使用Swift）。

下面是各部分如何结合在一起的架构图：

相关代码可访问GitHub找到，地址：https://github.com/sararob/tswift-detection

在开始之前，首先要解释一下我们即将用到的技术和术语：TensorFlow物体检测API是一个构建在TensorFlow上的框架，用于识别图像中特定的对象。例如，你可以用很多猫的照片训练它，一旦训练完毕，你可以输入一张猫的图像，它就会输出一个方框列表，认为图像中有一只猫。虽然它的名字中含有API，但是你可以将它更多地想象成用于迁移学习的一套便利的工具。

但是，训练模型识别图像中的对象是个费时费力的活。物体检测最酷的地方就是它支持五个预训练模型的迁移学习（transfer learning），那么什么是迁移学习呢？比如，当儿童学习第一门语言时，他们会接触大量的例子，如果有错就会立刻被纠正过来。例如当孩子们学习识别猫时，他们的父母会指着图片上的猫，并说出“猫”这个词，这种重复增强了他们的脑回路。当它们学习如何识别一只狗时，无需从头开始，这一过程与猫的识别类似，只是学习对象不同。这就是迁移学习的工作原理。

但我没有时间寻找并标记数千个Taylor Swift的图像，但是我可以通过修改最后几个图层、在数百万张图像上训练的模型中提取特征，应用于TSwift的检测。

第一步：预处理图像

首先要感谢Dat Tran写的关于浣熊检测器的博客，地址：https://towardsdatascience.com/how-to-train-your-own-object-detector-with-tensorflows-object-detector-api-bec72ecfe1d9

首先，我从谷歌图片上下载了200张Taylor Swift的照片，这里安利一个Chrome插件：Fatkun Batch Download Image，可以下载所有图片搜索结果。在打标签之前，我把图片分为两类：训练和测试。另外，我写了一个调整图片大小的脚本（https://github.com/sararob/tswift-detection/blob/master/resize.py），确保每张图的宽度不超过600px。

由于检测器会告诉我们图中的对象位置，所以你不能直接把图像和标签作为训练数据。你需要用边框将对象圈出来，以及将表框打上标签（在我们的数据集中，只需要一个标签tswift）。

打边框工具依然使用LabelImg，这是一个基于Python的程序，你只需输入带标签的图像，它就会输出一个xml文件，将每张照片都打上边框同时还有相关标签（不到一上午我就处理好200张图片了）。下面是它如何工作的（标签输入为tswift）：

然后LabelImg生成一个xml文件：

<annotation>
    <folder>Desktop</folder>
    <filename>tswift.jpg</filename>
    <path>/Desktop/tswift.jpg</path>
    <source>
        <database>Unknown</database>
    </source>
    <size>
        <width>1000</width>
        <height>667</height>
        <depth>3</depth>
    </size>
    <segmented>0</segmented>
    <object>
        <name>tswift</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>420</xmin>
            <ymin>80</ymin>
            <xmax>582</xmax>
            <ymax>291</ymax>
        </bndbox>
    </object>
</annotation>

现在我有了一张带有边框和标签的图片了，但是我还要把它转换成TensorFlow可接受的方式——一个数据的二进制表示TFRecord。关于这一方法可以在GitHub上查看。要运行我的脚本，你需要先下载一个tensorflow/models，从tensorflow/models/research本地直接运行脚本，带上以下参数（运行两次：一次用于训练数据，一次用于测试数据）

python convert_labels_to_tfrecords.py \
--output_path=train.record \ 
--images_dir=path/to/your/training/images/ \
--labels_dir=path/to/training/label/xml/

第二步：训练检测器

我可以在笔记本电脑上训练这个模型，但是时间会很长，而且占用大量的资源。并且一旦我需要用电脑做别的事，训练就会中断。所以，我选择了云！我们可以利用云来运行多个跨核心的训练，几个小时内就能完成整个工作，并且用Cloud ML engine的速度比GPU还要快。

设置Cloud ML Engine

我准备将所有TFRecord格式的数据上传到云并开始训练。首先，我在谷歌云端控制台中创建了一个项目，并启用了Cloud ML Engine：

然后，我将创建一个云存储bucket来打包模型的所有资源。确保在指定区域进行存储（不要选择多个区域）：

我将在这个bucket中/data子目录来放置训练和测试TFRecord的文件：

目标对象检测API还需要一个将标签映射到整数ID的pbtxt文件。由于我们只有一个标签，这个是非常短的：

item {
  id: 1
  name: 'tswift'
}

添加MobileNet检查点进行迁移学习

因为我并非从零开始训练这个模型，所以当我运行训练时，我需要指向我将要建立的预训练模型。我选择使用MobileNet模型——它是针对移动设备优化的一系列小模型。虽然我不会直接在移动设备上训练模型，但MobileNet将会快速训练，并允许更快的预测请求。我下载了这个MobileNet检查点用于训练，检查点是一个二进制文件，包含训练过程中特定点的TensorFlow模型的状态。下载并解压缩后，你可以看到它包含的三个文件：

以上所有都要用来训练模型，所以我将它们放在云存储bucket中的同一个data/目录中。

在开始训练之前，还需要添加一个文件。对象检测脚本需要一种方法查找模型的检查点、标签映射和训练数据。我们将用配置文件处理这一点。TF对象检测为五个预训练模型采集了样本配置文件。我们在这里为MobileNet使用一个，并且在云存储bucket的相应路径中更新了所有PATH_TO_BE_CONFIGURED占位符。除了将我的模型连接到云存储中的数据外，此文件还为我的模型配置了几个超参数，如卷积大小、激活函数和步骤。

以下是开始训练之前云存储bucket中我的/data中的所有文件：

我还会在bucket中创建train/和eval/子目录——这是TensorFlow在训练和评估时书写模型检查点文件的地方。

现在已经准备好训练了，通过执行gcloud命令开始。请注意，你需要在本地复制tensorflow/models/research并从该目录运行此训练脚本：

# Run this script from tensorflow/models/research: 
gcloud ml-engine jobs submit training ${YOUR_TRAINING_JOB_NAME} \
    --job-dir=${YOUR_GCS_BUCKET}/train \
    --packages dist/object_detection-0.1.tar.gz,slim/dist/slim-0.1.tar.gz \
    --module-name object_detection.train \
    --region us-central1 \
    --config object_detection/samples/cloud/cloud.yml \
    --runtime-version=1.4
    -- \
    --train_dir=${YOUR_GCS_BUCKET}/train \
    --pipeline_config_path=${YOUR_GCS_BUCKET}/data/ssd_mobilenet_v1_coco.config

训练的同时，我也开始了评估工作。我会使用之前从未见过的数据来评估模型的准确性：

# Run this script from tensorflow/models/research: 
gcloud ml-engine jobs submit training ${YOUR_EVAL_JOB_NAME} \
    --job-dir=${YOUR_GCS_BUCKET}/train \
    --packages dist/object_detection-0.1.tar.gz,slim/dist/slim-0.1.tar.gz \
    --module-name object_detection.eval \
    --region us-central1 \
    --scale-tier BASIC_GPU \
    --runtime-version=1.4
    -- \
    --checkpoint_dir=${YOUR_GCS_BUCKET}/train \
    --eval_dir=${YOUR_GCS_BUCKET}/eval \
    --pipeline_config_path=${YOUR_GCS_BUCKET}/data/ssd_mobilenet_v1_coco.config

你可以通过在云端控制台导航到ML Engine的“作业”部分来验证您的任务是否正确运行，并检查日志以查找特定作业：

第三步：部署预测模型

为了将模型部署到ML Engine，我需要将模型检查点转换为ProtoBuf。在我的train/bucket中，可以看到从几处保留的检查点文件：

文件的第一行告诉我最新的检查点路径——我应从该检查点本地下载3个文件。每个检查点应该有一个.index，.meta，和.data文件。将它们保存在本地目录中后，我可以使用对象检测的export_inference_graph脚本将它们转换为ProtoBuf。要运行以下脚本，你需要定义MobileNet配置文件的本地路径、训练时下载的模型检查点编号以及要导出的图形目录名称：

# Run this script from tensorflow/models/research: 
python object_detection/export_inference_graph.py \
    --input_type encoded_image_string_tensor \
    --pipeline_config_path ${LOCAL_PATH_TO_MOBILENET_CONFIG} \
    --trained_checkpoint_prefix model.ckpt-${CHECKPOINT_NUMBER} \
    --output_directory ${PATH_TO_YOUR_OUTPUT}.pb

这个脚本运行后，你将会在.pb输出目录中看到一个saved_model/目录。将saved_model.pb文件上传到你的云存储/data目录中（不要担心生成其他文件）。

现在你已经准备好将模型部署到ML Engine上了。首先，用gcloud创建你的模型：

gcloud ml-engine models create tswift_detector

然后，通过将模型指向刚刚上传到云存储的已保存的ProtoBuf来创建第一个模型版本：

gcloud ml-engine versions create v1 --model=tswift_detector --origin=gs://${YOUR_GCS_BUCKET}/data  --runtime-version=1.4

模型部署好后，我将用ML Engine的线上预测API生成新的预测图像。

第四步：使用Firebase函数和Swift构建预测客户端

我在Swift中编写了一个iOS客户端来对我的模型进行预测请求。Swift客户端将图像上传到云存储，云存储触发Firebase函数，在Node.js中发起预测请求，并将生成的预测图像和数据保存到云存储和Firebase中。

首先，在我的Swift客户端中，我添加了一个按钮，供用户访问设备的图片库。用户选择照片后，会触发将图像上传到云端存储的操作：

let firestore = Firestore.firestore()
func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [String : Any]) {
    let imageURL = info[UIImagePickerControllerImageURL] as? URL
    let imageName = imageURL?.lastPathComponent
    let storageRef = storage.reference().child("images").child(imageName!)
    storageRef.putFile(from: imageURL!, metadata: nil) { metadata, error in
        if let error = error {
            print(error)
        } else {
            print("Photo uploaded successfully!")
            // TODO: create a listener for the image's prediction data in Firestore
            }
        }
    }
    dismiss(animated: true, completion: nil)
}

接下来，我编写了在上传到云存储时触发的Firebase函数（https://github.com/sararob/tswift-detection/blob/master/firebase/functions/index.js）。下面的代码也包含了我向ML Engine预测API发出请求的函数部分：

function cmlePredict(b64img, callback) {
    return new Promise((resolve, reject) => {
        google.auth.getApplicationDefault(function (err, authClient, projectId) {
            if (err) {
                reject(err);
            }
            if (authClient.createScopedRequired && authClient.createScopedRequired()) {
                authClient = authClient.createScoped([
                    'https://www.googleapis.com/auth/cloud-platform'
                ]);
            }
            var ml = google.ml({
                version: 'v1'
            });
            const params = {
                auth: authClient,
                name: 'projects/sara-cloud-ml/models/tswift_detector',
                resource: {
                    instances: [
                    {
                        "inputs": {
                        "b64": b64img
                        }
                    }
                    ]
                }
            };
            ml.projects.predict(params, (err, result) => {
                if (err) {
                    reject(err);
                } else {
                    resolve(result);
                }
            });
        });
    });
}

在ML Engine的反应中，我们得到：

• detection_boxes：可以用来标出Taylor Swift周围的边框；

• detection_scores：为每个检测框架返回一个置信度值，其中只包括分数高于70%的检测；

• detection_classes：告诉我们与检测相关的ID。在这种情况下，因为只有一个标签所以该值总为1。

在函数中，如果检测到Taylor，则用detection_boxes在图像中绘制一个边框以及生成置信度分数。然后将新的带有边框的图像保存到云中，将图像的文件路径写入Cloud Firestore，一边在iOS应用程序中读取路径并下载新图像：

const admin = require('firebase-admin');
admin.initializeApp(functions.config().firebase);
const db = admin.firestore();
let outlinedImgPath = `outlined_img/${filePath.slice(7)}`;
let imageRef = db.collection('predicted_images').doc(filePath);
imageRef.set({
    image_path: outlinedImgPath,
    confidence: confidence
});
bucket.upload('/tmp/path/to/new/image', {destination: outlinedImgPath});

最后，在iOS应用程序中，我们可以监测图像Firestore路径的更新。如果检测到目标，我会下载这张图片并在应用程序中显示这张图以及可信度分数。这个函数将替换上一个代码片段中的注释：

self.firestore.collection("predicted_images").document(imageName!)
    .addSnapshotListener { documentSnapshot, error in
        if let error = error {
            print("error occurred\(error)")
        } else {
            if (documentSnapshot?.exists)! {
                let imageData = (documentSnapshot?.data())
                self.visualizePrediction(imgData: imageData)
            } else {
                print("waiting for prediction data...")
            }
        }
}

好了！现在我们有一款Taylor Swift检测器了！注意，由于模型只用了140张图像进行训练，所以准确度不够高，可能会把其他人误认为是Taylor。但是，如果有时间的话，我会收集更多贴有标签的图片，并更新模型，发布到应用商店里。

结语

这篇文章信息量有些大，涉及的外链较多。如果你想亲自动手试试，可以点击『阅读原文』查看分解步骤，其中包含重要的链接和资源。

原文地址：towardsdatascience.com/build-a-taylor-swift-detector-with-the-tensorflow-object-detection-api-ml-engine-and-swift-82707f5b4a56