| /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. |
| |
| Licensed under the Apache License, Version 2.0 (the "License"); |
| you may not use this file except in compliance with the License. |
| You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| ==============================================================================*/ |
| |
| // This file contains the logic of android model wrapper generation. |
| // |
| // At the beginning is the helper functions handling metadata and code writer. |
| // |
| // Codes are generated in every `Generate{FOO}` functions. Gradle and Manifest |
| // files are simple. The wrapper file generation is a bit complex so we divided |
| // it into several sub-functions. |
| // |
| // The structure of the wrapper file looks like: |
| // |
| // [ imports ] |
| // [ class ] |
| // [ inner "Outputs" class ] |
| // [ innner "Metadata" class ] |
| // [ APIs ] ( including ctors, public APIs and private APIs ) |
| // |
| // We tried to mostly write it in a "template-generation" way. `CodeWriter` does |
| // the job as a template renderer. To avoid repeatedly setting the token values, |
| // helper functions `SetCodeWriterWith{Foo}Info` set the token values with info |
| // structures (`TensorInfo` and `ModelInfo`) - the Info structures are |
| // intermediate datastructures between Metadata (represented in Flatbuffers) and |
| // generated code. |
| |
| #include "tensorflow_lite_support/codegen/android_java_generator.h" |
| |
| #include <ctype.h> |
| |
| #include <algorithm> |
| #include <memory> |
| #include <string> |
| #include <vector> |
| |
| #include "tensorflow_lite_support/codegen/code_generator.h" |
| #include "tensorflow_lite_support/codegen/metadata_helper.h" |
| #include "tensorflow_lite_support/codegen/utils.h" |
| #include "tensorflow_lite_support/metadata/metadata_schema_generated.h" |
| |
| namespace tflite { |
| namespace support { |
| namespace codegen { |
| |
| namespace { |
| |
| using details_android_java::ModelInfo; |
| using details_android_java::TensorInfo; |
| |
| // Helper class to organize the C++ code block as a generated code block. |
| // Using ctor and dtor to simulate an enter/exit schema like `with` in Python. |
| class AsBlock { |
| public: |
| AsBlock(CodeWriter* code_writer, const std::string& before, |
| bool trailing_blank_line = false) |
| : code_writer_(code_writer), trailing_blank_line_(trailing_blank_line) { |
| code_writer_->AppendNoNewLine(before); |
| code_writer_->Append(" {"); |
| code_writer_->Indent(); |
| } |
| ~AsBlock() { |
| code_writer_->Outdent(); |
| code_writer_->Append("}"); |
| if (trailing_blank_line_) { |
| code_writer_->NewLine(); |
| } |
| } |
| |
| private: |
| CodeWriter* code_writer_; |
| bool trailing_blank_line_; |
| }; |
| |
| // Declare the functions first, so that the functions can follow a logical |
| // order. |
| bool GenerateWrapperClass(CodeWriter*, const ModelInfo&, ErrorReporter*); |
| bool GenerateWrapperImports(CodeWriter*, const ModelInfo&, ErrorReporter*); |
| bool GenerateWrapperInputs(CodeWriter*, const ModelInfo&, ErrorReporter*); |
| bool GenerateWrapperOutputs(CodeWriter*, const ModelInfo&, ErrorReporter*); |
| bool GenerateWrapperMetadata(CodeWriter*, const ModelInfo&, ErrorReporter*); |
| bool GenerateWrapperAPI(CodeWriter*, const ModelInfo&, ErrorReporter*); |
| |
| std::string GetModelVersionedName(const ModelMetadata* metadata) { |
| std::string model_name = "MyModel"; |
| if (metadata->name() != nullptr && !(metadata->name()->str().empty())) { |
| model_name = metadata->name()->str(); |
| } |
| std::string model_version = "unknown"; |
| if (metadata->version() != nullptr && !(metadata->version()->str().empty())) { |
| model_version = metadata->version()->str(); |
| } |
| return model_name + " (Version: " + model_version + ")"; |
| } |
| |
| TensorInfo CreateTensorInfo(const TensorMetadata* metadata, |
| const std::string& name, bool is_input, int index, |
| ErrorReporter* err) { |
| TensorInfo tensor_info; |
| std::string tensor_identifier = is_input ? "input" : "output"; |
| tensor_identifier += " " + std::to_string(index); |
| tensor_info.associated_axis_label_index = FindAssociatedFile( |
| metadata, AssociatedFileType_TENSOR_AXIS_LABELS, tensor_identifier, err); |
| tensor_info.associated_value_label_index = FindAssociatedFile( |
| metadata, AssociatedFileType_TENSOR_VALUE_LABELS, tensor_identifier, err); |
| if (is_input && (tensor_info.associated_axis_label_index >= 0 || |
| tensor_info.associated_value_label_index >= 0)) { |
| err->Warning( |
| "Found label file on input tensor (%s). Label file for input " |
| "tensor is not supported yet. The " |
| "file will be ignored.", |
| tensor_identifier.c_str()); |
| } |
| if (tensor_info.associated_axis_label_index >= 0 && |
| tensor_info.associated_value_label_index >= 0) { |
| err->Warning( |
| "Found both axis label file and value label file for tensor (%s), " |
| "which is not supported. Only the axis label file will be used.", |
| tensor_identifier.c_str()); |
| } |
| tensor_info.is_input = is_input; |
| tensor_info.name = SnakeCaseToCamelCase(name); |
| tensor_info.upper_camel_name = tensor_info.name; |
| tensor_info.upper_camel_name[0] = toupper(tensor_info.upper_camel_name[0]); |
| tensor_info.normalization_unit = |
| FindNormalizationUnit(metadata, tensor_identifier, err); |
| if (metadata->content() != nullptr && |
| metadata->content()->content_properties() != nullptr) { |
| // Enter tensor wrapper type inferring |
| if (metadata->content()->content_properties_type() == |
| ContentProperties_ImageProperties) { |
| if (metadata->content() |
| ->content_properties_as_ImageProperties() |
| ->color_space() == ColorSpaceType_RGB) { |
| tensor_info.content_type = "image"; |
| tensor_info.wrapper_type = "TensorImage"; |
| tensor_info.processor_type = "ImageProcessor"; |
| return tensor_info; |
| } else { |
| err->Warning( |
| "Found Non-RGB image on tensor (%s). Codegen currently does not " |
| "support it, and regard it as a plain numeric tensor.", |
| tensor_identifier.c_str()); |
| } |
| } |
| } |
| tensor_info.content_type = "tensor"; |
| tensor_info.wrapper_type = "TensorBuffer"; |
| tensor_info.processor_type = "TensorProcessor"; |
| return tensor_info; |
| } |
| |
| ModelInfo CreateModelInfo(const ModelMetadata* metadata, |
| const std::string& package_name, |
| const std::string& model_class_name, |
| const std::string& model_asset_path, |
| ErrorReporter* err) { |
| ModelInfo model_info; |
| if (!CodeGenerator::VerifyMetadata(metadata, err)) { |
| // TODO(b/150116380): Create dummy model info. |
| err->Error("Validating metadata failed."); |
| return model_info; |
| } |
| model_info.package_name = package_name; |
| model_info.model_class_name = model_class_name; |
| model_info.model_asset_path = model_asset_path; |
| model_info.model_versioned_name = GetModelVersionedName(metadata); |
| const auto* graph = metadata->subgraph_metadata()->Get(0); |
| auto names = CodeGenerator::NameInputsAndOutputs( |
| graph->input_tensor_metadata(), graph->output_tensor_metadata()); |
| std::vector<std::string> input_tensor_names = std::move(names.first); |
| std::vector<std::string> output_tensor_names = std::move(names.second); |
| |
| for (int i = 0; i < input_tensor_names.size(); i++) { |
| model_info.inputs.push_back( |
| CreateTensorInfo(graph->input_tensor_metadata()->Get(i), |
| input_tensor_names[i], true, i, err)); |
| if (i < input_tensor_names.size() - 1) { |
| model_info.inputs_list += ", "; |
| model_info.input_type_param_list += ", "; |
| } |
| model_info.inputs_list += model_info.inputs[i].name; |
| model_info.input_type_param_list += |
| model_info.inputs[i].wrapper_type + " " + model_info.inputs[i].name; |
| } |
| for (int i = 0; i < output_tensor_names.size(); i++) { |
| model_info.outputs.push_back( |
| CreateTensorInfo(graph->output_tensor_metadata()->Get(i), |
| output_tensor_names[i], false, i, err)); |
| if (i < output_tensor_names.size() - 1) { |
| model_info.postprocessor_type_param_list += ", "; |
| model_info.postprocessors_list += ", "; |
| } |
| model_info.postprocessors_list += |
| model_info.outputs[i].name + "Postprocessor"; |
| model_info.postprocessor_type_param_list += |
| model_info.outputs[i].processor_type + " " + |
| model_info.outputs[i].name + "Postprocessor"; |
| } |
| return model_info; |
| } |
| |
| void SetCodeWriterWithTensorInfo(CodeWriter* code_writer, |
| const TensorInfo& tensor_info) { |
| code_writer->SetTokenValue("NAME", tensor_info.name); |
| code_writer->SetTokenValue("NAME_U", tensor_info.upper_camel_name); |
| code_writer->SetTokenValue("CONTENT_TYPE", tensor_info.content_type); |
| code_writer->SetTokenValue("WRAPPER_TYPE", tensor_info.wrapper_type); |
| std::string wrapper_name = tensor_info.wrapper_type; |
| wrapper_name[0] = tolower(wrapper_name[0]); |
| code_writer->SetTokenValue("WRAPPER_NAME", wrapper_name); |
| code_writer->SetTokenValue("PROCESSOR_TYPE", tensor_info.processor_type); |
| code_writer->SetTokenValue("NORMALIZATION_UNIT", |
| std::to_string(tensor_info.normalization_unit)); |
| code_writer->SetTokenValue( |
| "ASSOCIATED_AXIS_LABEL_INDEX", |
| std::to_string(tensor_info.associated_axis_label_index)); |
| code_writer->SetTokenValue( |
| "ASSOCIATED_VALUE_LABEL_INDEX", |
| std::to_string(tensor_info.associated_value_label_index)); |
| } |
| |
| void SetCodeWriterWithModelInfo(CodeWriter* code_writer, |
| const ModelInfo& model_info) { |
| code_writer->SetTokenValue("PACKAGE", model_info.package_name); |
| code_writer->SetTokenValue("MODEL_PATH", model_info.model_asset_path); |
| code_writer->SetTokenValue("MODEL_CLASS_NAME", model_info.model_class_name); |
| // Extra info, half generated. |
| code_writer->SetTokenValue("INPUT_TYPE_PARAM_LIST", |
| model_info.input_type_param_list); |
| code_writer->SetTokenValue("INPUTS_LIST", model_info.inputs_list); |
| code_writer->SetTokenValue("POSTPROCESSORS_LIST", |
| model_info.postprocessors_list); |
| code_writer->SetTokenValue("POSTPROCESSOR_TYPE_PARAM_LIST", |
| model_info.postprocessor_type_param_list); |
| } |
| |
| constexpr char JAVA_DEFAULT_PACKAGE[] = "default"; |
| |
| std::string ConvertPackageToPath(const std::string& package) { |
| if (package == JAVA_DEFAULT_PACKAGE) { |
| return ""; |
| } |
| std::string path = package; |
| std::replace(path.begin(), path.end(), '.', '/'); |
| return path; |
| } |
| |
| bool IsImageUsed(const ModelInfo& model) { |
| for (const auto& input : model.inputs) { |
| if (input.content_type == "image") { |
| return true; |
| } |
| } |
| for (const auto& output : model.outputs) { |
| if (output.content_type == "image") { |
| return true; |
| } |
| } |
| return false; |
| } |
| |
| // The following functions generates the wrapper Java code for a model. |
| |
| bool GenerateWrapperFileContent(CodeWriter* code_writer, const ModelInfo& model, |
| ErrorReporter* err) { |
| code_writer->Append("// Generated by TFLite Support."); |
| code_writer->Append("package {{PACKAGE}};"); |
| code_writer->NewLine(); |
| |
| if (!GenerateWrapperImports(code_writer, model, err)) { |
| err->Error("Fail to generate imports for wrapper class."); |
| return false; |
| } |
| if (!GenerateWrapperClass(code_writer, model, err)) { |
| err->Error("Fail to generate wrapper class."); |
| return false; |
| } |
| code_writer->NewLine(); |
| return true; |
| } |
| |
| bool GenerateWrapperImports(CodeWriter* code_writer, const ModelInfo& model, |
| ErrorReporter* err) { |
| const std::string support_pkg = "org.tensorflow.lite.support."; |
| std::vector<std::string> imports{ |
| "android.content.Context", |
| "java.io.IOException", |
| "java.nio.ByteBuffer", |
| "java.nio.FloatBuffer", |
| "java.util.Arrays", |
| "java.util.HashMap", |
| "java.util.List", |
| "java.util.Map", |
| "org.tensorflow.lite.DataType", |
| "org.tensorflow.lite.Tensor", |
| "org.tensorflow.lite.Tensor.QuantizationParams", |
| support_pkg + "common.FileUtil", |
| support_pkg + "common.TensorProcessor", |
| support_pkg + "common.ops.CastOp", |
| support_pkg + "common.ops.DequantizeOp", |
| support_pkg + "common.ops.NormalizeOp", |
| support_pkg + "common.ops.QuantizeOp", |
| support_pkg + "label.Category", |
| support_pkg + "label.TensorLabel", |
| support_pkg + "metadata.MetadataExtractor", |
| support_pkg + "metadata.schema.NormalizationOptions", |
| support_pkg + "model.Model", |
| support_pkg + "tensorbuffer.TensorBuffer", |
| }; |
| if (IsImageUsed(model)) { |
| for (const auto& target : |
| {"image.ImageProcessor", "image.TensorImage", "image.ops.ResizeOp", |
| "image.ops.ResizeOp.ResizeMethod"}) { |
| imports.push_back(support_pkg + target); |
| } |
| } |
| |
| std::sort(imports.begin(), imports.end()); |
| for (const auto& target : imports) { |
| code_writer->SetTokenValue("TARGET", target); |
| code_writer->Append("import {{TARGET}};"); |
| } |
| code_writer->NewLine(); |
| return true; |
| } |
| |
| bool GenerateWrapperClass(CodeWriter* code_writer, const ModelInfo& model, |
| ErrorReporter* err) { |
| code_writer->SetTokenValue("MODEL_VERSIONED_NAME", |
| model.model_versioned_name); |
| code_writer->Append( |
| R"(/** Wrapper class of model {{MODEL_VERSIONED_NAME}} */)"); |
| const auto code_block = |
| AsBlock(code_writer, "public class {{MODEL_CLASS_NAME}}"); |
| code_writer->Append(R"(private final Metadata metadata; |
| private final Model model; |
| private static final String MODEL_NAME = "{{MODEL_PATH}}";)"); |
| for (const auto& tensor : model.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append("private {{PROCESSOR_TYPE}} {{NAME}}Preprocessor;"); |
| } |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append("private {{PROCESSOR_TYPE}} {{NAME}}Postprocessor;"); |
| } |
| code_writer->NewLine(); |
| if (!GenerateWrapperOutputs(code_writer, model, err)) { |
| err->Error("Failed to generate output classes"); |
| return false; |
| } |
| code_writer->NewLine(); |
| if (!GenerateWrapperMetadata(code_writer, model, err)) { |
| err->Error("Failed to generate the metadata class"); |
| return false; |
| } |
| code_writer->NewLine(); |
| if (!GenerateWrapperAPI(code_writer, model, err)) { |
| err->Error("Failed to generate the common APIs"); |
| return false; |
| } |
| return true; |
| } |
| |
| bool GenerateWrapperOutputs(CodeWriter* code_writer, const ModelInfo& model, |
| ErrorReporter* err) { |
| code_writer->Append("/** Output wrapper of {@link {{MODEL_CLASS_NAME}}} */"); |
| auto class_block = AsBlock(code_writer, "public static class Outputs"); |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append("private final {{WRAPPER_TYPE}} {{NAME}};"); |
| if (tensor.associated_axis_label_index >= 0) { |
| code_writer->Append("private final List<String> {{NAME}}Labels;"); |
| } |
| code_writer->Append( |
| "private final {{PROCESSOR_TYPE}} {{NAME}}Postprocessor;"); |
| } |
| // Getters |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->NewLine(); |
| if (tensor.associated_axis_label_index >= 0) { |
| if (tensor.content_type == "tensor") { |
| code_writer->Append( |
| R"(public List<Category> get{{NAME_U}}AsCategoryList() { |
| return new TensorLabel({{NAME}}Labels, postprocess{{NAME_U}}({{NAME}})).getCategoryList(); |
| })"); |
| } else { // image |
| err->Warning( |
| "Axis label for images is not supported. The labels will " |
| "be ignored."); |
| } |
| } else { // no label |
| code_writer->Append( |
| R"(public {{WRAPPER_TYPE}} get{{NAME_U}}As{{WRAPPER_TYPE}}() { |
| return postprocess{{NAME_U}}({{NAME}}); |
| })"); |
| } |
| } |
| code_writer->NewLine(); |
| { |
| const auto ctor_block = AsBlock( |
| code_writer, |
| "Outputs(Metadata metadata, {{POSTPROCESSOR_TYPE_PARAM_LIST}})"); |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| if (tensor.content_type == "image") { |
| code_writer->Append( |
| R"({{NAME}} = new TensorImage(metadata.get{{NAME_U}}Type()); |
| {{NAME}}.load(TensorBuffer.createFixedSize(metadata.get{{NAME_U}}Shape(), metadata.get{{NAME_U}}Type()));)"); |
| } else { // FEATURE, UNKNOWN |
| code_writer->Append( |
| "{{NAME}} = " |
| "TensorBuffer.createFixedSize(metadata.get{{NAME_U}}Shape(), " |
| "metadata.get{{NAME_U}}Type());"); |
| } |
| if (tensor.associated_axis_label_index >= 0) { |
| code_writer->Append("{{NAME}}Labels = metadata.get{{NAME_U}}Labels();"); |
| } |
| code_writer->Append( |
| "this.{{NAME}}Postprocessor = {{NAME}}Postprocessor;"); |
| } |
| } |
| code_writer->NewLine(); |
| { |
| const auto get_buffer_block = |
| AsBlock(code_writer, "Map<Integer, Object> getBuffer()"); |
| code_writer->Append("Map<Integer, Object> outputs = new HashMap<>();"); |
| for (int i = 0; i < model.outputs.size(); i++) { |
| SetCodeWriterWithTensorInfo(code_writer, model.outputs[i]); |
| code_writer->SetTokenValue("ID", std::to_string(i)); |
| code_writer->Append("outputs.put({{ID}}, {{NAME}}.getBuffer());"); |
| } |
| code_writer->Append("return outputs;"); |
| } |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->NewLine(); |
| { |
| auto processor_block = |
| AsBlock(code_writer, |
| "private {{WRAPPER_TYPE}} " |
| "postprocess{{NAME_U}}({{WRAPPER_TYPE}} {{WRAPPER_NAME}})"); |
| code_writer->Append( |
| "return {{NAME}}Postprocessor.process({{WRAPPER_NAME}});"); |
| } |
| } |
| return true; |
| } |
| |
| bool GenerateWrapperMetadata(CodeWriter* code_writer, const ModelInfo& model, |
| ErrorReporter* err) { |
| code_writer->Append( |
| "/** Metadata accessors of {@link {{MODEL_CLASS_NAME}}} */"); |
| const auto class_block = AsBlock(code_writer, "public static class Metadata"); |
| for (const auto& tensor : model.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append(R"(private final int[] {{NAME}}Shape; |
| private final DataType {{NAME}}DataType; |
| private final QuantizationParams {{NAME}}QuantizationParams;)"); |
| if (tensor.normalization_unit >= 0) { |
| code_writer->Append(R"(private final float[] {{NAME}}Mean; |
| private final float[] {{NAME}}Stddev;)"); |
| } |
| } |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append(R"(private final int[] {{NAME}}Shape; |
| private final DataType {{NAME}}DataType; |
| private final QuantizationParams {{NAME}}QuantizationParams;)"); |
| if (tensor.normalization_unit >= 0) { |
| code_writer->Append(R"(private final float[] {{NAME}}Mean; |
| private final float[] {{NAME}}Stddev;)"); |
| } |
| if (tensor.associated_axis_label_index >= 0 || |
| tensor.associated_value_label_index >= 0) { |
| code_writer->Append("private final List<String> {{NAME}}Labels;"); |
| } |
| } |
| code_writer->NewLine(); |
| { |
| const auto ctor_block = AsBlock( |
| code_writer, |
| "public Metadata(ByteBuffer buffer, Model model) throws IOException"); |
| code_writer->Append( |
| "MetadataExtractor extractor = new MetadataExtractor(buffer);"); |
| for (int i = 0; i < model.inputs.size(); i++) { |
| SetCodeWriterWithTensorInfo(code_writer, model.inputs[i]); |
| code_writer->SetTokenValue("ID", std::to_string(i)); |
| code_writer->Append( |
| R"(Tensor {{NAME}}Tensor = model.getInputTensor({{ID}}); |
| {{NAME}}Shape = {{NAME}}Tensor.shape(); |
| {{NAME}}DataType = {{NAME}}Tensor.dataType(); |
| {{NAME}}QuantizationParams = {{NAME}}Tensor.quantizationParams();)"); |
| if (model.inputs[i].normalization_unit >= 0) { |
| code_writer->Append( |
| R"(NormalizationOptions {{NAME}}NormalizationOptions = |
| (NormalizationOptions) extractor.getInputTensorMetadata({{ID}}).processUnits({{NORMALIZATION_UNIT}}).options(new NormalizationOptions()); |
| FloatBuffer {{NAME}}MeanBuffer = {{NAME}}NormalizationOptions.meanAsByteBuffer().asFloatBuffer(); |
| {{NAME}}Mean = new float[{{NAME}}MeanBuffer.limit()]; |
| {{NAME}}MeanBuffer.get({{NAME}}Mean); |
| FloatBuffer {{NAME}}StddevBuffer = {{NAME}}NormalizationOptions.stdAsByteBuffer().asFloatBuffer(); |
| {{NAME}}Stddev = new float[{{NAME}}StddevBuffer.limit()]; |
| {{NAME}}StddevBuffer.get({{NAME}}Stddev);)"); |
| } |
| } |
| for (int i = 0; i < model.outputs.size(); i++) { |
| SetCodeWriterWithTensorInfo(code_writer, model.outputs[i]); |
| code_writer->SetTokenValue("ID", std::to_string(i)); |
| code_writer->Append( |
| R"(Tensor {{NAME}}Tensor = model.getOutputTensor({{ID}}); |
| {{NAME}}Shape = {{NAME}}Tensor.shape(); |
| {{NAME}}DataType = {{NAME}}Tensor.dataType(); |
| {{NAME}}QuantizationParams = {{NAME}}Tensor.quantizationParams();)"); |
| if (model.outputs[i].normalization_unit >= 0) { |
| code_writer->Append( |
| R"(NormalizationOptions {{NAME}}NormalizationOptions = |
| (NormalizationOptions) extractor.getInputTensorMetadata({{ID}}).processUnits({{NORMALIZATION_UNIT}}).options(new NormalizationOptions()); |
| FloatBuffer {{NAME}}MeanBuffer = {{NAME}}NormalizationOptions.meanAsByteBuffer().asFloatBuffer(); |
| {{NAME}}Mean = new float[{{NAME}}MeanBuffer.limit()]; |
| {{NAME}}MeanBuffer.get({{NAME}}Mean); |
| FloatBuffer {{NAME}}StddevBuffer = {{NAME}}NormalizationOptions.stdAsByteBuffer().asFloatBuffer(); |
| {{NAME}}Stddev = new float[{{NAME}}StddevBuffer.limit()]; |
| {{NAME}}StddevBuffer.get({{NAME}}Stddev);)"); |
| } |
| if (model.outputs[i].associated_axis_label_index >= 0) { |
| code_writer->Append(R"(String {{NAME}}LabelsFileName = |
| extractor.getOutputTensorMetadata({{ID}}).associatedFiles({{ASSOCIATED_AXIS_LABEL_INDEX}}).name(); |
| {{NAME}}Labels = FileUtil.loadLabels(extractor.getAssociatedFile({{NAME}}LabelsFileName));)"); |
| } else if (model.outputs[i].associated_value_label_index >= 0) { |
| code_writer->Append(R"(String {{NAME}}LabelsFileName = |
| extractor.getOutputTensorMetadata({{ID}}).associatedFiles({{ASSOCIATED_VALUE_LABEL_INDEX}}).name(); |
| {{NAME}}Labels = FileUtil.loadLabels(extractor.getAssociatedFile({{NAME}}LabelsFileName));)"); |
| } |
| } |
| } |
| for (const auto& tensor : model.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append(R"( |
| public int[] get{{NAME_U}}Shape() { |
| return Arrays.copyOf({{NAME}}Shape, {{NAME}}Shape.length); |
| } |
| |
| public DataType get{{NAME_U}}Type() { |
| return {{NAME}}DataType; |
| } |
| |
| public QuantizationParams get{{NAME_U}}QuantizationParams() { |
| return {{NAME}}QuantizationParams; |
| })"); |
| if (tensor.normalization_unit >= 0) { |
| code_writer->Append(R"( |
| public float[] get{{NAME_U}}Mean() { |
| return Arrays.copyOf({{NAME}}Mean, {{NAME}}Mean.length); |
| } |
| |
| public float[] get{{NAME_U}}Stddev() { |
| return Arrays.copyOf({{NAME}}Stddev, {{NAME}}Stddev.length); |
| })"); |
| } |
| } |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append(R"( |
| public int[] get{{NAME_U}}Shape() { |
| return Arrays.copyOf({{NAME}}Shape, {{NAME}}Shape.length); |
| } |
| |
| public DataType get{{NAME_U}}Type() { |
| return {{NAME}}DataType; |
| } |
| |
| public QuantizationParams get{{NAME_U}}QuantizationParams() { |
| return {{NAME}}QuantizationParams; |
| })"); |
| if (tensor.normalization_unit >= 0) { |
| code_writer->Append(R"( |
| public float[] get{{NAME_U}}Mean() { |
| return Arrays.copyOf({{NAME}}Mean, {{NAME}}Mean.length); |
| } |
| |
| public float[] get{{NAME_U}}Stddev() { |
| return Arrays.copyOf({{NAME}}Stddev, {{NAME}}Stddev.length); |
| })"); |
| } |
| if (tensor.associated_axis_label_index >= 0 || |
| tensor.associated_value_label_index >= 0) { |
| code_writer->Append(R"( |
| public List<String> get{{NAME_U}}Labels() { |
| return {{NAME}}Labels; |
| })"); |
| } |
| } |
| return true; |
| } |
| |
| bool GenerateWrapperAPI(CodeWriter* code_writer, const ModelInfo& model, |
| ErrorReporter* err) { |
| code_writer->Append(R"(public Metadata getMetadata() { |
| return metadata; |
| } |
| )"); |
| code_writer->Append(R"(/** |
| * Creates interpreter and loads associated files if needed. |
| * |
| * @throws IOException if an I/O error occurs when loading the tflite model. |
| */ |
| public static {{MODEL_CLASS_NAME}} newInstance(Context context) throws IOException { |
| return newInstance(context, MODEL_NAME, new Model.Options.Builder().build()); |
| } |
| |
| /** |
| * Creates interpreter and loads associated files if needed, but loading another model in the same |
| * input / output structure with the original one. |
| * |
| * @throws IOException if an I/O error occurs when loading the tflite model. |
| */ |
| public static {{MODEL_CLASS_NAME}} newInstance(Context context, String modelPath) throws IOException { |
| return newInstance(context, modelPath, new Model.Options.Builder().build()); |
| } |
| |
| /** |
| * Creates interpreter and loads associated files if needed, with running options configured. |
| * |
| * @throws IOException if an I/O error occurs when loading the tflite model. |
| */ |
| public static {{MODEL_CLASS_NAME}} newInstance(Context context, Model.Options runningOptions) throws IOException { |
| return newInstance(context, MODEL_NAME, runningOptions); |
| } |
| |
| /** |
| * Creates interpreter for a user-specified model. |
| * |
| * @throws IOException if an I/O error occurs when loading the tflite model. |
| */ |
| public static {{MODEL_CLASS_NAME}} newInstance(Context context, String modelPath, Model.Options runningOptions) throws IOException { |
| Model model = Model.createModel(context, modelPath, runningOptions); |
| Metadata metadata = new Metadata(model.getData(), model); |
| MyImageClassifier instance = new MyImageClassifier(model, metadata);)"); |
| for (const auto& tensor : model.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append( |
| R"( instance.reset{{NAME_U}}Preprocessor( |
| instance.buildDefault{{NAME_U}}Preprocessor());)"); |
| } |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append( |
| R"( instance.reset{{NAME_U}}Postprocessor( |
| instance.buildDefault{{NAME_U}}Postprocessor());)"); |
| } |
| code_writer->Append(R"( return instance; |
| } |
| )"); |
| |
| // Pre, post processor setters |
| for (const auto& tensor : model.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append(R"( |
| public void reset{{NAME_U}}Preprocessor({{PROCESSOR_TYPE}} processor) { |
| {{NAME}}Preprocessor = processor; |
| })"); |
| } |
| for (const auto& tensor : model.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append(R"( |
| public void reset{{NAME_U}}Postprocessor({{PROCESSOR_TYPE}} processor) { |
| {{NAME}}Postprocessor = processor; |
| })"); |
| } |
| // Process method |
| code_writer->Append(R"( |
| /** Triggers the model. */ |
| public Outputs process({{INPUT_TYPE_PARAM_LIST}}) { |
| Outputs outputs = new Outputs(metadata, {{POSTPROCESSORS_LIST}}); |
| Object[] inputBuffers = preprocessInputs({{INPUTS_LIST}}); |
| model.run(inputBuffers, outputs.getBuffer()); |
| return outputs; |
| } |
| |
| /** Closes the model. */ |
| public void close() { |
| model.close(); |
| } |
| )"); |
| { |
| auto block = |
| AsBlock(code_writer, |
| "private {{MODEL_CLASS_NAME}}(Model model, Metadata metadata)"); |
| code_writer->Append(R"(this.model = model; |
| this.metadata = metadata;)"); |
| } |
| for (const auto& tensor : model.inputs) { |
| code_writer->NewLine(); |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| auto block = AsBlock( |
| code_writer, |
| "private {{PROCESSOR_TYPE}} buildDefault{{NAME_U}}Preprocessor()"); |
| code_writer->Append( |
| "{{PROCESSOR_TYPE}}.Builder builder = new " |
| "{{PROCESSOR_TYPE}}.Builder()"); |
| if (tensor.content_type == "image") { |
| code_writer->Append(R"( .add(new ResizeOp( |
| metadata.get{{NAME_U}}Shape()[1], |
| metadata.get{{NAME_U}}Shape()[2], |
| ResizeMethod.NEAREST_NEIGHBOR)))"); |
| } |
| if (tensor.normalization_unit >= 0) { |
| code_writer->Append( |
| R"( .add(new NormalizeOp(metadata.get{{NAME_U}}Mean(), metadata.get{{NAME_U}}Stddev())))"); |
| } |
| code_writer->Append( |
| R"( .add(new QuantizeOp( |
| metadata.get{{NAME_U}}QuantizationParams().getZeroPoint(), |
| metadata.get{{NAME_U}}QuantizationParams().getScale())) |
| .add(new CastOp(metadata.get{{NAME_U}}Type())); |
| return builder.build();)"); |
| } |
| for (const auto& tensor : model.outputs) { |
| code_writer->NewLine(); |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| auto block = AsBlock( |
| code_writer, |
| "private {{PROCESSOR_TYPE}} buildDefault{{NAME_U}}Postprocessor()"); |
| code_writer->AppendNoNewLine( |
| R"({{PROCESSOR_TYPE}}.Builder builder = new {{PROCESSOR_TYPE}}.Builder() |
| .add(new DequantizeOp( |
| metadata.get{{NAME_U}}QuantizationParams().getZeroPoint(), |
| metadata.get{{NAME_U}}QuantizationParams().getScale())))"); |
| if (tensor.normalization_unit >= 0) { |
| code_writer->AppendNoNewLine(R"( |
| .add(new NormalizeOp(metadata.get{{NAME_U}}Mean(), metadata.get{{NAME_U}}Stddev())))"); |
| } |
| code_writer->Append(R"(; |
| return builder.build();)"); |
| } |
| code_writer->NewLine(); |
| { |
| const auto block = |
| AsBlock(code_writer, |
| "private Object[] preprocessInputs({{INPUT_TYPE_PARAM_LIST}})"); |
| CodeWriter param_list_gen(err); |
| for (const auto& tensor : model.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, tensor); |
| code_writer->Append("{{NAME}} = {{NAME}}Preprocessor.process({{NAME}});"); |
| SetCodeWriterWithTensorInfo(¶m_list_gen, tensor); |
| param_list_gen.AppendNoNewLine("{{NAME}}.getBuffer(), "); |
| } |
| param_list_gen.Backspace(2); |
| code_writer->AppendNoNewLine("return new Object[] {"); |
| code_writer->AppendNoNewLine(param_list_gen.ToString()); |
| code_writer->Append("};"); |
| } |
| return true; |
| } |
| |
| bool GenerateBuildGradleContent(CodeWriter* code_writer, |
| const ModelInfo& model_info) { |
| code_writer->Append(R"(buildscript { |
| repositories { |
| google() |
| jcenter() |
| } |
| dependencies { |
| classpath 'com.android.tools.build:gradle:3.2.1' |
| } |
| } |
| |
| allprojects { |
| repositories { |
| google() |
| jcenter() |
| flatDir { |
| dirs 'libs' |
| } |
| } |
| } |
| |
| apply plugin: 'com.android.library' |
| |
| android { |
| compileSdkVersion 29 |
| defaultConfig { |
| targetSdkVersion 29 |
| versionCode 1 |
| versionName "1.0" |
| } |
| aaptOptions { |
| noCompress "tflite" |
| } |
| compileOptions { |
| sourceCompatibility = '1.8' |
| targetCompatibility = '1.8' |
| } |
| lintOptions { |
| abortOnError false |
| } |
| } |
| |
| configurations { |
| libMetadata |
| } |
| |
| dependencies { |
| libMetadata 'org.tensorflow:tensorflow-lite-support:0.0.0-experimental-metadata-monolithic' |
| } |
| |
| task downloadLibs(type: Sync) { |
| from configurations.libMetadata |
| into "$buildDir/libs" |
| rename 'tensorflow-lite-support-0.0.0-experimental-metadata-monolithic.jar', "tensorflow-lite-support-metadata.jar" |
| } |
| |
| preBuild.dependsOn downloadLibs |
| |
| dependencies { |
| compileOnly 'org.checkerframework:checker-qual:2.5.8' |
| api 'org.tensorflow:tensorflow-lite:0.0.0-nightly' |
| api 'org.tensorflow:tensorflow-lite-support:0.0.0-nightly' |
| api files("$buildDir/libs/tensorflow-lite-support-metadata.jar") |
| implementation 'org.apache.commons:commons-compress:1.19' |
| })"); |
| return true; |
| } |
| |
| bool GenerateAndroidManifestContent(CodeWriter* code_writer, |
| const ModelInfo& model_info) { |
| code_writer->Append(R"(<?xml version="1.0" encoding="utf-8"?> |
| <manifest xmlns:android="http://schemas.android.com/apk/res/android" |
| package="{{PACKAGE}}"> |
| </manifest>)"); |
| return true; |
| } |
| |
| bool GenerateDocContent(CodeWriter* code_writer, const ModelInfo& model_info) { |
| code_writer->Append("# {{MODEL_CLASS_NAME}} Usage"); |
| // TODO(b/158651848) Generate imports for TFLS util types like TensorImage. |
| code_writer->AppendNoNewLine(R"( |
| ``` |
| import {{PACKAGE}}.{{MODEL_CLASS_NAME}}; |
| |
| // 1. Initialize the Model |
| {{MODEL_CLASS_NAME}} model = null; |
| |
| try { |
| model = {{MODEL_CLASS_NAME}}.newInstance(context); // android.content.Context |
| } catch (IOException e) { |
| e.printStackTrace(); |
| } |
| |
| if (model != null) { |
| |
| // 2. Set the inputs)"); |
| for (const auto& t : model_info.inputs) { |
| SetCodeWriterWithTensorInfo(code_writer, t); |
| if (t.content_type == "image") { |
| code_writer->Append(R"( |
| // Prepare tensor "{{NAME}}" from a Bitmap with ARGB_8888 format. |
| Bitmap bitmap = ...; |
| TensorImage {{NAME}} = TensorImage.fromBitmap(bitmap); |
| // Alternatively, load the input tensor "{{NAME}}" from pixel values. |
| // Check out TensorImage documentation to load other image data structures. |
| // int[] pixelValues = ...; |
| // int[] shape = ...; |
| // TensorImage {{NAME}} = new TensorImage(); |
| // {{NAME}}.load(pixelValues, shape);)"); |
| } else { |
| code_writer->Append(R"( |
| // Prepare input tensor "{{NAME}}" from an array. |
| // Check out TensorBuffer documentation to load other data structures. |
| TensorBuffer {{NAME}} = ...; |
| int[] values = ...; |
| int[] shape = ...; |
| {{NAME}}.load(values, shape);)"); |
| } |
| } |
| code_writer->Append(R"( |
| // 3. Run the model |
| {{MODEL_CLASS_NAME}}.Outputs outputs = model.process({{INPUTS_LIST}});)"); |
| code_writer->Append(R"( |
| // 4. Retrieve the results)"); |
| for (const auto& t : model_info.outputs) { |
| SetCodeWriterWithTensorInfo(code_writer, t); |
| if (t.associated_axis_label_index >= 0) { |
| code_writer->SetTokenValue("WRAPPER_TYPE", "List<Category>"); |
| code_writer->Append( |
| " List<Category> {{NAME}} = " |
| "outputs.get{{NAME_U}}AsCategoryList();"); |
| } else { |
| code_writer->Append( |
| " {{WRAPPER_TYPE}} {{NAME}} = " |
| "outputs.get{{NAME_U}}As{{WRAPPER_TYPE}}();"); |
| } |
| } |
| code_writer->Append(R"(} |
| ```)"); |
| return true; |
| } |
| |
| GenerationResult::File GenerateWrapperFile(const std::string& module_root, |
| const ModelInfo& model_info, |
| ErrorReporter* err) { |
| const auto java_path = JoinPath(module_root, "src/main/java"); |
| const auto package_path = |
| JoinPath(java_path, ConvertPackageToPath(model_info.package_name)); |
| const auto file_path = |
| JoinPath(package_path, model_info.model_class_name + JAVA_EXT); |
| |
| CodeWriter code_writer(err); |
| code_writer.SetIndentString(" "); |
| SetCodeWriterWithModelInfo(&code_writer, model_info); |
| |
| if (!GenerateWrapperFileContent(&code_writer, model_info, err)) { |
| err->Error("Generating Java wrapper content failed."); |
| } |
| |
| const auto java_file = code_writer.ToString(); |
| return GenerationResult::File{file_path, java_file}; |
| } |
| |
| GenerationResult::File GenerateBuildGradle(const std::string& module_root, |
| const ModelInfo& model_info, |
| ErrorReporter* err) { |
| const auto file_path = JoinPath(module_root, "build.gradle"); |
| CodeWriter code_writer(err); |
| SetCodeWriterWithModelInfo(&code_writer, model_info); |
| if (!GenerateBuildGradleContent(&code_writer, model_info)) { |
| err->Error("Generating build.gradle failed."); |
| } |
| const auto content = code_writer.ToString(); |
| return GenerationResult::File{file_path, content}; |
| } |
| |
| GenerationResult::File GenerateAndroidManifest(const std::string& module_root, |
| const ModelInfo& model_info, |
| ErrorReporter* err) { |
| const auto file_path = JoinPath(module_root, "src/main/AndroidManifest.xml"); |
| CodeWriter code_writer(err); |
| SetCodeWriterWithModelInfo(&code_writer, model_info); |
| if (!GenerateAndroidManifestContent(&code_writer, model_info)) { |
| err->Error("Generating AndroidManifest.xml failed."); |
| } |
| return GenerationResult::File{file_path, code_writer.ToString()}; |
| } |
| |
| GenerationResult::File GenerateDoc(const std::string& module_root, |
| const ModelInfo& model_info, |
| ErrorReporter* err) { |
| std::string lower = model_info.model_class_name; |
| for (int i = 0; i < lower.length(); i++) { |
| lower[i] = std::tolower(lower[i]); |
| } |
| const auto file_path = JoinPath(module_root, lower + ".md"); |
| CodeWriter code_writer(err); |
| SetCodeWriterWithModelInfo(&code_writer, model_info); |
| if (!GenerateDocContent(&code_writer, model_info)) { |
| err->Error("Generating doc failed."); |
| } |
| return GenerationResult::File{file_path, code_writer.ToString()}; |
| } |
| |
| } // namespace |
| |
| AndroidJavaGenerator::AndroidJavaGenerator(const std::string& module_root) |
| : CodeGenerator(), module_root_(module_root) {} |
| |
| GenerationResult AndroidJavaGenerator::Generate( |
| const Model* model, const std::string& package_name, |
| const std::string& model_class_name, const std::string& model_asset_path) { |
| GenerationResult result; |
| if (model == nullptr) { |
| err_.Error( |
| "Cannot read model from the buffer. Codegen will generate nothing."); |
| return result; |
| } |
| const ModelMetadata* metadata = GetMetadataFromModel(model); |
| if (metadata == nullptr) { |
| err_.Error( |
| "Cannot find TFLite Metadata in the model. Codegen will generate " |
| "nothing."); |
| return result; |
| } |
| details_android_java::ModelInfo model_info = CreateModelInfo( |
| metadata, package_name, model_class_name, model_asset_path, &err_); |
| result.files.push_back(GenerateWrapperFile(module_root_, model_info, &err_)); |
| result.files.push_back(GenerateBuildGradle(module_root_, model_info, &err_)); |
| result.files.push_back( |
| GenerateAndroidManifest(module_root_, model_info, &err_)); |
| result.files.push_back(GenerateDoc(module_root_, model_info, &err_)); |
| return result; |
| } |
| |
| GenerationResult AndroidJavaGenerator::Generate( |
| const char* model_storage, const std::string& package_name, |
| const std::string& model_class_name, const std::string& model_asset_path) { |
| const Model* model = GetModel(model_storage); |
| return Generate(model, package_name, model_class_name, model_asset_path); |
| } |
| |
| std::string AndroidJavaGenerator::GetErrorMessage() { |
| return err_.GetMessage(); |
| } |
| |
| } // namespace codegen |
| } // namespace support |
| } // namespace tflite |
