split drawnet into module code and script

Don't run scripts in the module dir to avoid import collisions between
io and caffe.io.
This commit is contained in:
Evan Shelhamer 2014-05-18 17:14:53 -07:00
Родитель 50d0b6d9c6
Коммит 6b85fd006d
2 изменённых файлов: 25 добавлений и 15 удалений

Просмотреть файл

@ -1,4 +1,3 @@
#!/usr/bin/env python
"""
Caffe network visualization: draw the NetParameter protobuffer.
@ -10,8 +9,6 @@ Caffe.
from caffe.proto import caffe_pb2
from google.protobuf import text_format
import pydot
import os
import sys
# Internal layer and blob styles.
LAYER_STYLE = {'shape': 'record', 'fillcolor': '#6495ED',
@ -77,15 +74,3 @@ def draw_net_to_file(caffe_net, filename):
ext = filename[filename.rfind('.')+1:]
with open(filename, 'wb') as fid:
fid.write(draw_net(caffe_net, ext))
if __name__ == '__main__':
if len(sys.argv) != 3:
print 'Usage: %s input_net_proto_file output_image_file' % \
os.path.basename(sys.argv[0])
else:
net = caffe_pb2.NetParameter()
text_format.Merge(open(sys.argv[1]).read(), net)
print 'Drawing net to %s' % sys.argv[2]
draw_net_to_file(net, sys.argv[2])

25
python/draw_net.py Executable file
Просмотреть файл

@ -0,0 +1,25 @@
#!/usr/bin/env python
"""
Draw a graph of the net architecture.
"""
import os
from google.protobuf import text_format
import caffe
from caffe.proto import caffe_pb2
def main(argv):
if len(argv) != 3:
print 'Usage: %s input_net_proto_file output_image_file' % \
os.path.basename(sys.argv[0])
else:
net = caffe_pb2.NetParameter()
text_format.Merge(open(sys.argv[1]).read(), net)
print 'Drawing net to %s' % sys.argv[2]
draw_net_to_file(net, sys.argv[2])
if __name__ == '__main__':
import sys
main(sys.argv)