Project import generated by Copybara.

GitOrigin-RevId: 19a829ffd755edb43e54d20c0e7b9348512d5108

mediapipe/util/BUILD

@@ -35,6 +35,12 @@ mediapipe_proto_library(
     visibility = ["//visibility:public"],
 )
 
+mediapipe_proto_library(
+    name = "label_map_proto",
+    srcs = ["label_map.proto"],
+    visibility = ["//visibility:public"],
+)
+
 mediapipe_proto_library(
     name = "render_data_proto",
     srcs = ["render_data.proto"],
@@ -124,6 +130,19 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "label_map_util",
+    srcs = ["label_map_util.cc"],
+    hdrs = ["label_map_util.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":label_map_cc_proto",
+        "//mediapipe/framework/port:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
 cc_library(
     name = "annotation_renderer",
     srcs = ["annotation_renderer.cc"],

mediapipe/util/label_map.proto

@@ -0,0 +1,40 @@
+// Copyright 2022 The MediaPipe Authors.
+//
+// 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.
+
+syntax = "proto2";
+
+package mediapipe;
+
+// Mapping a numerical class index output to a Knowledge Graph entity
+// ID or any other string label representing this class. Optionally it is
+// possible to specify an additional display name (in a given language) which is
+// typically used for display purposes.
+message LabelMapItem {
+  // Label name.
+  // E.g. name = "/m/02xwb"
+  optional string name = 1;
+
+  // Display name.
+  // E.g. display_name = "Fruit"
+  optional string display_name = 2;
+
+  // Optional list of children (e.g. subcategories) used to represent a
+  // hierarchy.
+  repeated string child_name = 3;
+}
+
+// Mapping from index to a label map item.
+message LabelMap {
+  map<int64, LabelMapItem> index_to_item = 1;
+}

mediapipe/util/label_map_util.cc

@@ -0,0 +1,78 @@
+// Copyright 2022 The MediaPipe Authors.
+//
+// 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.
+
+#include "mediapipe/util/label_map_util.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/strings/str_format.h"
+#include "absl/strings/str_split.h"
+#include "absl/strings/string_view.h"
+#include "mediapipe/framework/port/statusor.h"
+#include "mediapipe/util/label_map.pb.h"
+
+namespace mediapipe {
+
+absl::StatusOr<LabelMap> BuildLabelMapFromFiles(
+    absl::string_view labels_file_contents,
+    absl::string_view display_names_file) {
+  if (labels_file_contents.empty()) {
+    return absl::InvalidArgumentError("Expected non-empty labels file.");
+  }
+  std::vector<absl::string_view> labels =
+      absl::StrSplit(labels_file_contents, '\n');
+  // In most cases, there is an empty line (i.e. newline character) at the end
+  // of the file that needs to be ignored. In such a situation, StrSplit() will
+  // produce a vector with an empty string as final element. Also note that in
+  // case `labels_file_contents` is entirely empty, StrSplit() will produce a
+  // vector with one single empty substring, so there's no out-of-range risk
+  // here.
+  if (labels[labels.size() - 1].empty()) {
+    labels.pop_back();
+  }
+
+  std::vector<LabelMapItem> label_map_items;
+  label_map_items.reserve(labels.size());
+  for (int i = 0; i < labels.size(); ++i) {
+    LabelMapItem item;
+    item.set_name(std::string(labels[i]));
+    label_map_items.emplace_back(item);
+  }
+
+  if (!display_names_file.empty()) {
+    std::vector<std::string> display_names =
+        absl::StrSplit(display_names_file, '\n');
+    // In most cases, there is an empty line (i.e. newline character) at the end
+    // of the file that needs to be ignored. See above.
+    if (display_names[display_names.size() - 1].empty()) {
+      display_names.pop_back();
+    }
+    if (display_names.size() != labels.size()) {
+      return absl::InvalidArgumentError(absl::StrFormat(
+          "Mismatch between number of labels (%d) and display names (%d).",
+          labels.size(), display_names.size()));
+    }
+    for (int i = 0; i < display_names.size(); ++i) {
+      label_map_items[i].set_display_name(display_names[i]);
+    }
+  }
+  LabelMap label_map;
+  for (int i = 0; i < label_map_items.size(); ++i) {
+    (*label_map.mutable_index_to_item())[i] = label_map_items[i];
+  }
+  return label_map;
+}
+
+}  // namespace mediapipe

mediapipe/util/label_map_util.h

@@ -0,0 +1,34 @@
+// Copyright 2022 The MediaPipe Authors.
+//
+// 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.
+
+#ifndef MEDIAPIPE_UTIL_LABEL_MAP_UTIL_H_
+#define MEDIAPIPE_UTIL_LABEL_MAP_UTIL_H_
+
+#include "absl/strings/string_view.h"
+#include "mediapipe/framework/port/statusor.h"
+#include "mediapipe/util/label_map.pb.h"
+
+namespace mediapipe {
+
+// Builds a label map from labels and (optional) display names file contents,
+// both expected to contain one label per line.
+// Returns an error e.g. if there's a mismatch between the number of labels and
+// display names.
+absl::StatusOr<mediapipe::LabelMap> BuildLabelMapFromFiles(
+    absl::string_view labels_file_contents,
+    absl::string_view display_names_file);
+
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_UTIL_LABEL_MAP_UTIL_H_

mediapipe/util/tflite/BUILD

@@ -50,6 +50,17 @@ cc_library(
     alwayslink = 1,
 )
 
+cc_library(
+    name = "error_reporter",
+    srcs = ["error_reporter.cc"],
+    hdrs = ["error_reporter.h"],
+    deps = [
+        "@org_tensorflow//tensorflow/lite:minimal_logging",
+        "@org_tensorflow//tensorflow/lite:stateful_error_reporter",
+        "@org_tensorflow//tensorflow/lite/core/api:error_reporter",
+    ],
+)
+
 cc_library(
     name = "op_resolver",
     srcs = ["op_resolver.cc"],
@@ -108,6 +119,7 @@ cc_library(
     name = "tflite_model_loader",
     srcs = ["tflite_model_loader.cc"],
     hdrs = ["tflite_model_loader.h"],
+    visibility = ["//visibility:public"],
     deps = [
         "//mediapipe/framework/api2:packet",
         "//mediapipe/framework/port:ret_check",

mediapipe/util/tflite/error_reporter.cc

@@ -0,0 +1,49 @@
+// Copyright 2022 The MediaPipe Authors.
+//
+// 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.
+
+#include "mediapipe/util/tflite/error_reporter.h"
+
+#include <cstdarg>
+#include <cstdio>
+#include <cstring>
+#include <string>
+
+#include "tensorflow/lite/minimal_logging.h"
+
+namespace mediapipe {
+namespace util {
+namespace tflite {
+
+ErrorReporter::ErrorReporter() {
+  message_[0] = '\0';
+  previous_message_[0] = '\0';
+}
+
+int ErrorReporter::Report(const char* format, va_list args) {
+  std::strcpy(previous_message_, message_);  // NOLINT
+  message_[0] = '\0';
+  int num_characters = vsnprintf(message_, kBufferSize, format, args);
+  // To mimic tflite::StderrReporter.
+  ::tflite::logging_internal::MinimalLogger::Log(::tflite::TFLITE_LOG_ERROR,
+                                                 "%s", message_);
+  return num_characters;
+}
+
+std::string ErrorReporter::message() { return message_; }
+
+std::string ErrorReporter::previous_message() { return previous_message_; }
+
+}  // namespace tflite
+}  // namespace util
+}  // namespace mediapipe

mediapipe/util/tflite/error_reporter.h

@@ -0,0 +1,52 @@
+// Copyright 2022 The MediaPipe Authors.
+//
+// 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.
+
+#ifndef MEDIAPIPE_UTIL_TFLITE_ERROR_REPORTER_H_
+#define MEDIAPIPE_UTIL_TFLITE_ERROR_REPORTER_H_
+
+#include <string>
+
+#include "tensorflow/lite/core/api/error_reporter.h"
+#include "tensorflow/lite/stateful_error_reporter.h"
+
+namespace mediapipe {
+namespace util {
+namespace tflite {
+
+// An ErrorReporter that logs to stderr and captures the last two messages.
+class ErrorReporter : public ::tflite::StatefulErrorReporter {
+ public:
+  ErrorReporter();
+
+  // We declared two functions with name 'Report', so that the variadic Report
+  // function in tflite::ErrorReporter is hidden.
+  // See https://isocpp.org/wiki/faq/strange-inheritance#hiding-rule.
+  using ::tflite::ErrorReporter::Report;
+
+  int Report(const char* format, std::va_list args) override;
+
+  std::string message() override;
+  std::string previous_message();
+
+ private:
+  static constexpr int kBufferSize = 1024;
+  char message_[kBufferSize];
+  char previous_message_[kBufferSize];
+};
+
+}  // namespace tflite
+}  // namespace util
+}  // namespace mediapipe
+
+#endif  // MEDIAPIPE_UTIL_TFLITE_ERROR_REPORTER_H_