Skip to content

Commit

Permalink
Add nearest-neighbor resizing to FrameBufferUtils
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 482315989
  • Loading branch information
@tensorflower-gardener @tflite-support-robot
tensorflower-gardener authored and tflite-support-robot committed Oct 19, 2022
1 parent 508139c commit 5c1a046
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 19 deletions.
6 changes: 6 additions & 0 deletions tensorflow_lite_support/cc/task/vision/utils/frame_buffer_utils.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,12 @@ absl::Status FrameBufferUtils::Resize(const FrameBuffer& buffer,
return utils_->Resize(buffer, output_buffer);
}

absl::Status FrameBufferUtils::ResizeNearestNeighbor(
const FrameBuffer& buffer, FrameBuffer* output_buffer) {
TFLITE_DCHECK(utils_ != nullptr);
return utils_->ResizeNearestNeighbor(buffer, output_buffer);
}

absl::Status FrameBufferUtils::Rotate(const FrameBuffer& buffer,
RotationDegree rotation,
FrameBuffer* output_buffer) {
Expand Down
11 changes: 10 additions & 1 deletion tensorflow_lite_support/cc/task/vision/utils/frame_buffer_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -215,14 +215,23 @@ class FrameBufferUtils {
absl::Status Crop(const FrameBuffer& buffer, int x0, int y0, int x1, int y1,
FrameBuffer* output_buffer);

// Performs resizing operation.
// Performs resizing operation with bilinear interpolation.
//
// The resize dimension is determined based on output_buffer's size metadata.
//
// The output_buffer should have metadata populated and its backing buffer
// should be big enough to store the operation result.
absl::Status Resize(const FrameBuffer& buffer, FrameBuffer* output_buffer);

// Performs resizing operation with nearest-neighbor interpolation.
//
// The resize dimension is determined based on output_buffer's size metadata.
//
// The output_buffer should have metadata populated and its backing buffer
// should be big enough to store the operation result.
absl::Status ResizeNearestNeighbor(const FrameBuffer& buffer,
FrameBuffer* output_buffer);

// Performs rotation operation.
//
// The rotation is specified in counter-clockwise direction.
Expand Down
13 changes: 12 additions & 1 deletion tensorflow_lite_support/cc/task/vision/utils/frame_buffer_utils_interface.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,8 @@ class FrameBufferUtilsInterface {
virtual absl::Status Crop(const FrameBuffer& buffer, int x0, int y0, int x1,
int y1, FrameBuffer* output_buffer) = 0;

// Resizes `buffer` to the size of the given `output_buffer`.
// Resizes `buffer` to the size of the given `output_buffer` using bilinear
// interpolation.
//
// The resize dimension is determined based on the size of `output_buffer`.
//
Expand All @@ -49,6 +50,16 @@ class FrameBufferUtilsInterface {
virtual absl::Status Resize(const FrameBuffer& buffer,
FrameBuffer* output_buffer) = 0;

// Resizes `buffer` to the size of the given `output_buffer` using
// nearest-neighbor interpolation.
//
// The resize dimension is determined based on the size of `output_buffer`.
//
// The `output_buffer` should have metadata populated and its backing buffer
// should be big enough to store the operation result.
virtual absl::Status ResizeNearestNeighbor(const FrameBuffer& buffer,
FrameBuffer* output_buffer) = 0;

// Rotates `buffer` counter-clockwise by the given `angle_deg` (in degrees).
//
// When rotating by 90 degrees, the top-right corner of `buffer` becomes
Expand Down
67 changes: 51 additions & 16 deletions tensorflow_lite_support/cc/task/vision/utils/libyuv_frame_buffer_utils.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include "absl/strings/str_format.h" // from @com_google_absl
#include "libyuv.h" // from @libyuv
#include "libyuv/convert_argb.h" // from @libyuv
#include "libyuv/scale.h" // from @libyuv
#include "tensorflow_lite_support/cc/common.h"
#include "tensorflow_lite_support/cc/port/integral_types.h"
#include "tensorflow_lite_support/cc/port/status_macros.h"
Expand Down Expand Up @@ -329,7 +330,9 @@ absl::Status ConvertFromYv(const FrameBuffer& buffer,
}

// Resizes YV12/YV21 `buffer` to the target `output_buffer`.
absl::Status ResizeYv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
absl::Status ResizeYv(
const FrameBuffer& buffer, FrameBuffer* output_buffer,
libyuv::FilterMode interpolation = libyuv::FilterMode::kFilterBilinear) {
ASSIGN_OR_RETURN(FrameBuffer::YuvData input_data,
FrameBuffer::GetYuvDataFromFrameBuffer(buffer));
ASSIGN_OR_RETURN(FrameBuffer::YuvData output_data,
Expand All @@ -344,7 +347,7 @@ absl::Status ResizeYv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
const_cast<uint8_t*>(output_data.u_buffer), output_data.uv_row_stride,
const_cast<uint8_t*>(output_data.v_buffer), output_data.uv_row_stride,
output_buffer->dimension().width, output_buffer->dimension().height,
libyuv::FilterMode::kFilterBilinear);
interpolation);
if (ret != 0) {
return CreateStatusWithPayload(
StatusCode::kUnknown, "Libyuv I420Scale operation failed.",
Expand All @@ -354,7 +357,9 @@ absl::Status ResizeYv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
}

// Resizes NV12/NV21 `buffer` to the target `output_buffer`.
absl::Status ResizeNv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
absl::Status ResizeNv(
const FrameBuffer& buffer, FrameBuffer* output_buffer,
libyuv::FilterMode interpolation = libyuv::FilterMode::kFilterBilinear) {
ASSIGN_OR_RETURN(FrameBuffer::YuvData input_data,
FrameBuffer::GetYuvDataFromFrameBuffer(buffer));
ASSIGN_OR_RETURN(FrameBuffer::YuvData output_data,
Expand All @@ -372,7 +377,7 @@ absl::Status ResizeNv(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
buffer.dimension().height, const_cast<uint8_t*>(output_data.y_buffer),
output_data.y_row_stride, const_cast<uint8_t*>(dst_uv),
output_data.uv_row_stride, output_buffer->dimension().width,
output_buffer->dimension().height, libyuv::FilterMode::kFilterBilinear);
output_buffer->dimension().height, interpolation);

if (ret != 0) {
return CreateStatusWithPayload(
Expand Down Expand Up @@ -1094,7 +1099,9 @@ absl::Status FlipHorizontallyPlane(const FrameBuffer& buffer,
return absl::OkStatus();
}

absl::Status ResizeRgb(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
absl::Status ResizeRgb(
const FrameBuffer& buffer, FrameBuffer* output_buffer,
libyuv::FilterMode interpolation = libyuv::FilterMode::kFilterBilinear) {
if (buffer.plane_count() > 1) {
return CreateStatusWithPayload(
StatusCode::kInternal,
Expand Down Expand Up @@ -1125,7 +1132,7 @@ absl::Status ResizeRgb(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
argb_buffer.get(), argb_row_bytes, buffer.dimension().width,
buffer.dimension().height, resized_argb_buffer.get(),
resized_argb_row_bytes, output_buffer->dimension().width,
output_buffer->dimension().height, libyuv::FilterMode::kFilterBilinear);
output_buffer->dimension().height, interpolation);
if (ret != 0) {
return CreateStatusWithPayload(
StatusCode::kUnknown, "Libyuv ARGBScale operation failed.",
Expand Down Expand Up @@ -1166,7 +1173,9 @@ absl::Status FlipHorizontallyRgb(const FrameBuffer& buffer,
#endif // LIBYUV_VERSION >= 1747
}

absl::Status ResizeRgba(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
absl::Status ResizeRgba(
const FrameBuffer& buffer, FrameBuffer* output_buffer,
libyuv::FilterMode interpolation = libyuv::FilterMode::kFilterBilinear) {
if (buffer.plane_count() > 1) {
return CreateStatusWithPayload(
StatusCode::kInternal,
Expand All @@ -1180,7 +1189,7 @@ absl::Status ResizeRgba(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
const_cast<uint8*>(output_buffer->plane(0).buffer),
output_buffer->plane(0).stride.row_stride_bytes,
output_buffer->dimension().width, output_buffer->dimension().height,
libyuv::FilterMode::kFilterBilinear);
interpolation);
if (ret != 0) {
return CreateStatusWithPayload(
StatusCode::kUnknown, "Libyuv ARGBScale operation failed.",
Expand Down Expand Up @@ -1289,21 +1298,23 @@ absl::Status FlipVerticallyYv(const FrameBuffer& buffer,

// Resize `buffer` to metadata defined in `output_buffer`. This
// method assumes buffer has pixel stride equals to 1 (grayscale equivalent).
absl::Status ResizeGray(const FrameBuffer& buffer, FrameBuffer* output_buffer) {
absl::Status ResizeGray(
const FrameBuffer& buffer, FrameBuffer* output_buffer,
libyuv::FilterMode interpolation = libyuv::FilterMode::kFilterBilinear) {
if (buffer.plane_count() > 1) {
return CreateStatusWithPayload(
StatusCode::kInternal,
absl::StrFormat("Only single plane is supported for format %i.",
buffer.format()),
TfLiteSupportStatus::kImageProcessingError);
}
libyuv::ScalePlane(
buffer.plane(0).buffer, buffer.plane(0).stride.row_stride_bytes,
buffer.dimension().width, buffer.dimension().height,
const_cast<uint8*>(output_buffer->plane(0).buffer),
output_buffer->plane(0).stride.row_stride_bytes,
output_buffer->dimension().width, output_buffer->dimension().height,
libyuv::FilterMode::kFilterBilinear);
libyuv::ScalePlane(buffer.plane(0).buffer,
buffer.plane(0).stride.row_stride_bytes,
buffer.dimension().width, buffer.dimension().height,
const_cast<uint8*>(output_buffer->plane(0).buffer),
output_buffer->plane(0).stride.row_stride_bytes,
output_buffer->dimension().width,
output_buffer->dimension().height, interpolation);
return absl::OkStatus();
}

Expand Down Expand Up @@ -1394,6 +1405,30 @@ absl::Status LibyuvFrameBufferUtils::Resize(const FrameBuffer& buffer,
}
}

absl::Status LibyuvFrameBufferUtils::ResizeNearestNeighbor(
const FrameBuffer& buffer, FrameBuffer* output_buffer) {
RETURN_IF_ERROR(ValidateResizeBufferInputs(buffer, *output_buffer));
switch (buffer.format()) {
case FrameBuffer::Format::kYV12:
case FrameBuffer::Format::kYV21:
return ResizeYv(buffer, output_buffer, libyuv::FilterMode::kFilterNone);
case FrameBuffer::Format::kNV12:
case FrameBuffer::Format::kNV21:
return ResizeNv(buffer, output_buffer, libyuv::FilterMode::kFilterNone);
case FrameBuffer::Format::kRGB:
return ResizeRgb(buffer, output_buffer, libyuv::FilterMode::kFilterNone);
case FrameBuffer::Format::kRGBA:
return ResizeRgba(buffer, output_buffer, libyuv::FilterMode::kFilterNone);
case FrameBuffer::Format::kGRAY:
return ResizeGray(buffer, output_buffer, libyuv::FilterMode::kFilterNone);
default:
return CreateStatusWithPayload(
StatusCode::kInternal,
absl::StrFormat("Format %i is not supported.", buffer.format()),
TfLiteSupportStatus::kImageProcessingError);
}
}

absl::Status LibyuvFrameBufferUtils::Rotate(const FrameBuffer& buffer,
int angle_deg,
FrameBuffer* output_buffer) {
Expand Down
8 changes: 7 additions & 1 deletion tensorflow_lite_support/cc/task/vision/utils/libyuv_frame_buffer_utils.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,16 @@ class LibyuvFrameBufferUtils : public FrameBufferUtilsInterface {
absl::Status Crop(const FrameBuffer& buffer, int x0, int y0, int x1, int y1,
FrameBuffer* output_buffer) override;

// Resizes `buffer` to the size of the given `output_buffer`.
// Resizes `buffer` to the size of the given `output_buffer` using bilinear
// interpolation.
absl::Status Resize(const FrameBuffer& buffer,
FrameBuffer* output_buffer) override;

// Resizes `buffer` to the size of the given `output_buffer` using
// nearest-neighbor interpolation.
absl::Status ResizeNearestNeighbor(const FrameBuffer& buffer,
FrameBuffer* output_buffer) override;

// Rotates `buffer` counter-clockwise by the given `angle_deg` (in degrees).
//
// The given angle must be a multiple of 90 degrees.
Expand Down

0 comments on commit 5c1a046

Please sign in to comment.