blob: ed6c81df9727b3475086d72b9284d5e2e22b1a3f (
plain)
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
|
# SPDX-FileCopyrightText: Copyright 2023, Arm Limited and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
import os
import flatbuffers
from tensorflow.lite.python import schema_py_generated as schema_fb
def load(input_tflite_file):
if not os.path.exists(input_tflite_file):
raise RuntimeError("TFLite file not found at %r\n" % input_tflite_file)
with open(input_tflite_file, "rb") as file_handle:
file_data = bytearray(file_handle.read())
model_obj = schema_fb.Model.GetRootAsModel(file_data, 0)
model = schema_fb.ModelT.InitFromObj(model_obj)
return model
def save(model, output_tflite_file):
builder = flatbuffers.Builder(1024) # Initial size of the buffer, which
# will grow automatically if needed
model_offset = model.Pack(builder)
builder.Finish(model_offset, file_identifier=b"TFL3")
model_data = builder.Output()
with open(output_tflite_file, "wb") as out_file:
out_file.write(model_data)
|
