/
caffe2keras_converter.py
106 lines (87 loc) · 3.96 KB
/
caffe2keras_converter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
""" Converts a caffe model into a Keras equivalent one """
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import models
from tensorflow.keras import layers
import os.path
from pathlib import Path
import argparse
import sys
sys.path.insert(1, 'caffe2keras')
sys.path.insert(1, 'caffe2keras/caffe_weight_converter')
import create_nn_struct
def write_beginning(outfile):
"""
Writes the initial part of the source file to generate the Keras model (imports etc)
Arguments:
outfile: a file object to write the code in
"""
import_modules = {
'tensorflow': ['keras'],
'tensorflow.keras': ['models'],
'tensorflow.keras.layers': ['*']
}
outfile.write('import tensorflow as tf\n')
for k in import_modules:
for module in import_modules[k]:
outfile.write('from {} import {}\n'.format(k, module))
outfile.write("\n\ndef keras_model():\n")
def write_src(prototxt_path, output_dir):
"""
Writes the Python source file defining the Keras model
Arguments:
prototxt_path: the path to the model's prototxt file
outsrc_path: the path to the output Python source file defining the Keras model
"""
# Create the output directory (if it doens't exist)
Path(output_dir).mkdir(parents=True, exist_ok=True)
outsrc_path = output_dir + 'net_caffe2keras.py'
with open(outsrc_path, 'w') as outfile:
write_beginning(outfile)
create_nn_struct.write_nn_struct_code_keras(prototxt_path, outsrc_path)
with open(outsrc_path, 'a') as outfile:
outfile.write("\n\n\treturn keras_model\n\n\n")
outfile.write(f"if __name__ == '__main__':\n\tkeras_model()")
def caffe2keras(prototxt_path, caffemodel_path, output_dir, verbose):
"""
Converts a caffe model into a Keras qìequivalent one
Arguments:
prototxt_path: the path to the model's prototxt file
outsrc_path: the path to the output Python source file defining the Keras model
"""
output_dir = output_dir + ('/' if output_dir[-1] != '/' else '')
# Write the Python source file containing the keras model
write_src(prototxt_path, output_dir)
# Import the source file we just created
import sys
sys.path.insert(1, output_dir)
import net_caffe2keras
# Take the model from the source file
keras_model = net_caffe2keras.keras_model()
if verbose: keras_model.summary()
# Export weights from the Caffe net, load them into the Keras model and save the model
args = argparse.Namespace()
args.out_file = output_dir + 'weights'
args.prototxt = prototxt_path
args.caffemodel = caffemodel_path
args.format = 'hdf5'
args.include_non_weight = True
args.skip_unknown = True
args.backend = 'tf'
if verbose: args.verbose = True
else: args.verbose = False
import caffe_weight_converter
caffe_weight_converter.convert_caffe_weights(args)
keras_model.load_weights(output_dir + 'weights.h5')
keras_model.save(output_dir + 'keras_model.h5')
if __name__ == '__main__':
parser = argparse.ArgumentParser(
description="Reads the Caffe network's defnition from its prototxt, the parameters from its caffemodel \
and generates a Python file containing the network (architecture + parameters) in Keras"
)
parser.add_argument('prototxt', action='store', help="The filename (full path including file extension) of the '.prototxt' file that defines the Caffe model.")
parser.add_argument('caffemodel', action='store', help="The filename (full path including file extension) of the '.caffemodel' file that contains the network's parameters")
parser.add_argument('output_dir', action='store', help="The path to the directory where to save the Keras model and the file where you want the code to be written in.")
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()
caffe2keras(args.prototxt, args.caffemodel,args.output_dir, args.verbose)