xref: /external/tensorflow/tensorflow/examples/wav_to_spectrogram/wav_to_spectrogram.cc

/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.

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 "tensorflow/examples/wav_to_spectrogram/wav_to_spectrogram.h"

#include <vector>

#include "tensorflow/cc/ops/audio_ops.h"
#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/cc/ops/image_ops.h"
#include "tensorflow/cc/ops/standard_ops.h"
#include "tensorflow/core/framework/graph.pb.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/graph/default_device.h"
#include "tensorflow/core/graph/graph_def_builder.h"
#include "tensorflow/core/lib/core/errors.h"
#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/stringprintf.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/types.h"
#include "tensorflow/core/public/session.h"
#include "tensorflow/core/util/command_line_flags.h"

using tensorflow::DT_FLOAT;
using tensorflow::DT_UINT8;
using tensorflow::Output;
using tensorflow::TensorShape;

// Runs a TensorFlow graph to convert an audio file into a visualization.
tensorflow::Status WavToSpectrogram(const tensorflow::string& input_wav,
                                    tensorflow::int32 window_size,
                                    tensorflow::int32 stride, float brightness,
                                    const tensorflow::string& output_image) {
  auto root = tensorflow::Scope::NewRootScope();
  using namespace tensorflow::ops;  // NOLINT(build/namespaces)
  // The following block creates a TensorFlow graph that:
  //  - Reads and decodes the audio file into a tensor of float samples.
  //  - Creates a float spectrogram from those samples.
  //  - Scales, clamps, and converts that spectrogram to 0 to 255 uint8's.
  //  - Reshapes the tensor so that it's [height, width, 1] for imaging.
  //  - Encodes it as a PNG stream and saves it out to a file.
  Output file_reader =
      tensorflow::ops::ReadFile(root.WithOpName("input_wav"), input_wav);
  DecodeWav wav_decoder =
      DecodeWav(root.WithOpName("wav_decoder"), file_reader);
  Output spectrogram = AudioSpectrogram(root.WithOpName("spectrogram"),
                                        wav_decoder.audio, window_size, stride);
  Output brightness_placeholder =
      Placeholder(root.WithOpName("brightness_placeholder"), DT_FLOAT,
                  Placeholder::Attrs().Shape(TensorShape({})));
  Output mul = Mul(root.WithOpName("mul"), spectrogram, brightness_placeholder);
  Output min_const = Const(root.WithOpName("min_const"), 255.0f);
  Output min = Minimum(root.WithOpName("min"), mul, min_const);
  Output cast = Cast(root.WithOpName("cast"), min, DT_UINT8);
  Output expand_dims_const = Const(root.WithOpName("expand_dims_const"), -1);
  Output expand_dims =
      ExpandDims(root.WithOpName("expand_dims"), cast, expand_dims_const);
  Output squeeze = Squeeze(root.WithOpName("squeeze"), expand_dims,
                           Squeeze::Attrs().Axis({0}));
  Output png_encoder = EncodePng(root.WithOpName("png_encoder"), squeeze);
  tensorflow::ops::WriteFile file_writer = tensorflow::ops::WriteFile(
      root.WithOpName("output_image"), output_image, png_encoder);
  tensorflow::GraphDef graph;
  TF_RETURN_IF_ERROR(root.ToGraphDef(&graph));

  // Build a session object from this graph definition. The power of TensorFlow
  // is that you can reuse complex computations like this, so usually we'd run a
  // lot of different inputs through it. In this example, we're just doing a
  // one-off run, so we'll create it and then use it immediately.
  std::unique_ptr<tensorflow::Session> session(
      tensorflow::NewSession(tensorflow::SessionOptions()));
  TF_RETURN_IF_ERROR(session->Create(graph));

  // We're passing in the brightness as an input, so create a tensor to hold the
  // value.
  tensorflow::Tensor brightness_tensor(DT_FLOAT, TensorShape({}));
  brightness_tensor.scalar<float>()() = brightness;

  // Run the session to analyze the audio and write out the file.
  TF_RETURN_IF_ERROR(
      session->Run({{"brightness_placeholder", brightness_tensor}}, {},
                   {"output_image"}, nullptr));
  return tensorflow::Status::OK();
}