30秒轻松实现TensorFlow物体检测

当前位置:

首页 > Python基础教程 >

30秒轻松实现TensorFlow物体检测

Google发布了新的TensorFlow物体检测API，包含了预训练模型，一个发布模型的jupyter notebook，一些可用于使用自己数据集对模型进行重新训练的有用脚本。

使用该API可以快速的构建一些图片中物体检测的应用。这里我们一步一步来看如何使用预训练模型来检测图像中的物体。

首先我们载入一些会使用的库

import numpy as np 

import os 

import six.moves.urllib as urllib 

import sys 

import tarfile 

import tensorflow as tf 

import zipfile 

from collections import defaultdict 

from io import StringIO 

from matplotlib import pyplot as plt 

from PIL import Image

接下来进行环境设置

1 2	`%matplotlib inline` `sys.path.append("..")`

物体检测载入

1 2	`from` `utils` `import` `label_map_util` `from` `utils` `import` `visualization_utils as vis_util`

准备模型

变量任何使用export_inference_graph.py工具输出的模型可以在这里载入，只需简单改变PATH_TO_CKPT指向一个新的.pb文件。这里我们使用“移动网SSD”模型。

MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017'

MODEL_FILE = MODEL_NAME + '.tar.gz'

DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt') 

NUM_CLASSES = 90

栏目列表