Showing the changes to tools/python/convert.py (1 of 2 files). Show all
Unified · Split
| @@ -27,6 +27,7 @@ def get_arguments(): · 26 unmodified lines | |||
| 27 | parser.add_argument('--outdir', type=str, default='./', help='where to put generated files') | 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') | 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') | 29 | parser.add_argument('infile', help='path to the deep learning model with weights') |
| 30 | parser.add_argument('--dump4tb', type=str, default='no', help='dump file for visualization in tensorboard') | ||
| 30 | 31 | ||
| 31 | return parser.parse_args() | 32 | return parser.parse_args() |
| 32 | 33 | ||
| @@ -44,9 +45,12 @@ def main(): · 11 unmodified lines | |||
| 44 | basefile = os.path.split(args.infile)[1] | 45 | basefile = os.path.split(args.infile)[1] |
| 45 | basefile = os.path.splitext(basefile)[0] | 46 | basefile = os.path.splitext(basefile)[0] |
| 46 | outfile = os.path.join(args.outdir, basefile) + '.model' | 47 | outfile = os.path.join(args.outdir, basefile) + '.model' |
| 48 | dump4tb = False | ||
| 49 | if args.dump4tb.lower() in ('yes', 'true', 't', 'y', '1'): | ||
| 50 | dump4tb = True | ||
| 47 | 51 | ||
| 48 | if args.infmt == 'tensorflow': | 52 | if args.infmt == 'tensorflow': |
| 49 | convert_from_tensorflow(args.infile, outfile) | 53 | convert_from_tensorflow(args.infile, outfile, dump4tb) |
| 50 | 54 | ||
| 51 | if __name__ == '__main__': | 55 | if __name__ == '__main__': |
| 52 | main() | 56 | main() |