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