tomrunia / TF_FeatureExtraction

Convenient wrapper for TensorFlow feature extraction from pre-trained models using tf.contrib.slim
180 stars 61 forks source link
feature-extraction pre-trained slim tensorfow

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

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/
  1. 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"
  1. 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

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