Python3 version of pb2t

[robinvanemden]

For a side project, we needed a basic Python3 version of pb2t. I thought I'd share our code here - maybe it is of use to others as well. Tested & working for ONNC v1.3.0.

import sys
import base64
import onnx   # install: https://pypi.org/project/onnx/

if len(sys.argv) <= 2:
    exit("usage: " + sys.argv[0] + " input.pb output.tensor")

input_file = sys.argv[1]
output_file = sys.argv[2]

try:
    reader = onnx.TensorProto()
    with open(input_file, 'rb') as f:
        reader.ParseFromString(f.read())
    f.close()
except IOError:
    print("Could not open file: " + sys.argv[1])

assert (reader and hasattr(reader, 'data_type') and reader.data_type == onnx.TensorProto.FLOAT)

input_data = reader.raw_data
input_data_length = len(input_data)

filemagic_byte_array = base64.b16decode(b'2E5453520000000001000000000000002000000000000000')
length_byte_array = input_data_length.to_bytes(8, 'little')

header_byte_array = filemagic_byte_array + length_byte_array

tsr_file = header_byte_array + input_data
try:
    with open(output_file, 'wb') as f:
        f.write(tsr_file)
        f.close()
except IOError:
    print("Could not write file: " + sys.argv[2])