def _populate_metadata_buffer(self): """Populates the metadata buffer (in bytearray) into the model file. Inserts metadata_buf into the metadata field of schema.Model. If the MetadataPopulator object is created using the method, with_model_file(model_file), the model file will be updated. Existing metadata buffer (if applied) will be overridden by the new metadata buffer. """ with open(self._model_file, "rb") as f: model_buf = f.read() model = _schema_fb.ModelT.InitFromObj( _schema_fb.Model.GetRootAsModel(model_buf, 0)) buffer_field = _schema_fb.BufferT() buffer_field.data = self._metadata_buf is_populated = False if not model.metadata: model.metadata = [] else: # Check if metadata has already been populated. for meta in model.metadata: if meta.name.decode("utf-8") == self.METADATA_FIELD_NAME: is_populated = True model.buffers[meta.buffer] = buffer_field if not is_populated: if not model.buffers: model.buffers = [] model.buffers.append(buffer_field) # Creates a new metadata field. metadata_field = _schema_fb.MetadataT() metadata_field.name = self.METADATA_FIELD_NAME metadata_field.buffer = len(model.buffers) - 1 model.metadata.append(metadata_field) # Packs model back to a flatbuffer binaray file. b = flatbuffers.Builder(0) b.Finish(model.Pack(b), self.TFLITE_FILE_IDENTIFIER) model_buf = b.Output() # Saves the updated model buffer to model file. # Gets files that have been packed to self._model_file. packed_files = self.get_packed_associated_file_list() if packed_files: # Writes the updated model buffer and associated files into a new model # file. Then overwrites the original model file. with tempfile.NamedTemporaryFile() as temp: new_file = temp.name with open(new_file, "wb") as f: f.write(model_buf) self._copy_archived_files(self._model_file, new_file, packed_files) shutil.copy(new_file, self._model_file) os.remove(new_file) else: with open(self._model_file, "wb") as f: f.write(model_buf)
def _create_model_file_with_metadata_and_buf_fields(self): metadata_field = _schema_fb.MetadataT() metadata_field.name = "meta" buffer_field = _schema_fb.BufferT() model = _schema_fb.ModelT() model.metadata = [metadata_field, metadata_field] model.buffers = [buffer_field, buffer_field, buffer_field] model_builder = flatbuffers.Builder(0) model_builder.Finish( model.Pack(model_builder), _metadata.MetadataPopulator.TFLITE_FILE_IDENTIFIER) mnodel_file = self.create_tempfile().full_path with open(mnodel_file, "wb") as f: f.write(model_builder.Output()) return mnodel_file
def _populate_metadata_with_identifier(self, model_buf, metadata_buf, identifier): # For testing purposes only. MetadataPopulator cannot populate metadata with # wrong identifiers. model = _schema_fb.ModelT.InitFromObj( _schema_fb.Model.GetRootAsModel(model_buf, 0)) buffer_field = _schema_fb.BufferT() buffer_field.data = metadata_buf model.buffers = [buffer_field] # Creates a new metadata field. metadata_field = _schema_fb.MetadataT() metadata_field.name = _metadata.MetadataPopulator.METADATA_FIELD_NAME metadata_field.buffer = len(model.buffers) - 1 model.metadata = [metadata_field] b = flatbuffers.Builder(0) b.Finish(model.Pack(b), identifier) return b.Output()