diff options
Diffstat (limited to 'tensorflow/examples/label_image/main.cc')
| -rw-r--r-- | tensorflow/examples/label_image/main.cc | 7 |
1 files changed, 4 insertions, 3 deletions
diff --git a/tensorflow/examples/label_image/main.cc b/tensorflow/examples/label_image/main.cc index 63bc39de6c..baa65d3243 100644 --- a/tensorflow/examples/label_image/main.cc +++ b/tensorflow/examples/label_image/main.cc @@ -49,6 +49,7 @@ limitations under the License. #include "tensorflow/core/lib/core/stringpiece.h" #include "tensorflow/core/lib/core/threadpool.h" #include "tensorflow/core/lib/io/path.h" +#include "tensorflow/core/lib/strings/str_util.h" #include "tensorflow/core/lib/strings/stringprintf.h" #include "tensorflow/core/platform/env.h" #include "tensorflow/core/platform/init_main.h" @@ -137,15 +138,15 @@ Status ReadTensorFromImageFile(const string& file_name, const int input_height, // Now try to figure out what kind of file it is and decode it. const int wanted_channels = 3; tensorflow::Output image_reader; - if (tensorflow::StringPiece(file_name).ends_with(".png")) { + if (tensorflow::str_util::EndsWith(file_name, ".png")) { image_reader = DecodePng(root.WithOpName("png_reader"), file_reader, DecodePng::Channels(wanted_channels)); - } else if (tensorflow::StringPiece(file_name).ends_with(".gif")) { + } else if (tensorflow::str_util::EndsWith(file_name, ".gif")) { // gif decoder returns 4-D tensor, remove the first dim image_reader = Squeeze(root.WithOpName("squeeze_first_dim"), DecodeGif(root.WithOpName("gif_reader"), file_reader)); - } else if (tensorflow::StringPiece(file_name).ends_with(".bmp")) { + } else if (tensorflow::str_util::EndsWith(file_name, ".bmp")) { image_reader = DecodeBmp(root.WithOpName("bmp_reader"), file_reader); } else { // Assume if it's neither a PNG nor a GIF then it must be a JPEG. |