aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/ffmpeg/encode_audio_op.cc
blob: c00cccd84615c2fccaf5c8053c162d8b980eb7fe (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
// Copyright 2016 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 <limits>

#include "tensorflow/contrib/ffmpeg/ffmpeg_lib.h"
#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"

namespace tensorflow {
namespace ffmpeg {
namespace {

/*
 * Encoding implementation, shared across V1 and V2 ops. Creates a new
 * output in the context.
 */
void Encode(OpKernelContext* context, const Tensor& contents,
            const string& file_format, const int32 bits_per_second,
            const int32 samples_per_second) {
  std::vector<float> samples;
  samples.reserve(contents.NumElements());
  for (int32 i = 0; i < contents.NumElements(); ++i) {
    samples.push_back(contents.flat<float>()(i));
  }
  const int32 channel_count = contents.dim_size(1);
  string encoded_audio;
  OP_REQUIRES_OK(
      context, CreateAudioFile(file_format, bits_per_second, samples_per_second,
                               channel_count, samples, &encoded_audio));

  // Copy the encoded audio file to the output tensor.
  Tensor* output = nullptr;
  OP_REQUIRES_OK(context, context->allocate_output(0, TensorShape(), &output));
  output->scalar<string>()() = encoded_audio;
}

}  // namespace

/*
 * Supersedes `EncodeAudioOp`. Allows all parameters to be inputs
 * instead of attributes, so that the sample rate (and, probably less
 * usefully, the output file format) can be given as tensors rather than
 * constants only.
 */
class EncodeAudioOpV2 : public OpKernel {
 public:
  explicit EncodeAudioOpV2(OpKernelConstruction* context) : OpKernel(context) {}

  void Compute(OpKernelContext* context) override {
    OP_REQUIRES(
        context, context->num_inputs() == 4,
        errors::InvalidArgument("EncodeAudio requires exactly four inputs."));

    const Tensor& contents = context->input(0);
    const Tensor& file_format_tensor = context->input(1);
    const Tensor& samples_per_second_tensor = context->input(2);
    const Tensor& bits_per_second_tensor = context->input(3);

    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(contents.shape()),
                errors::InvalidArgument(
                    "sampled_audio must be a rank-2 tensor but got shape ",
                    contents.shape().DebugString()));
    OP_REQUIRES(
        context, contents.NumElements() <= std::numeric_limits<int32>::max(),
        errors::InvalidArgument(
            "sampled_audio cannot have more than 2^31 entries. Shape = ",
            contents.shape().DebugString()));
    OP_REQUIRES(context, TensorShapeUtils::IsScalar(file_format_tensor.shape()),
                errors::InvalidArgument(
                    "file_format must be a rank-0 tensor but got shape ",
                    file_format_tensor.shape().DebugString()));
    OP_REQUIRES(context,
                TensorShapeUtils::IsScalar(samples_per_second_tensor.shape()),
                errors::InvalidArgument(
                    "samples_per_second must be a rank-0 tensor but got shape ",
                    samples_per_second_tensor.shape().DebugString()));
    OP_REQUIRES(context,
                TensorShapeUtils::IsScalar(bits_per_second_tensor.shape()),
                errors::InvalidArgument(
                    "bits_per_second must be a rank-0 tensor but got shape ",
                    bits_per_second_tensor.shape().DebugString()));

    const string file_format =
        str_util::Lowercase(file_format_tensor.scalar<string>()());
    const int32 samples_per_second =
        samples_per_second_tensor.scalar<int32>()();
    const int32 bits_per_second = bits_per_second_tensor.scalar<int32>()();

    OP_REQUIRES(context, file_format == "wav",
                errors::InvalidArgument(
                    "file_format must be \"wav\", but got: ", file_format));
    OP_REQUIRES(context, samples_per_second > 0,
                errors::InvalidArgument(
                    "samples_per_second must be positive, but got: ",
                    samples_per_second));
    OP_REQUIRES(
        context, bits_per_second > 0,
        errors::InvalidArgument("bits_per_second must be positive, but got: ",
                                bits_per_second));

    Encode(context, contents, file_format, bits_per_second, samples_per_second);
  }
};

REGISTER_KERNEL_BUILDER(Name("EncodeAudioV2").Device(DEVICE_CPU),
                        EncodeAudioOpV2);

REGISTER_OP("EncodeAudioV2")
    .Input("sampled_audio: float")
    .Input("file_format: string")
    .Input("samples_per_second: int32")
    .Input("bits_per_second: int32")
    .Output("contents: string")
    .SetShapeFn(shape_inference::ScalarShape)
    .Doc(R"doc(
Processes a `Tensor` containing sampled audio with the number of channels
and length of the audio specified by the dimensions of the `Tensor`. The
audio is converted into a string that, when saved to disk, will be equivalent
to the audio in the specified audio format.

The input audio has one row of the tensor for each channel in the audio file.
Each channel contains audio samples starting at the beginning of the audio and
having `1/samples_per_second` time between them. The output file will contain
all of the audio channels contained in the tensor.

sampled_audio: A rank-2 float tensor containing all tracks of the audio.
    Dimension 0 is time and dimension 1 is the channel.
file_format: A string or rank-0 string tensor describing the audio file
    format. This value must be `"wav"`.
samples_per_second: The number of samples per second that the audio should
    have, as an int or rank-0 `int32` tensor. This value must be
    positive.
bits_per_second: The approximate bitrate of the encoded audio file, as
    an int or rank-0 `int32` tensor. This is ignored by the "wav" file
    format.
contents: The binary audio file contents, as a rank-0 string tensor.
)doc");

/*
 * Deprecated in favor of EncodeAudioOpV2.
 */
class EncodeAudioOp : public OpKernel {
 public:
  explicit EncodeAudioOp(OpKernelConstruction* context) : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("file_format", &file_format_));
    file_format_ = str_util::Lowercase(file_format_);
    OP_REQUIRES(context, file_format_ == "wav",
                errors::InvalidArgument("file_format arg must be \"wav\"."));

    OP_REQUIRES_OK(
        context, context->GetAttr("samples_per_second", &samples_per_second_));
    OP_REQUIRES(context, samples_per_second_ > 0,
                errors::InvalidArgument("samples_per_second must be > 0."));
    OP_REQUIRES_OK(context,
                   context->GetAttr("bits_per_second", &bits_per_second_));
  }

  void Compute(OpKernelContext* context) override {
    // Get and verify the input data.
    OP_REQUIRES(
        context, context->num_inputs() == 1,
        errors::InvalidArgument("EncodeAudio requires exactly one input."));
    const Tensor& contents = context->input(0);
    OP_REQUIRES(context, TensorShapeUtils::IsMatrix(contents.shape()),
                errors::InvalidArgument(
                    "sampled_audio must be a rank 2 tensor but got shape ",
                    contents.shape().DebugString()));
    OP_REQUIRES(
        context, contents.NumElements() <= std::numeric_limits<int32>::max(),
        errors::InvalidArgument(
            "sampled_audio cannot have more than 2^31 entries. Shape = ",
            contents.shape().DebugString()));

    Encode(context, contents, file_format_, bits_per_second_,
           samples_per_second_);
  }

 private:
  string file_format_;
  int32 samples_per_second_;
  int32 bits_per_second_;
};

REGISTER_KERNEL_BUILDER(Name("EncodeAudio").Device(DEVICE_CPU), EncodeAudioOp);

REGISTER_OP("EncodeAudio")
    .Input("sampled_audio: float")
    .Output("contents: string")
    .Attr("file_format: string")
    .Attr("samples_per_second: int")
    .Attr("bits_per_second: int = 192000")
    .SetShapeFn(shape_inference::ScalarShape)
    .Doc(R"doc(
Processes a `Tensor` containing sampled audio with the number of channels
and length of the audio specified by the dimensions of the `Tensor`. The
audio is converted into a string that, when saved to disk, will be equivalent
to the audio in the specified audio format.

The input audio has one row of the tensor for each channel in the audio file.
Each channel contains audio samples starting at the beginning of the audio and
having `1/samples_per_second` time between them. The output file will contain
all of the audio channels contained in the tensor.

sampled_audio: A rank 2 tensor containing all tracks of the audio. Dimension 0
    is time and dimension 1 is the channel.
contents: The binary audio file contents.
file_format: A string describing the audio file format. This must be "wav".
samples_per_second: The number of samples per second that the audio should have.
bits_per_second: The approximate bitrate of the encoded audio file. This is
    ignored by the "wav" file format.
)doc");

}  // namespace ffmpeg
}  // namespace tensorflow