一个基于 TensorFlow 的图片分类器

0x00 前言

备注:本文训练效果太差,所以直接太监了。

这年头,不会写爬虫不会写网站,那基本上不能算是一个 Python 程序员,但是 2017 年的 Google IO 之后,作为一个 Pythonist 你不会点数据分析和机器学习,也不好见人了。
所以,本文教你在什么机器学习概念都不懂的情况下,做出一个基于 TensorFlow 的图片分类功能,入个门。仅此而已。
本文的代码和文章内容主要源于我在 Github 上无意间翻到的一个 Repo, 链接地址,我所做的功夫就是在这基础上将代码改为了 Python 3 / TensorFlow 1.1.0 的环境,将这个流程梳理一下,不算是代码的生产者最多搬运工,仅此而已。

本文的目的是:

通过 TensorFlow 训练一个小型的分类器,用这个分类器通过照片识别出明星的姓名。

即我们要训练一个小 AI, 这个小 AI 能分辨欢乐颂里面的五美:

  • 乔欣
  • 刘涛
  • 王子文
  • 杨紫
  • 蒋欣

注意:本文主要在 MacOS 上进行测试。
本文的目录结构如下:

  • 准备训练数据和测试数据
  • TensorFlow 环境配置
  • 开始训练图片
  • 对图片进行分类
  • 文章回顾
  • 参考链接

首先克隆我的 Repo,(我的 Repo 在这里)[https://github.com/twocucao/the-machine] .

仓库结构大致如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
├── README.md
├── compose
│   └── tensorflow
│   ├── Dockerfile
│   └── Dockerfile-dev
├── config
├── doc
├── image_classifier
│   ├── __init__.py
│   ├── label_dir.py
│   ├── label_image.py
│   ├── retrain.py
│   └── train.sh
└── bootstrap.sh

0x01 准备训练数据和测试数据

准备训练数据,数据从哪里来呢?从百度来。我们从百度抓取大约 5 组图片,每组图片大约 1000 张图片,并且从每组里面留下 900 组作为训练数据,抽出 100 作为测试数据。

1.1 抓取图片

抓取代码放在代码仓库中,比较简单,下载即可用。

在 crawl_baidu_images.py 中填入五美的姓名,运行脚本即可。
脚本会请求百度图片的图片,然后下载下来,程序跑完会有如下的图片数据。
好,抓取图片我们就完成了。

1.2 归类训练数据

归类训练数据,其实就是把刚刚下载下来的图片,分类为五美,也就是把刘涛的照片放到刘涛文件夹中。看一下现在的文件夹,似乎已经分类完毕了,是不是这样呢?显然不是,因为:

  • 下载下来的图片貌似 JPG 结尾的图片,但是文件内容是不是 JPEG 的格式就不好说了,也可 GIF 也可能是 PNG.
  • 当你搜索刘涛的时候,百度图片给出的不仅仅是刘涛。也可能有胡歌《琅琊榜》, 也可能有胡军《天龙八部》

所以下面需要做的是:

  1. 移除非 JPG 的文件格式
  2. 人工确 (jian) 定 (huang) 文件夹中的图片。

在做上面这两步之前,我们先新建文件夹 image_classifier_train ( 笔者放在 /Users/twocucao/Codes/Repos/image_classifier_train ), 注意,这个文件夹不要放在代码下面,把五美的文件夹放到这个文件夹下面的 data 文件夹下。并且用拼音命名。

如下:

我们先移除非 JPEG 的图片。如果是 mac 系统需要先安装 jpeginfo , brew install jpeginfo 即可。
进入 /Users/twocucao/Codes/Repos/image_classifier_train 执行下面脚本 去除非 JPG 的图片

1
2
3
find . -iname "*.jpg" -exec jpeginfo -c {} \; | grep -E "WARNING|ERROR" > need_delete.sh
cat need_delete.sh | awk '{print $1}' | xargs rm
rm need_delete.sh

好了接下来,我们需要到每一个文件夹下进行人工鉴黄检验图片是不是五美,比如,到 liutao 文件夹下检查,删除图片基本上没有清晰面容的照片。
为了速度,我们把图片转成缩略图大致看一下,去除明显不是五美的照片,我们进行下一步的筛选。

1.3 找出对应的头像

我们需要从图片库中选出五美的头像来

1
2
3
4
pyenv global system
brew reinstall boost-python --with-python3 --without-python
conda install -c menpo dlib=19.4

0x02 TensorFlow 环境配置

2.1 Docker 的安装和镜像加速

本文需要使用 Docker 作为环境配置,也正是因为如此,我们可以在很快的时间内搭建起来 tensorflow 的运行环境。目测,Docker 也是未来几年内搭建环境分发环境的首选。
Docker 下载不必多说,需要补充一句的是,我们可以在阿里云账户上使用一个 registry 对 Docker 镜像进行加速。
在阿里云的容器界面获取加速链接填到 docker 里面即可。如图。

2.2 构建镜像文件并且构建镜像

建议在执行构建镜像之前,务必先完成本文的第二小节的图片准备。然后执行下面的命令,将镜像文件构建成镜像。

1
2
cd /Users/twocucao/Codes/Repos/the-machine
docker build -f compose/tensorflow/Dockerfile-dev -t twocucao/tensorflow .

该行命令使用 compose/tensorflow/Dockerfile-dev 作为 Dockerfile 文件,构建镜像名称为 twocucao/tensorflow , 传入的 context 为 当前路径。

2.3 测试 Tensorflow 容器

1
2
docker run -it twocucao/tensorflow /bin/bash
echo 'hello tensorflow'

如果运行正常,则一切正常。可以进行下一步骤了。

0x03 开始训练

执行命令开始训练。

1
./train.sh /Users/twocucao/Codes/Repos/image_classifier_train

我设置的训练次数为 20000, 在我的本子上基本两个小时,可能时间有些长,没有耐心的童鞋可以吧训练次数调整低一些。然后重新构建镜像。

那么,当 TensorFlow 在训练的时候,我们要谈些什么?

Google 开源了 Inception 模型,这个模型从 ImageNet 的上千个分类的图片训练而来,而我们所做的工作,便是在此基础上做最后的增量训练。然而,我们只用来区分女明星,似乎这个 Inception 的模型有点大材小用?好,训练结束之后我们查看一下文件夹 /Users/twocucao/Codes/Repos/image_classifier_train 下,

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
├── bottlenecks
│   ├── jiangxin
│   ├── liutao
│   ├── qiaoxin
│   ├── wangziwen
│   └── yangzi
├── data
│   ├── jiangxin
│   ├── liutao
│   ├── qiaoxin
│   ├── wangziwen
│   └── yangzi
├── inception
│   ├── LICENSE
│   ├── classify_image_graph_def.pb
│   ├── cropped_panda.jpg
│   ├── imagenet_2012_challenge_label_map_proto.pbtxt
│   ├── imagenet_synset_to_human_label_map.txt
│   └── inception-2015-12-05.tgz
├── retrained_graph.pb
├── retrained_labels.txt
└── test_data
├── src
└── target

Bottlenecks 文件夹为我们将每一张 JPG 转成矩阵的文本
Inception 为 Google 开源的模型文件。
retrained_graph.pb 为训练出来的图的模型文件。
retrained_labels.txt 为标签。

如下图不断刷出的一坨坨的文字是什么呢?

  • 时间 , 无需多说
  • 当前训练次数
  • Train accuracy = 87.0%
  • Cross entropy = 0.499145
  • Validation accuracy = 52.0% (N=100)

看到 Train accuracy >> Validation accuracy 估计是模型过度拟合了,嗯,看来这个模型还是有点问题的。

0x04 对图片进行分类

4.1 开始分类

4.2 对分类结果进行评估

0x05 文章回顾

0x06 参考链接