| 48 | |
| 49 | template <typename Kernel> |
| 50 | Status ComputeSpansCore(OpKernelContext* context, const Kernel& kernel, |
| 51 | const int64 output_size, const int64 input_size, |
| 52 | const float scale, const float translate, |
| 53 | const bool antialias, Spans* spans) { |
| 54 | // When sampling, we need the inverse scale and translation, to map from an |
| 55 | // output to an input pixel. |
| 56 | const float inv_scale = 1.0 / scale; |
| 57 | const float inv_translate = -inv_scale * translate; |
| 58 | // When downsampling the kernel should be scaled since we want to low pass |
| 59 | // filter and interpolate, but when upsampling it should not be since we only |
| 60 | // want to interpolate. |
| 61 | const float kernel_scale = antialias ? std::max(inv_scale, 1.0f) : 1.0f; |
| 62 | spans->span_size = std::min( |
| 63 | 2 * static_cast<int>(std::ceil(kernel.Radius() * kernel_scale)) + 1, |
| 64 | static_cast<int>(input_size)); |
| 65 | AllocatorAttributes alloc_attr; |
| 66 | alloc_attr.set_on_host(true); |
| 67 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 68 | tensorflow::DT_INT32, tensorflow::TensorShape({output_size}), |
| 69 | &spans->starts, alloc_attr)); |
| 70 | auto starts_vec = spans->starts.vec<int32>(); |
| 71 | TF_RETURN_IF_ERROR(context->allocate_temp( |
| 72 | tensorflow::DT_FLOAT, |
| 73 | tensorflow::TensorShape({spans->span_size * output_size}), |
| 74 | &spans->weights, alloc_attr)); |
| 75 | auto weights_vec = spans->weights.vec<float>(); |
| 76 | weights_vec.setZero(); |
| 77 | |
| 78 | const float one_over_kernel_scale = 1.0f / kernel_scale; |
| 79 | int max_span_size = 0; |
| 80 | std::vector<float> temp_weights; |
| 81 | for (int x = 0; x < output_size; ++x) { |
| 82 | const float col_f = x + 0.5f; |
| 83 | const float sample_f = col_f * inv_scale + inv_translate; |
| 84 | |
| 85 | // Don't sample when the sampling location is outside the source image. |
| 86 | if (sample_f < 0 || sample_f > input_size) { |
| 87 | // Add an empty span. |
| 88 | starts_vec(x) = 0; |
| 89 | continue; |
| 90 | } |
| 91 | int64 span_start = |
| 92 | std::ceil(sample_f - kernel.Radius() * kernel_scale - 0.5f); |
| 93 | int64 span_end = |
| 94 | std::floor(sample_f + kernel.Radius() * kernel_scale - 0.5f); |
| 95 | span_start = Clamp(static_cast<int64>(0), input_size - 1, span_start); |
| 96 | span_end = Clamp(static_cast<int64>(0), input_size - 1, span_end) + 1; |
| 97 | const int this_span_size = span_end - span_start; |
| 98 | if (this_span_size > spans->span_size) { |
| 99 | return errors::Internal(Printf("Span is too large: %d vs %d.", |
| 100 | this_span_size, spans->span_size)); |
| 101 | } |
| 102 | float total_weight_sum = 0.0f; |
| 103 | temp_weights.clear(); |
| 104 | for (int source = span_start; source < span_end; ++source) { |
| 105 | float kernel_pos = static_cast<float>(source) + 0.5f - sample_f; |
| 106 | float weight = kernel(std::abs(kernel_pos * one_over_kernel_scale)); |
| 107 | total_weight_sum += weight; |
no test coverage detected