• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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