def _create_metadata_file(self): associated_file1 = _metadata_fb.AssociatedFileT() associated_file1.name = b"file1" associated_file2 = _metadata_fb.AssociatedFileT() associated_file2.name = b"file2" self.expected_recorded_files = [ six.ensure_str(associated_file1.name), six.ensure_str(associated_file2.name) ] input_meta = _metadata_fb.TensorMetadataT() output_meta = _metadata_fb.TensorMetadataT() output_meta.associatedFiles = [associated_file2] subgraph = _metadata_fb.SubGraphMetadataT() # Create a model with two inputs and one output. subgraph.inputTensorMetadata = [input_meta, input_meta] subgraph.outputTensorMetadata = [output_meta] model_meta = _metadata_fb.ModelMetadataT() model_meta.name = "Mobilenet_quantized" model_meta.associatedFiles = [associated_file1] model_meta.subgraphMetadata = [subgraph] b = flatbuffers.Builder(0) b.Finish( model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) metadata_file = self.create_tempfile().full_path with open(metadata_file, "wb") as f: f.write(b.Output()) return metadata_file
def testGetRecordedAssociatedFileListWithSubgraphProcessUnits( self, tensor_type, tokenizer_type): # Creates a metadata with the tokenizer in the subgraph process units. tokenizer, expected_files = self._create_tokenizer(tokenizer_type) # Create the subgraph with process units. subgraph = _metadata_fb.SubGraphMetadataT() if tensor_type is TensorType.INPUT: subgraph.inputProcessUnits = [tokenizer] elif tensor_type is TensorType.OUTPUT: subgraph.outputProcessUnits = [tokenizer] else: raise ValueError( "The tensor type, {0}, is unsupported.".format(tensor_type)) # Creates the input and output tensor meta to match self._model_file. dummy_tensor_meta = _metadata_fb.TensorMetadataT() subgraph.inputTensorMetadata = [dummy_tensor_meta, dummy_tensor_meta] subgraph.outputTensorMetadata = [dummy_tensor_meta] # Create a model metadata with the subgraph metadata meta_buffer = self._create_model_meta_with_subgraph_meta(subgraph) # Creates the tempfiles. tempfiles = self._create_tempfiles(expected_files) # Creates the MetadataPopulator object. populator = _metadata.MetadataPopulator.with_model_file(self._model_file) populator.load_metadata_buffer(meta_buffer) populator.load_associated_files(tempfiles) populator.populate() recorded_files = populator.get_recorded_associated_file_list() self.assertEqual(set(recorded_files), set(expected_files))
def testPopulateMetadataFileToModelWithMetadataAndAssociatedFiles(self): # First, creates a dummy metadata different from self._metadata_file. It # needs to have the same input/output tensor numbers as self._model_file. # Populates it and the associated files into the model. input_meta = _metadata_fb.TensorMetadataT() output_meta = _metadata_fb.TensorMetadataT() subgraph = _metadata_fb.SubGraphMetadataT() # Create a model with two inputs and one output. subgraph.inputTensorMetadata = [input_meta, input_meta] subgraph.outputTensorMetadata = [output_meta] model_meta = _metadata_fb.ModelMetadataT() model_meta.subgraphMetadata = [subgraph] b = flatbuffers.Builder(0) b.Finish( model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) metadata_buf = b.Output() # Populate the metadata. populator1 = _metadata.MetadataPopulator.with_model_file(self._model_file) populator1.load_metadata_buffer(metadata_buf) populator1.load_associated_files([self._file1, self._file2]) populator1.populate() # Then, populate the metadata again. populator2 = _metadata.MetadataPopulator.with_model_file(self._model_file) populator2.load_metadata_file(self._metadata_file) populator2.populate() # Test if the metadata is populated correctly. self._assert_golden_metadata(self._model_file)
def testPopulatedFullPathAssociatedFileShouldSucceed(self): # Create AssociatedFileT using the full path file name. associated_file = _metadata_fb.AssociatedFileT() associated_file.name = self._file1 # Create model metadata with the associated file. subgraph = _metadata_fb.SubGraphMetadataT() subgraph.associatedFiles = [associated_file] # Creates the input and output tensor metadata to match self._model_file. dummy_tensor = _metadata_fb.TensorMetadataT() subgraph.inputTensorMetadata = [dummy_tensor, dummy_tensor] subgraph.outputTensorMetadata = [dummy_tensor] md_buffer = self._create_model_meta_with_subgraph_meta(subgraph) # Populate the metadata to a model. populator = _metadata.MetadataPopulator.with_model_file( self._model_file) populator.load_metadata_buffer(md_buffer) populator.load_associated_files([self._file1]) populator.populate() # The recorded file name in metadata should only contain file basename; file # directory should not be included. recorded_files = populator.get_recorded_associated_file_list() self.assertEqual(set(recorded_files), set([os.path.basename(self._file1)]))
def create_from_metadata( cls, model_buffer: bytearray, model_metadata: Optional[_metadata_fb.ModelMetadataT] = None, input_metadata: Optional[List[ _metadata_fb.TensorMetadataT]] = None, output_metadata: Optional[List[ _metadata_fb.TensorMetadataT]] = None, associated_files: Optional[List[str]] = None): """Creates MetadataWriter based on the metadata Flatbuffers Python Objects. Args: model_buffer: valid buffer of the model file. model_metadata: general model metadata [1]. The subgraph_metadata will be refreshed with input_metadata and output_metadata. input_metadata: a list of metadata of the input tensors [2]. output_metadata: a list of metadata of the output tensors [3]. associated_files: path to the associated files to be populated. [1]: https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L640-L681 [2]: https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L590 [3]: https://github.com/tensorflow/tflite-support/blob/b80289c4cd1224d0e1836c7654e82f070f9eefaa/tensorflow_lite_support/metadata/metadata_schema.fbs#L599 Returns: A MetadataWriter Object. """ # Create empty tensor metadata when input_metadata/output_metadata are None # to bypass MetadataPopulator verification. if not input_metadata: model = _schema_fb.Model.GetRootAsModel(model_buffer, 0) num_input_tensors = model.Subgraphs(0).InputsLength() input_metadata = [_metadata_fb.TensorMetadataT() ] * num_input_tensors if not output_metadata: model = _schema_fb.Model.GetRootAsModel(model_buffer, 0) num_output_tensors = model.Subgraphs(0).OutputsLength() output_metadata = [_metadata_fb.TensorMetadataT() ] * num_output_tensors subgraph_metadata = _metadata_fb.SubGraphMetadataT() subgraph_metadata.inputTensorMetadata = input_metadata subgraph_metadata.outputTensorMetadata = output_metadata if model_metadata is None: model_metadata = _metadata_fb.ModelMetadataT() model_metadata.subgraphMetadata = [subgraph_metadata] b = flatbuffers.Builder(0) b.Finish(model_metadata.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) return cls(model_buffer, b.Output(), associated_files)
def _create_dummy_model_metadata( tensor_metadata: _metadata_fb.TensorMetadataT) -> bytes: # Create a dummy model using the tensor metadata. subgraph_metadata = _metadata_fb.SubGraphMetadataT() subgraph_metadata.inputTensorMetadata = [tensor_metadata] model_metadata = _metadata_fb.ModelMetadataT() model_metadata.subgraphMetadata = [subgraph_metadata] # Create the Flatbuffers object and convert it to the json format. builder = flatbuffers.Builder(0) builder.Finish(model_metadata.Pack(builder), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) return bytes(builder.Output())
def testLoadMetadataBufferWithWrongOutputMetaNumberThrowsException(self): # Create a dummy metadata with no output tensor metadata, while the expected # number is 1. input_meta = _metadata_fb.TensorMetadataT() subgprah_meta = _metadata_fb.SubGraphMetadataT() subgprah_meta.inputTensorMetadata = [input_meta, input_meta] model_meta = _metadata_fb.ModelMetadataT() model_meta.subgraphMetadata = [subgprah_meta] builder = flatbuffers.Builder(0) builder.Finish( model_meta.Pack(builder), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) meta_buf = builder.Output() populator = _metadata.MetadataPopulator.with_model_buffer(self._model_buf) with self.assertRaises(ValueError) as error: populator.load_metadata_buffer(meta_buf) self.assertEqual( ("The number of output tensors (1) should match the number of " "output tensor metadata (0)"), str(error.exception))
def create_from_metadata_info( cls, model_buffer: bytearray, general_md: Optional[metadata_info.GeneralMd] = None, input_md: Optional[metadata_info.InputImageTensorMd] = None, output_location_md: Optional[metadata_info.TensorMd] = None, output_category_md: Optional[ metadata_info.CategoryTensorMd] = None, output_score_md: Union[ None, metadata_info.TensorMd, metadata_info.ClassificationTensorMd] = None, output_number_md: Optional[metadata_info.TensorMd] = None): """Creates MetadataWriter based on general/input/outputs information. Args: model_buffer: valid buffer of the model file. general_md: general infromation about the model. input_md: input image tensor informaton. output_location_md: output location tensor informaton. The location tensor is a multidimensional array of [N][4] floating point values between 0 and 1, the inner arrays representing bounding boxes in the form [top, left, bottom, right]. output_category_md: output category tensor information. The category tensor is an array of N integers (output as floating point values) each indicating the index of a class label from the labels file. output_score_md: output score tensor information. The score tensor is an array of N floating point values between 0 and 1 representing probability that a class was detected. Use ClassificationTensorMd to calibrate score. output_number_md: output number of detections tensor information. This tensor is an integer value of N. Returns: A MetadataWriter object. """ if general_md is None: general_md = metadata_info.GeneralMd( name=_MODEL_NAME, description=_MODEL_DESCRIPTION) if input_md is None: input_md = metadata_info.InputImageTensorMd( name=_INPUT_NAME, description=_INPUT_DESCRIPTION, color_space_type=_metadata_fb.ColorSpaceType.RGB) warn_message_format = ( "The output name isn't the default string \"%s\". This may cause the " "model not work in the TFLite Task Library since the tensor name will " "be used to handle the output order in the TFLite Task Library.") if output_location_md is None: output_location_md = metadata_info.TensorMd( name=_OUTPUT_LOCATION_NAME, description=_OUTPUT_LOCATION_DESCRIPTION) elif output_location_md.name != _OUTPUT_LOCATION_NAME: logging.warning(warn_message_format, _OUTPUT_LOCATION_NAME) if output_category_md is None: output_category_md = metadata_info.CategoryTensorMd( name=_OUTPUT_CATRGORY_NAME, description=_OUTPUT_CATEGORY_DESCRIPTION) elif output_category_md.name != _OUTPUT_CATRGORY_NAME: logging.warning(warn_message_format, _OUTPUT_CATRGORY_NAME) if output_score_md is None: output_score_md = metadata_info.ClassificationTensorMd( name=_OUTPUT_SCORE_NAME, description=_OUTPUT_SCORE_DESCRIPTION, ) elif output_score_md.name != _OUTPUT_SCORE_NAME: logging.warning(warn_message_format, _OUTPUT_SCORE_NAME) if output_number_md is None: output_number_md = metadata_info.TensorMd( name=_OUTPUT_NUMBER_NAME, description=_OUTPUT_NUMBER_DESCRIPTION) elif output_number_md.name != _OUTPUT_NUMBER_NAME: logging.warning(warn_message_format, _OUTPUT_NUMBER_NAME) # Create output tensor group info. group = _metadata_fb.TensorGroupT() group.name = _GROUP_NAME group.tensorNames = [ output_location_md.name, output_category_md.name, output_score_md.name ] # Gets the tensor inidces of tflite outputs and then gets the order of the # output metadata by the value of tensor indices. For instance, if the # output indices are [601, 599, 598, 600], tensor names and indices aligned # are: # - location: 598 # - category: 599 # - score: 600 # - number of detections: 601 # because of the op's ports of TFLITE_DETECTION_POST_PROCESS # (https://github.com/tensorflow/tensorflow/blob/a4fe268ea084e7d323133ed7b986e0ae259a2bc7/tensorflow/lite/kernels/detection_postprocess.cc#L47-L50). # Thus, the metadata of tensors are sorted in this way, according to # output_tensor_indicies correctly. output_tensor_indices = _get_tflite_outputs(model_buffer) metadata_list = [ _create_location_metadata(output_location_md), _create_metadata_with_value_range(output_category_md), _create_metadata_with_value_range(output_score_md), output_number_md.create_metadata() ] # Align indices with tensors. sorted_indices = sorted(output_tensor_indices) indices_to_tensors = dict(zip(sorted_indices, metadata_list)) # Output metadata according to output_tensor_indices. output_metadata = [ indices_to_tensors[i] for i in output_tensor_indices ] # Create subgraph info. subgraph_metadata = _metadata_fb.SubGraphMetadataT() subgraph_metadata.inputTensorMetadata = [input_md.create_metadata()] subgraph_metadata.outputTensorMetadata = output_metadata subgraph_metadata.outputTensorGroups = [group] # Create model metadata model_metadata = general_md.create_metadata() model_metadata.subgraphMetadata = [subgraph_metadata] b = flatbuffers.Builder(0) b.Finish(model_metadata.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) associated_files = [] _extend_new_files(associated_files, output_category_md.associated_files) _extend_new_files(associated_files, output_score_md.associated_files) return cls(model_buffer, b.Output(), associated_files=associated_files)
def create_from_metadata_info( cls, model_buffer: bytearray, general_md: Optional[metadata_info.GeneralMd] = None, input_md: Optional[metadata_info.InputImageTensorMd] = None, output_location_md: Optional[metadata_info.TensorMd] = None, output_category_md: Optional[metadata_info.CategoryTensorMd] = None, output_score_md: Optional[metadata_info.TensorMd] = None, output_number_md: Optional[metadata_info.TensorMd] = None): """Creates MetadataWriter based on general/input/outputs information. Args: model_buffer: valid buffer of the model file. general_md: general infromation about the model. input_md: input image tensor informaton. output_location_md: output location tensor informaton. The location tensor is a multidimensional array of [N][4] floating point values between 0 and 1, the inner arrays representing bounding boxes in the form [top, left, bottom, right]. output_category_md: output category tensor information. The category tensor is an array of N integers (output as floating point values) each indicating the index of a class label from the labels file. output_score_md: output score tensor information. The score tensor is an array of N floating point values between 0 and 1 representing probability that a class was detected. output_number_md: output number of dections tensor information. This tensor is an integer value of N. Returns: A MetadataWriter object. """ if general_md is None: general_md = metadata_info.GeneralMd( name=_MODEL_NAME, description=_MODEL_DESCRIPTION) if input_md is None: input_md = metadata_info.InputImageTensorMd( name=_INPUT_NAME, description=_INPUT_DESCRIPTION, color_space_type=_metadata_fb.ColorSpaceType.RGB) if output_location_md is None: output_location_md = metadata_info.TensorMd( name=_OUTPUT_LOCATION_NAME, description=_OUTPUT_LOCATION_DESCRIPTION) if output_category_md is None: output_category_md = metadata_info.CategoryTensorMd( name=_OUTPUT_CATRGORY_NAME, description=_OUTPUT_CATEGORY_DESCRIPTION) if output_score_md is None: output_score_md = metadata_info.TensorMd( name=_OUTPUT_SCORE_NAME, description=_OUTPUT_SCORE_DESCRIPTION) if output_number_md is None: output_number_md = metadata_info.TensorMd( name=_OUTPUT_NUMBER_NAME, description=_OUTPUT_NUMBER_DESCRIPTION) if output_category_md.associated_files is None: output_category_md.associated_files = [] # Create output tensor group info. group = _metadata_fb.TensorGroupT() group.name = _GROUP_NAME group.tensorNames = [ output_location_md.name, output_category_md.name, output_score_md.name ] # Create subgraph info. subgraph_metadata = _metadata_fb.SubGraphMetadataT() subgraph_metadata.inputTensorMetadata = [input_md.create_metadata()] subgraph_metadata.outputTensorMetadata = [ _create_location_metadata(output_location_md), _create_metadata_with_value_range(output_category_md), _create_metadata_with_value_range(output_score_md), output_number_md.create_metadata() ] subgraph_metadata.outputTensorGroups = [group] # Create model metadata model_metadata = general_md.create_metadata() model_metadata.subgraphMetadata = [subgraph_metadata] b = flatbuffers.Builder(0) b.Finish( model_metadata.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) return cls( model_buffer, b.Output(), associated_files=[ file.file_path for file in output_category_md.associated_files ])