A lightweight python module to load TensorFlow frozen model (a single pb file).

These instructions will get you a copy of the project up and running on your local machine for development and testing purposes.

TensorFlow

# For CPU
python3 -m pip install tensorflow

# or, if you have a CUDA compatible GPU
python3 -m pip install tensorflow-gpu

OpenCV, only if you want to run the demo with video/camera.

python3 -m pip install opencv-python

From your favorite development directory:

git clone https://github.com/yinguobing/butterfly.git

A demonstration of how to use this module through the demo.py script. Before we start make sure you have your frozen model named frozen_graph.pb ready.

Most of the time you should know the input and output nodes' name. In case you don't, we can output all the node names in the graph.

python3 demo.py --model frozen_model.pb --list_ops true

and the output looks like:

butterfly/image_tensor
butterfly/map/Shape
butterfly/map/strided_slice/stack
...
butterfly/model/final_dense
butterfly/ArgMax/dimension
butterfly/ArgMax
butterfly/softmax_tensor

Here the input node is butterfly/image_tensor and the output node is butterfly/ArgMax. We will use these later.

NOTE: If the model has more than one output node, just separate them with ','.

Use the --image argument.

python3 demo.py \
    --model frozen_model.pb \
    --input_node butterfly/image_tensor \
    --output_node butterfly/ArgMax \
    --image image.jpg

For video file, use the --video argument.

python3 demo.py \
    --model frozen_model.pb \
    --input_node butterfly/image_tensor \
    --output_node butterfly/ArgMax \
    --video video.mp4

For webcam, use the --cam argument.

python3 demo.py \
    --model frozen_model.pb \
    --input_node butterfly/image_tensor \
    --output_node butterfly/ArgMax \
    --cam 0

butterfly.py is sufficient for your development.

• yinguobing (尹国冰) - yinguobing

This project is licensed under the MIT License - see the LICENSE.md file for details