zengsn/TF_FeatureExtraction

Convenient wrapper for TensorFlow feature extraction from pre-trained models using tf.contrib.slim

TensorFlow Feature Extractor

This is a convenient wrapper for feature extraction or classification in TensorFlow. Given well known pre-trained models on ImageNet, the extractor runs over a list or directory of images. Optionally, features can be saved as HDF5 file. It supports all the pre-trained models listed on the official page.

TensorFlow models tested:

1. Inception v1-v4
2. ResNet v1 and v2
3. VGG 16-19

Requirements

• TensorFlow (tested with version 1.3)
• TensorFlow Models
• The usual suspects: numpy, scipy.
• Optionally h5py for saving features to HDF5 file

Setup

1. Checkout the TensorFlow models repository somewhere on your machine. The path where you checkout the repository will be denoted <checkout_dir>/models

git clone https://github.com/tensorflow/models/

2. Add the directory <checkout_dir>/research/slim to the$PYTHONPATH variable. Or add a line to your .bashrc file.

export PYTHONPATH="<checkout_dir>/research/slim:$PYTHONPATH"

3. Download the model checkpoints from the official page.

Usage

There are two example files, one for classification and one for feature extraction.

Feature Extraction

ResNet-v1-101

example_feat_extract.py 
--network resnet_v1_101 
--checkpoint ./checkpoints/resnet_v1_101.ckpt 
--image_path ./images_dir/ 
--out_file ./features.h5
--num_classes 1000 
--layer_names resnet_v1_101/logits

ResNet-v2-101

example_feat_extract.py 
--network resnet_v2_101 
--checkpoint ./checkpoints/resnet_v2_101.ckpt 
--image_path ./images_dir/
--out_file ./features.h5 
--layer_names resnet_v2_101/logits 
--preproc_func inception

Inception-v4

example_feat_extract.py 
--network inception_v4 
--checkpoint ./checkpoints/inception_v4.ckpt 
--image_path ./images_dir/
--out_file ./features.h5 
--layer_names Logits

Inception-ResNet-v2

example_feat_extract.py \
--network inception_resnet_v2  \
--checkpoint ../../checkpoints/inception_resnet_v2_2016_08_30.ckpt \
--image_path ../../data/FashionMNIST_28x28/ \
--out_file ../../data/FashionMNIST_28x28.inception_resnet_v2.h5 \
--num_classes 1001 \
--preproc_func inception \
--layer_names Logits \   

VGG-16

example_feat_extract.py \
--network vgg_16  \
--checkpoint ../../checkpoints/vgg_16.ckpt \
--image_path ../../data/FashionMNIST_28x28/ \
--out_file ../../data/FashionMNIST_28x28.vgg_16.h5 \
--num_classes 1000 \
--layer_names vgg_16/fc6 \  

Feature Extraction class by class for large-scale datasets

ResNet-v1-101

example_feat_extract.py 
--network resnet_v1_101 
--checkpoint ./checkpoints/resnet_v1_101.ckpt 
--image_path ./imagenet/ 
--out_file ./imagenet.resnet_v1_101.global_pool.h5
--num_classes 1000 
--layer_names global_pool
--class_by_class yes

Image Classification

example_classification.py
--network resnet_v1_101 
--checkpoint ./checkpoints/resnet_v1_101.ckpt 
--image_path ./images_dir/
--num_classes 1000 
--logits_name resnet_v1_101/logits

Work in Progress

1. Save image file names to HDF5 file
2. Support for multi-threaded preprocessing