| 37 | } |
| 38 | |
| 39 | void Compute(OpKernelContext* context) override { |
| 40 | const Tensor& input = context->input(0); |
| 41 | OP_REQUIRES(context, input.dims() == 2, |
| 42 | errors::InvalidArgument("input must be 2-dimensional", |
| 43 | input.shape().DebugString())); |
| 44 | Spectrogram spectrogram; |
| 45 | OP_REQUIRES(context, spectrogram.Initialize(window_size_, stride_), |
| 46 | errors::InvalidArgument( |
| 47 | "Spectrogram initialization failed for window size ", |
| 48 | window_size_, " and stride ", stride_)); |
| 49 | |
| 50 | const auto input_as_matrix = input.matrix<float>(); |
| 51 | |
| 52 | const int64 sample_count = input.dim_size(0); |
| 53 | const int64 channel_count = input.dim_size(1); |
| 54 | |
| 55 | const int64 output_width = spectrogram.output_frequency_channels(); |
| 56 | const int64 length_minus_window = (sample_count - window_size_); |
| 57 | int64 output_height; |
| 58 | if (length_minus_window < 0) { |
| 59 | output_height = 0; |
| 60 | } else { |
| 61 | output_height = 1 + (length_minus_window / stride_); |
| 62 | } |
| 63 | const int64 output_slices = channel_count; |
| 64 | |
| 65 | Tensor* output_tensor = nullptr; |
| 66 | OP_REQUIRES_OK( |
| 67 | context, |
| 68 | context->allocate_output( |
| 69 | 0, TensorShape({output_slices, output_height, output_width}), |
| 70 | &output_tensor)); |
| 71 | auto output_flat = output_tensor->flat<float>().data(); |
| 72 | |
| 73 | std::vector<float> input_for_channel(sample_count); |
| 74 | for (int64 channel = 0; channel < channel_count; ++channel) { |
| 75 | float* output_slice = |
| 76 | output_flat + (channel * output_height * output_width); |
| 77 | for (int i = 0; i < sample_count; ++i) { |
| 78 | input_for_channel[i] = input_as_matrix(i, channel); |
| 79 | } |
| 80 | std::vector<std::vector<float>> spectrogram_output; |
| 81 | OP_REQUIRES(context, |
| 82 | spectrogram.ComputeSquaredMagnitudeSpectrogram( |
| 83 | input_for_channel, &spectrogram_output), |
| 84 | errors::InvalidArgument("Spectrogram compute failed")); |
| 85 | OP_REQUIRES(context, (spectrogram_output.size() == output_height), |
| 86 | errors::InvalidArgument( |
| 87 | "Spectrogram size calculation failed: Expected height ", |
| 88 | output_height, " but got ", spectrogram_output.size())); |
| 89 | OP_REQUIRES(context, |
| 90 | spectrogram_output.empty() || |
| 91 | (spectrogram_output[0].size() == output_width), |
| 92 | errors::InvalidArgument( |
| 93 | "Spectrogram size calculation failed: Expected width ", |
| 94 | output_width, " but got ", spectrogram_output[0].size())); |
| 95 | for (int row_index = 0; row_index < output_height; ++row_index) { |
| 96 | const std::vector<float>& spectrogram_row = |
nothing calls this directly
no test coverage detected