1
2 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
3
4 Licensed under the Apache License, Version 2.0 (the "License");
5 you may not use this file except in compliance with the License.
6 You may obtain a copy of the License at
7
8 http://www.apache.org/licenses/LICENSE-2.0
9
10 Unless required by applicable law or agreed to in writing, software
11 distributed under the License is distributed on an "AS IS" BASIS,
12 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 See the License for the specific language governing permissions and
14 limitations under the License.
15 ==============================================================================*/
16 #include "tensorflow_lite_support/cc/task/vision/core/label_map_item.h"
17
18 #include "absl/strings/str_format.h"
19 #include "absl/strings/str_split.h"
20 #include "tensorflow_lite_support/cc/common.h"
21
22 namespace tflite {
23 namespace task {
24 namespace vision {
25
26 using ::absl::StatusCode;
27 using ::tflite::support::CreateStatusWithPayload;
28 using ::tflite::support::StatusOr;
29 using ::tflite::support::TfLiteSupportStatus;
30
BuildLabelMapFromFiles(absl::string_view labels_file,absl::string_view display_names_file)31 StatusOr<std::vector<LabelMapItem>> BuildLabelMapFromFiles(
32 absl::string_view labels_file, absl::string_view display_names_file) {
33 if (labels_file.empty()) {
34 return CreateStatusWithPayload(StatusCode::kInvalidArgument,
35 "Expected non-empty labels file.",
36 TfLiteSupportStatus::kInvalidArgumentError);
37 }
38 std::vector<absl::string_view> labels = absl::StrSplit(labels_file, '\n');
39 // In most cases, there is an empty line (i.e. newline character) at the end
40 // of the file that needs to be ignored. In such a situation, StrSplit() will
41 // produce a vector with an empty string as final element. Also note that in
42 // case `labels_file` is entirely empty, StrSplit() will produce a vector with
43 // one single empty substring, so there's no out-of-range risk here.
44 if (labels[labels.size() - 1].empty()) {
45 labels.pop_back();
46 }
47
48 std::vector<LabelMapItem> label_map_items;
49 label_map_items.reserve(labels.size());
50 for (int i = 0; i < labels.size(); ++i) {
51 label_map_items.emplace_back(LabelMapItem{.name = std::string(labels[i])});
52 }
53
54 if (!display_names_file.empty()) {
55 std::vector<std::string> display_names =
56 absl::StrSplit(display_names_file, '\n');
57 // In most cases, there is an empty line (i.e. newline character) at the end
58 // of the file that needs to be ignored. See above.
59 if (display_names[display_names.size() - 1].empty()) {
60 display_names.pop_back();
61 }
62 if (display_names.size() != labels.size()) {
63 return CreateStatusWithPayload(
64 StatusCode::kInvalidArgument,
65 absl::StrFormat(
66 "Mismatch between number of labels (%d) and display names (%d).",
67 labels.size(), display_names.size()),
68 TfLiteSupportStatus::kMetadataNumLabelsMismatchError);
69 }
70 for (int i = 0; i < display_names.size(); ++i) {
71 label_map_items[i].display_name = display_names[i];
72 }
73 }
74 return label_map_items;
75 }
76
InitializeFromLabelMap(std::vector<LabelMapItem> label_map_items)77 absl::Status LabelHierarchy::InitializeFromLabelMap(
78 std::vector<LabelMapItem> label_map_items) {
79 parents_map_.clear();
80 for (const LabelMapItem& label : label_map_items) {
81 for (const std::string& child_name : label.child_name) {
82 parents_map_[child_name].insert(label.name);
83 }
84 }
85 if (parents_map_.empty()) {
86 return CreateStatusWithPayload(StatusCode::kInvalidArgument,
87 "Input labelmap is not hierarchical: there "
88 "is no parent-child relationship.",
89 TfLiteSupportStatus::kInvalidArgumentError);
90 }
91 return absl::OkStatus();
92 }
93
HaveAncestorDescendantRelationship(const std::string & ancestor_name,const std::string & descendant_name) const94 bool LabelHierarchy::HaveAncestorDescendantRelationship(
95 const std::string& ancestor_name,
96 const std::string& descendant_name) const {
97 absl::flat_hash_set<std::string> ancestors;
98 GetAncestors(descendant_name, &ancestors);
99 return ancestors.contains(ancestor_name);
100 }
101
GetParents(const std::string & name) const102 absl::flat_hash_set<std::string> LabelHierarchy::GetParents(
103 const std::string& name) const {
104 absl::flat_hash_set<std::string> parents;
105 auto it = parents_map_.find(name);
106 if (it != parents_map_.end()) {
107 for (const std::string& parent_name : it->second) {
108 parents.insert(parent_name);
109 }
110 }
111 return parents;
112 }
113
GetAncestors(const std::string & name,absl::flat_hash_set<std::string> * ancestors) const114 void LabelHierarchy::GetAncestors(
115 const std::string& name,
116 absl::flat_hash_set<std::string>* ancestors) const {
117 const absl::flat_hash_set<std::string> parents = GetParents(name);
118 for (const std::string& parent_name : parents) {
119 auto it = ancestors->insert(parent_name);
120 if (it.second) {
121 GetAncestors(parent_name, ancestors);
122 }
123 }
124 }
125
126 } // namespace vision
127 } // namespace task
128 } // namespace tflite
129