For example, given TensorFlow model file espcn.pb, to generate native model file espcn.model, just run: python convert.py espcn.pb In current implementation, the native model file is generated for specific dnn network with hard-code python scripts maintained out of ffmpeg. For example, srcnn network used by vf_sr is generated with https://github.com/HighVoltageRocknRoll/sr/blob/master/generate_header_and_model.py#L85 In this patch, the script is designed as a general solution which converts general TensorFlow model .pb file into .model file. The script now has some tricky to be compatible with current implemention, will be refined step by step. The script is also added into ffmpeg source tree. It is expected there will be many more patches and community needs the ownership of it. Another technical direction is to do the conversion in c/c++ code within ffmpeg source tree. While .pb file is organized with protocol buffers, it is not easy to do such work with tiny c/c++ code, see more discussion at http://ffmpeg.org/pipermail/ffmpeg-devel/2019-May/244496.html. So, choose the python script.
Showing the changes to tools/python/convert.py (1 of 3 files). Show all
Unified · Split
| @@ -0,0 +1,52 @@ | ||
| 1 | +# Copyright (c) 2019 Guo Yejun | |
| 2 | +# | |
| 3 | +# This file is part of FFmpeg. | |
| 4 | +# | |
| 5 | +# FFmpeg is free software; you can redistribute it and/or | |
| 6 | +# modify it under the terms of the GNU Lesser General Public | |
| 7 | +# License as published by the Free Software Foundation; either | |
| 8 | +# version 2.1 of the License, or (at your option) any later version. | |
| 9 | +# | |
| 10 | +# FFmpeg is distributed in the hope that it will be useful, | |
| 11 | +# but WITHOUT ANY WARRANTY; without even the implied warranty of | |
| 12 | +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU | |
| 13 | +# Lesser General Public License for more details. | |
| 14 | +# | |
| 15 | +# You should have received a copy of the GNU Lesser General Public | |
| 16 | +# License along with FFmpeg; if not, write to the Free Software | |
| 17 | +# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA | |
| 18 | +# ============================================================================== | |
| 19 | + | |
| 20 | +# verified with Python 3.5.2 on Ubuntu 16.04 | |
| 21 | +import argparse | |
| 22 | +import os | |
| 23 | +from convert_from_tensorflow import * | |
| 24 | + | |
| 25 | +def get_arguments(): | |
| 26 | + parser = argparse.ArgumentParser(description='generate native mode model with weights from deep learning model') | |
| 27 | + parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') | |
| 28 | + parser.add_argument('--infmt', type=str, default='tensorflow', help='format of the deep learning model') | |
| 29 | + parser.add_argument('infile', help='path to the deep learning model with weights') | |
| 30 | + | |
| 31 | + return parser.parse_args() | |
| 32 | + | |
| 33 | +def main(): | |
| 34 | + args = get_arguments() | |
| 35 | + | |
| 36 | + if not os.path.isfile(args.infile): | |
| 37 | + print('the specified input file %s does not exist' % args.infile) | |
| 38 | + exit(1) | |
| 39 | + | |
| 40 | + if not os.path.exists(args.outdir): | |
| 41 | + print('create output directory %s' % args.outdir) | |
| 42 | + os.mkdir(args.outdir) | |
| 43 | + | |
| 44 | + basefile = os.path.split(args.infile)[1] | |
| 45 | + basefile = os.path.splitext(basefile)[0] | |
| 46 | + outfile = os.path.join(args.outdir, basefile) + '.model' | |
| 47 | + | |
| 48 | + if args.infmt == 'tensorflow': | |
| 49 | + convert_from_tensorflow(args.infile, outfile) | |
| 50 | + | |
| 51 | +if __name__ == '__main__': | |
| 52 | + main() | |