diff --git a/tensorflow/lite/tools/flatbuffer_utils.py b/tensorflow/lite/tools/flatbuffer_utils.py index 25251ed925d..391132307df 100644 --- a/tensorflow/lite/tools/flatbuffer_utils.py +++ b/tensorflow/lite/tools/flatbuffer_utils.py @@ -158,7 +158,7 @@ def randomize_weights(model, random_seed=0): def xxd_output_to_bytes(input_cc_file): """Converts xxd output C++ source file to bytes (immutable) - + Args: input_cc_file: Full path name to th C++ source file dumped by xxd diff --git a/tensorflow/lite/tools/flatbuffer_utils_test.py b/tensorflow/lite/tools/flatbuffer_utils_test.py index 60235b06bc8..13a63c4f8b0 100644 --- a/tensorflow/lite/tools/flatbuffer_utils_test.py +++ b/tensorflow/lite/tools/flatbuffer_utils_test.py @@ -159,5 +159,34 @@ class RandomizeWeightsTest(test_util.TensorFlowTestCase): self.assertNotEqual(initial_buffer.data[j], final_buffer.data[j]) +class XxdOutputToBytesTest(test_util.TensorFlowTestCase): + + def testXxdOutputToBytes(self): + # 1. SETUP + # Define the initial model + initial_model = test_utils.build_mock_model() + initial_bytes = flatbuffer_utils.convert_object_to_bytearray(initial_model) + + # Define temporary files + tmp_dir = self.get_temp_dir() + model_filename = os.path.join(tmp_dir, 'model.tflite') + + # 2. INVOKE + # Invoke the write_model and read_model functions + flatbuffer_utils.write_model(initial_model, model_filename) + + # 3. DUMP WITH xxd + input_cc_file = os.path.join(tmp_dir, 'model.cc') + + command = "xxd -i {} > {}".format(model_filename, input_cc_file) + subprocess.call(command, shell=True) + + # 4. VALIDATE + final_bytes = flatbuffer_utils.xxd_output_to_bytes(input_cc_file) + + # Validate that the initial and final bytearray are the same + self.assertEqual(initial_bytes, final_bytes) + + if __name__ == '__main__': test.main()