1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow_lite_support/cc/task/vision/core/frame_buffer.h"
17
18 namespace tflite {
19 namespace task {
20 namespace vision {
21
22 using ::tflite::support::StatusOr;
23
24 namespace {
25
26 // Returns whether the input `format` is a supported YUV format.
IsSupportedYuvFormat(FrameBuffer::Format format)27 bool IsSupportedYuvFormat(FrameBuffer::Format format) {
28 return format == FrameBuffer::Format::kNV21 ||
29 format == FrameBuffer::Format::kNV12 ||
30 format == FrameBuffer::Format::kYV12 ||
31 format == FrameBuffer::Format::kYV21;
32 }
33
34 // Returns supported 1-plane FrameBuffer in YuvData structure.
GetYuvDataFromOnePlaneFrameBuffer(const FrameBuffer & source)35 StatusOr<FrameBuffer::YuvData> GetYuvDataFromOnePlaneFrameBuffer(
36 const FrameBuffer& source) {
37 if (!IsSupportedYuvFormat(source.format())) {
38 return absl::InvalidArgumentError(
39 "The source FrameBuffer format is not part of YUV420 family.");
40 }
41
42 FrameBuffer::YuvData result;
43 const int y_buffer_size =
44 source.plane(0).stride.row_stride_bytes * source.dimension().height;
45 const int uv_buffer_size =
46 ((source.plane(0).stride.row_stride_bytes + 1) / 2) *
47 ((source.dimension().height + 1) / 2);
48 result.y_buffer = source.plane(0).buffer;
49 result.y_row_stride = source.plane(0).stride.row_stride_bytes;
50 result.uv_row_stride = result.y_row_stride;
51
52 if (source.format() == FrameBuffer::Format::kNV21) {
53 result.v_buffer = result.y_buffer + y_buffer_size;
54 result.u_buffer = result.v_buffer + 1;
55 result.uv_pixel_stride = 2;
56 // If y_row_stride equals to the frame width and is an odd value,
57 // uv_row_stride = y_row_stride + 1, otherwise uv_row_stride = y_row_stride.
58 if (result.y_row_stride == source.dimension().width &&
59 result.y_row_stride % 2 == 1) {
60 result.uv_row_stride = (result.y_row_stride + 1) / 2 * 2;
61 }
62 } else if (source.format() == FrameBuffer::Format::kNV12) {
63 result.u_buffer = result.y_buffer + y_buffer_size;
64 result.v_buffer = result.u_buffer + 1;
65 result.uv_pixel_stride = 2;
66 // If y_row_stride equals to the frame width and is an odd value,
67 // uv_row_stride = y_row_stride + 1, otherwise uv_row_stride = y_row_stride.
68 if (result.y_row_stride == source.dimension().width &&
69 result.y_row_stride % 2 == 1) {
70 result.uv_row_stride = (result.y_row_stride + 1) / 2 * 2;
71 }
72 } else if (source.format() == FrameBuffer::Format::kYV21) {
73 result.u_buffer = result.y_buffer + y_buffer_size;
74 result.v_buffer = result.u_buffer + uv_buffer_size;
75 result.uv_pixel_stride = 1;
76 result.uv_row_stride = (result.y_row_stride + 1) / 2;
77 } else if (source.format() == FrameBuffer::Format::kYV12) {
78 result.v_buffer = result.y_buffer + y_buffer_size;
79 result.u_buffer = result.v_buffer + uv_buffer_size;
80 result.uv_pixel_stride = 1;
81 result.uv_row_stride = (result.y_row_stride + 1) / 2;
82 }
83 return result;
84 }
85
86 // Returns supported 2-plane FrameBuffer in YuvData structure.
GetYuvDataFromTwoPlaneFrameBuffer(const FrameBuffer & source)87 StatusOr<FrameBuffer::YuvData> GetYuvDataFromTwoPlaneFrameBuffer(
88 const FrameBuffer& source) {
89 if (source.format() != FrameBuffer::Format::kNV12 &&
90 source.format() != FrameBuffer::Format::kNV21) {
91 return absl::InvalidArgumentError("Unsupported YUV planar format.");
92 }
93
94 FrameBuffer::YuvData result;
95 // Y plane
96 result.y_buffer = source.plane(0).buffer;
97 // All plane strides
98 result.y_row_stride = source.plane(0).stride.row_stride_bytes;
99 result.uv_row_stride = source.plane(1).stride.row_stride_bytes;
100 result.uv_pixel_stride = 2;
101
102 if (source.format() == FrameBuffer::Format::kNV12) {
103 // Y and UV interleaved format
104 result.u_buffer = source.plane(1).buffer;
105 result.v_buffer = result.u_buffer + 1;
106 } else {
107 // Y and VU interleaved format
108 result.v_buffer = source.plane(1).buffer;
109 result.u_buffer = result.v_buffer + 1;
110 }
111 return result;
112 }
113
114 // Returns supported 3-plane FrameBuffer in YuvData structure. Note that NV21
115 // and NV12 are included in the supported Yuv formats. Technically, NV21 and
116 // NV12 should not be described by the 3-plane format. Historically, NV21 is
117 // used loosely such that it can also be used to describe YV21 format. For
118 // backwards compatibility, FrameBuffer supports NV21/NV12 with 3-plane format
119 // but such usage is discouraged
GetYuvDataFromThreePlaneFrameBuffer(const FrameBuffer & source)120 StatusOr<FrameBuffer::YuvData> GetYuvDataFromThreePlaneFrameBuffer(
121 const FrameBuffer& source) {
122 if (!IsSupportedYuvFormat(source.format())) {
123 return absl::InvalidArgumentError(
124 "The source FrameBuffer format is not part of YUV420 family.");
125 }
126
127 if (source.plane(1).stride.row_stride_bytes !=
128 source.plane(2).stride.row_stride_bytes ||
129 source.plane(1).stride.pixel_stride_bytes !=
130 source.plane(2).stride.pixel_stride_bytes) {
131 return absl::InternalError("Unsupported YUV planar format.");
132 }
133 FrameBuffer::YuvData result;
134 if (source.format() == FrameBuffer::Format::kNV21 ||
135 source.format() == FrameBuffer::Format::kYV12) {
136 // Y follow by VU order. The VU chroma planes can be interleaved or
137 // planar.
138 result.y_buffer = source.plane(0).buffer;
139 result.v_buffer = source.plane(1).buffer;
140 result.u_buffer = source.plane(2).buffer;
141 result.y_row_stride = source.plane(0).stride.row_stride_bytes;
142 result.uv_row_stride = source.plane(1).stride.row_stride_bytes;
143 result.uv_pixel_stride = source.plane(1).stride.pixel_stride_bytes;
144 } else {
145 // Y follow by UV order. The UV chroma planes can be interleaved or
146 // planar.
147 result.y_buffer = source.plane(0).buffer;
148 result.u_buffer = source.plane(1).buffer;
149 result.v_buffer = source.plane(2).buffer;
150 result.y_row_stride = source.plane(0).stride.row_stride_bytes;
151 result.uv_row_stride = source.plane(1).stride.row_stride_bytes;
152 result.uv_pixel_stride = source.plane(1).stride.pixel_stride_bytes;
153 }
154 return result;
155 }
156
157 } // namespace
158
GetYuvDataFromFrameBuffer(const FrameBuffer & source)159 StatusOr<FrameBuffer::YuvData> FrameBuffer::GetYuvDataFromFrameBuffer(
160 const FrameBuffer& source) {
161 if (!IsSupportedYuvFormat(source.format())) {
162 return absl::InvalidArgumentError(
163 "The source FrameBuffer format is not part of YUV420 family.");
164 }
165
166 if (source.plane_count() == 1) {
167 return GetYuvDataFromOnePlaneFrameBuffer(source);
168 } else if (source.plane_count() == 2) {
169 return GetYuvDataFromTwoPlaneFrameBuffer(source);
170 } else if (source.plane_count() == 3) {
171 return GetYuvDataFromThreePlaneFrameBuffer(source);
172 }
173 return absl::InvalidArgumentError(
174 "The source FrameBuffer must be consisted by 1, 2, or 3 planes");
175 }
176
177 } // namespace vision
178 } // namespace task
179 } // namespace tflite
180