Skip to content

Commit 99779a4

Browse files
ivanseidelivanseidel
committed
Using Example file "label_image"
1 parent 7f9a089 commit 99779a4

1 file changed

Lines changed: 277 additions & 10 deletions

File tree

‎src/api.cc‎

Lines changed: 277 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,19 +2,286 @@
22
#include <memory>
33

44
#include "tensorflow/core/public/version.h"
5-
#include "tensorflow/core/public/tensor_c_api.h"
65

7-
#include "tensorflow/core/lib/core/coding.h"
6+
7+
#include <fstream>
8+
9+
#include "tensorflow/cc/ops/const_op.h"
10+
#include "tensorflow/cc/ops/image_ops.h"
11+
#include "tensorflow/cc/ops/standard_ops.h"
12+
#include "tensorflow/core/framework/graph.pb.h"
13+
#include "tensorflow/core/graph/default_device.h"
14+
#include "tensorflow/core/graph/graph_def_builder.h"
15+
#include "tensorflow/core/lib/core/command_line_flags.h"
816
#include "tensorflow/core/lib/core/errors.h"
917
#include "tensorflow/core/lib/core/stringpiece.h"
10-
#include "tensorflow/core/lib/gtl/array_slice.h"
11-
12-
#include "tensorflow/core/platform/port.h"
13-
#include "tensorflow/core/platform/protobuf.h"
14-
// #include "tensorflow/core/public/session.h"
15-
// #include "tensorflow/core/public/status.h"
16-
// #include "tensorflow/core/public/tensor.h"
17-
// #include "tensorflow/core/public/tensor_shape.h"
18+
#include "tensorflow/core/lib/core/threadpool.h"
19+
#include "tensorflow/core/lib/io/path.h"
20+
#include "tensorflow/core/lib/strings/stringprintf.h"
21+
#include "tensorflow/core/platform/init_main.h"
22+
#include "tensorflow/core/platform/logging.h"
23+
#include "tensorflow/core/public/session.h"
24+
#include "tensorflow/core/public/tensor.h"
25+
26+
// These are all common classes it's handy to reference with no namespace.
27+
using tensorflow::Tensor;
28+
using tensorflow::Status;
29+
using tensorflow::string;
30+
using tensorflow::int32;
31+
32+
// These are the command-line flags the program can understand.
33+
// They define where the graph and input data is located, and what kind of
34+
// input the model expects. If you train your own model, or use something
35+
// other than GoogLeNet you'll need to update these.
36+
TF_DEFINE_string(image,
37+
"tensorflow/examples/label_image/data/grace_hopper.jpg",
38+
"The image to classify (JPEG or PNG).");
39+
TF_DEFINE_string(graph,
40+
"tensorflow/examples/label_image/data/googlenet_graph.pb",
41+
"The location of the GraphDef file containing the protobuf"
42+
" definition of the network.");
43+
TF_DEFINE_string(labels,
44+
"tensorflow/examples/label_image/data/googlenet_labels.txt",
45+
"A text file containing the labels of all the categories, one"
46+
" per line.");
47+
TF_DEFINE_int32(input_width, 224, "Width of the image the network expects.");
48+
TF_DEFINE_int32(input_height, 224, "Height of the image the network expects.");
49+
TF_DEFINE_int32(input_mean, 117, "How much to subtract from input values.");
50+
TF_DEFINE_int32(input_std, 1, "What to divide the input values by.");
51+
TF_DEFINE_string(input_layer, "input", "The name of the input node.");
52+
TF_DEFINE_string(output_layer, "softmax2", "The name of the output node.");
53+
TF_DEFINE_bool(self_test, false, "Whether to run a sanity check on the results.");
54+
TF_DEFINE_string(root_dir, "", "The directory at the root of the data files.");
55+
56+
// Takes a file name, and loads a list of labels from it, one per line, and
57+
// returns a vector of the strings. It pads with empty strings so the length
58+
// of the result is a multiple of 16, because our model expects that.
59+
Status ReadLabelsFile(string file_name, std::vector<string>* result) {
60+
std::ifstream file(file_name);
61+
result->clear();
62+
string line;
63+
while (std::getline(file, line)) {
64+
result->push_back(line);
65+
}
66+
const int padding = 16;
67+
while (result->size() % padding) {
68+
result->emplace_back();
69+
}
70+
return Status::OK();
71+
}
72+
73+
// Given an image file name, read in the data, try to decode it as an image,
74+
// resize it to the requested size, and then scale the values as desired.
75+
Status ReadTensorFromImageFile(string file_name, const int input_height,
76+
const int input_width, const float input_mean,
77+
const float input_std,
78+
std::vector<Tensor>* out_tensors) {
79+
tensorflow::GraphDefBuilder b;
80+
string input_name = "file_reader";
81+
string output_name = "normalized";
82+
tensorflow::Node* file_reader =
83+
tensorflow::ops::ReadFile(tensorflow::ops::Const(file_name, b.opts()),
84+
b.opts().WithName(input_name));
85+
// Now try to figure out what kind of file it is and decode it.
86+
const int wanted_channels = 3;
87+
tensorflow::Node* image_reader;
88+
if (tensorflow::StringPiece(file_name).ends_with(".png")) {
89+
image_reader = tensorflow::ops::DecodePng(
90+
file_reader,
91+
b.opts().WithAttr("channels", wanted_channels).WithName("png_reader"));
92+
} else {
93+
// Assume if it's not a PNG then it must be a JPEG.
94+
image_reader = tensorflow::ops::DecodeJpeg(
95+
file_reader,
96+
b.opts().WithAttr("channels", wanted_channels).WithName("jpeg_reader"));
97+
}
98+
// Now cast the image data to float so we can do normal math on it.
99+
tensorflow::Node* float_caster = tensorflow::ops::Cast(
100+
image_reader, tensorflow::DT_FLOAT, b.opts().WithName("float_caster"));
101+
// The convention for image ops in TensorFlow is that all images are expected
102+
// to be in batches, so that they're four-dimensional arrays with indices of
103+
// [batch, height, width, channel]. Because we only have a single image, we
104+
// have to add a batch dimension of 1 to the start with ExpandDims().
105+
tensorflow::Node* dims_expander = tensorflow::ops::ExpandDims(
106+
float_caster, tensorflow::ops::Const(0, b.opts()), b.opts());
107+
// Bilinearly resize the image to fit the required dimensions.
108+
tensorflow::Node* resized = tensorflow::ops::ResizeBilinear(
109+
dims_expander, tensorflow::ops::Const({input_height, input_width},
110+
b.opts().WithName("size")),
111+
b.opts());
112+
// Subtract the mean and divide by the scale.
113+
tensorflow::ops::Div(
114+
tensorflow::ops::Sub(
115+
resized, tensorflow::ops::Const({input_mean}, b.opts()), b.opts()),
116+
tensorflow::ops::Const({input_std}, b.opts()),
117+
b.opts().WithName(output_name));
118+
119+
// This runs the GraphDef network definition that we've just constructed, and
120+
// returns the results in the output tensor.
121+
tensorflow::GraphDef graph;
122+
TF_RETURN_IF_ERROR(b.ToGraphDef(&graph));
123+
std::unique_ptr<tensorflow::Session> session(
124+
tensorflow::NewSession(tensorflow::SessionOptions()));
125+
TF_RETURN_IF_ERROR(session->Create(graph));
126+
TF_RETURN_IF_ERROR(session->Run({}, {output_name}, {}, out_tensors));
127+
return Status::OK();
128+
}
129+
130+
// Reads a model graph definition from disk, and creates a session object you
131+
// can use to run it.
132+
Status LoadGraph(string graph_file_name,
133+
std::unique_ptr<tensorflow::Session>* session) {
134+
tensorflow::GraphDef graph_def;
135+
Status load_graph_status =
136+
ReadBinaryProto(tensorflow::Env::Default(), graph_file_name, &graph_def);
137+
if (!load_graph_status.ok()) {
138+
return tensorflow::errors::NotFound("Failed to load compute graph at '",
139+
graph_file_name, "'");
140+
}
141+
142+
session->reset(tensorflow::NewSession(tensorflow::SessionOptions()));
143+
Status session_create_status = (*session)->Create(graph_def);
144+
if (!session_create_status.ok()) {
145+
return session_create_status;
146+
}
147+
return Status::OK();
148+
}
149+
150+
// Analyzes the output of the Inception graph to retrieve the highest scores and
151+
// their positions in the tensor, which correspond to categories.
152+
Status GetTopLabels(const std::vector<Tensor>& outputs, int how_many_labels,
153+
Tensor* indices, Tensor* scores) {
154+
tensorflow::GraphDefBuilder b;
155+
string output_name = "top_k";
156+
tensorflow::ops::TopK(tensorflow::ops::Const(outputs[0], b.opts()),
157+
how_many_labels, b.opts().WithName(output_name));
158+
// This runs the GraphDef network definition that we've just constructed, and
159+
// returns the results in the output tensors.
160+
tensorflow::GraphDef graph;
161+
TF_RETURN_IF_ERROR(b.ToGraphDef(&graph));
162+
std::unique_ptr<tensorflow::Session> session(
163+
tensorflow::NewSession(tensorflow::SessionOptions()));
164+
TF_RETURN_IF_ERROR(session->Create(graph));
165+
// The TopK node returns two outputs, the scores and their original indices,
166+
// so we have to append :0 and :1 to specify them both.
167+
std::vector<Tensor> out_tensors;
168+
TF_RETURN_IF_ERROR(session->Run({}, {output_name + ":0", output_name + ":1"},
169+
{}, &out_tensors));
170+
*scores = out_tensors[0];
171+
*indices = out_tensors[1];
172+
return Status::OK();
173+
}
174+
175+
// Given the output of a model run, and the name of a file containing the labels
176+
// this prints out the top five highest-scoring values.
177+
Status PrintTopLabels(const std::vector<Tensor>& outputs,
178+
string labels_file_name) {
179+
std::vector<string> labels;
180+
Status read_labels_status = ReadLabelsFile(labels_file_name, &labels);
181+
if (!read_labels_status.ok()) {
182+
LOG(ERROR) << read_labels_status;
183+
return read_labels_status;
184+
}
185+
const int how_many_labels = 5;
186+
Tensor indices;
187+
Tensor scores;
188+
TF_RETURN_IF_ERROR(GetTopLabels(outputs, how_many_labels, &indices, &scores));
189+
tensorflow::TTypes<float>::Flat scores_flat = scores.flat<float>();
190+
tensorflow::TTypes<int32>::Flat indices_flat = indices.flat<int32>();
191+
for (int pos = 0; pos < how_many_labels; ++pos) {
192+
const int label_index = indices_flat(pos);
193+
const float score = scores_flat(pos);
194+
LOG(INFO) << labels[label_index] << " (" << label_index << "): " << score;
195+
}
196+
return Status::OK();
197+
}
198+
199+
// This is a testing function that returns whether the top label index is the
200+
// one that's expected.
201+
Status CheckTopLabel(const std::vector<Tensor>& outputs, int expected,
202+
bool* is_expected) {
203+
*is_expected = false;
204+
Tensor indices;
205+
Tensor scores;
206+
const int how_many_labels = 1;
207+
TF_RETURN_IF_ERROR(GetTopLabels(outputs, how_many_labels, &indices, &scores));
208+
tensorflow::TTypes<int32>::Flat indices_flat = indices.flat<int32>();
209+
if (indices_flat(0) != expected) {
210+
LOG(ERROR) << "Expected label #" << expected << " but got #"
211+
<< indices_flat(0);
212+
*is_expected = false;
213+
} else {
214+
*is_expected = true;
215+
}
216+
return Status::OK();
217+
}
218+
219+
int main(int argc, char* argv[]) {
220+
// We need to call this to set up global state for TensorFlow.
221+
tensorflow::port::InitMain(argv[0], &argc, &argv);
222+
Status s = tensorflow::ParseCommandLineFlags(&argc, argv);
223+
if (!s.ok()) {
224+
LOG(ERROR) << "Error parsing command line flags: " << s.ToString();
225+
return -1;
226+
}
227+
228+
// First we load and initialize the model.
229+
std::unique_ptr<tensorflow::Session> session;
230+
string graph_path = tensorflow::io::JoinPath(FLAGS_root_dir, FLAGS_graph);
231+
Status load_graph_status = LoadGraph(graph_path, &session);
232+
if (!load_graph_status.ok()) {
233+
LOG(ERROR) << load_graph_status;
234+
return -1;
235+
}
236+
237+
// Get the image from disk as a float array of numbers, resized and normalized
238+
// to the specifications the main graph expects.
239+
std::vector<Tensor> resized_tensors;
240+
string image_path = tensorflow::io::JoinPath(FLAGS_root_dir, FLAGS_image);
241+
Status read_tensor_status = ReadTensorFromImageFile(
242+
image_path, FLAGS_input_height, FLAGS_input_width, FLAGS_input_mean,
243+
FLAGS_input_std, &resized_tensors);
244+
if (!read_tensor_status.ok()) {
245+
LOG(ERROR) << read_tensor_status;
246+
return -1;
247+
}
248+
const Tensor& resized_tensor = resized_tensors[0];
249+
250+
// Actually run the image through the model.
251+
std::vector<Tensor> outputs;
252+
Status run_status = session->Run({{FLAGS_input_layer, resized_tensor}},
253+
{FLAGS_output_layer}, {}, &outputs);
254+
if (!run_status.ok()) {
255+
LOG(ERROR) << "Running model failed: " << run_status;
256+
return -1;
257+
}
258+
259+
// This is for automated testing to make sure we get the expected result with
260+
// the default settings. We know that label 866 (military uniform) should be
261+
// the top label for the Admiral Hopper image.
262+
if (FLAGS_self_test) {
263+
bool expected_matches;
264+
Status check_status = CheckTopLabel(outputs, 866, &expected_matches);
265+
if (!check_status.ok()) {
266+
LOG(ERROR) << "Running check failed: " << check_status;
267+
return -1;
268+
}
269+
if (!expected_matches) {
270+
LOG(ERROR) << "Self-test failed!";
271+
return -1;
272+
}
273+
}
274+
275+
// Do something interesting with the results we've generated.
276+
Status print_status = PrintTopLabels(outputs, FLAGS_labels);
277+
if (!print_status.ok()) {
278+
LOG(ERROR) << "Running print failed: " << print_status;
279+
return -1;
280+
}
281+
282+
return 0;
283+
}
284+
18285

19286
using namespace v8;
20287

0 commit comments

Comments
 (0)