Arafatk/tensorflow.rb

# This file is useful for reading the contents of the ops generated by ruby.
# You can read any graph defination in pb/pbtxt format generated by ruby
# or by python and then convert it back and forth from human readable to
# binary format.

import tensorflow as tf
import shutil
import os
from google.protobuf import text_format
from tensorflow.python.platform import gfile

# Converts protobuf file from pbtxt to binary wire format


def pbtxt_to_graphdef(filename):
    with open(filename, 'r') as f:
        graph_def = tf.GraphDef()
        file_content = f.read()
        text_format.Merge(file_content, graph_def)
        tf.import_graph_def(graph_def, name='')
        tf.train.write_graph(graph_def, 'pbtxt/', 'protobuf.pb', as_text=False)

# Converts protobuf file from binary wire format to pbtxt


def graphdef_to_pbtxt(filename):
    with gfile.FastGFile(filename, 'rb') as f:
        graph_defe = tf.GraphDef()
        graph_defe.ParseFromString(f.read())
        tf.import_graph_def(graph_defe, name='')
        tf.train.write_graph(graph_defe, 'pbtxt/',
                             'protobuf.pbtxt', as_text=True)
    return


# A simple check to remove the directory if its already present.
directory = "/home/arafat/Desktop/tensorflow/tensorflow.rb/ext/sciruby/tensorflow_c/pbtxt"
if (os.path.isdir(directory)):
    shutil.rmtree(directory)

# Name of the file to be converted
filename = 'image_int64.pb'
graphdef_to_pbtxt(filename)