Skip to content

Commit

Permalink
Support YUV image format in Task Vision Library
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 355204226
  • Loading branch information
@lu-wang-g @tflite-support-robot
lu-wang-g authored and tflite-support-robot committed Feb 2, 2021
1 parent 799bd17 commit efac3d1
Show file tree
Hide file tree
Showing 4 changed files with 74 additions and 12 deletions.
23 changes: 21 additions & 2 deletions ...ite_support/java/src/java/org/tensorflow/lite/task/vision/classifier/ImageClassifier.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -374,8 +374,18 @@ private ImageClassifierOptions(Builder builder) {
/**
* Performs actual classification on the provided image.
*
* @param image a {@link TensorImage} object that represents an RGB image
* <p>{@link ImageClassifier} supports the following {@link TensorImage} color space types:
*
* <ul>
* <li>{@link ColorSpaceType#RGB}
* <li>{@link ColorSpaceType#NV21}
* <li>{@link ColorSpaceType#YV12}
* <li>{@link ColorSpaceType#YV21}
* </ul>
*
* @param image a UINT8 {@link TensorImage} object that represents an RGB or YUV image
* @throws AssertionError if error occurs when classifying the image from the native code
* @throws IllegalArgumentException if the color space type of image is unsupported
*/
public List<Classifications> classify(TensorImage image) {
return classify(image, ImageProcessingOptions.builder().build());
Expand All @@ -393,7 +403,16 @@ public List<Classifications> classify(TensorImage image) {
* defaults to {@link ImageProcessingOptions#Orientation#TOP_LEFT}.
* </ul>
*
* @param image a {@link TensorImage} object that represents an RGB image
* <p>{@link ImageClassifier} supports the following {@link TensorImage} color space types:
*
* <ul>
* <li>{@link ColorSpaceType#RGB}
* <li>{@link ColorSpaceType#NV21}
* <li>{@link ColorSpaceType#YV12}
* <li>{@link ColorSpaceType#YV21}
* </ul>
*
* @param image a UINT8 {@link TensorImage} object that represents an RGB or YUV image
* @throws AssertionError if error occurs when classifying the image from the native code
* @throws IllegalArgumentException if the color space type of image is unsupported
*/
Expand Down
5 changes: 5 additions & 0 deletions ...ow_lite_support/java/src/java/org/tensorflow/lite/task/vision/core/BaseVisionTaskApi.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,11 @@ private FrameBufferData createFrameBuffer(TensorImage image, int orientation) {
ColorSpaceType colorSpaceType = image.getColorSpaceType();
switch (colorSpaceType) {
case RGB:
// TODO(b/178035553): Disable NV12 due to FrameBuffer issue.
// case NV12:
case NV21:
case YV12:
case YV21:
break;
default:
throw new IllegalArgumentException(
Expand Down
31 changes: 25 additions & 6 deletions ...w_lite_support/java/src/java/org/tensorflow/lite/task/vision/detector/ObjectDetector.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,8 +118,8 @@ public static ObjectDetector createFromFile(File modelFile) throws IOException {
* Creates an {@link ObjectDetector} instance with a model buffer and the default {@link
* ObjectDetectorOptions}.
*
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the
* classification model
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the detection
* model
* @throws AssertionError if error occurs when creating {@link ObjectDetector} from the native
* code
* @throws IllegalArgumentException if the model buffer is not a direct {@link ByteBuffer} or a
Expand Down Expand Up @@ -190,8 +190,8 @@ public long createHandle() {
* Creates an {@link ObjectDetector} instance with a model buffer and {@link
* ObjectDetectorOptions}.
*
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the
* classification model
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the detection
* model
* @throws AssertionError if error occurs when creating {@link ObjectDetector} from the native
* code
* @throws IllegalArgumentException if the model buffer is not a direct {@link ByteBuffer} or a
Expand Down Expand Up @@ -394,8 +394,18 @@ private ObjectDetectorOptions(Builder builder) {
/**
* Performs actual detection on the provided image.
*
* @param image a {@link TensorImage} object that represents a RGB image
* <p>{@link ObjectDetector} supports the following {@link TensorImage} color space types:
*
* <ul>
* <li>{@link ColorSpaceType#RGB}
* <li>{@link ColorSpaceType#NV21}
* <li>{@link ColorSpaceType#YV12}
* <li>{@link ColorSpaceType#YV21}
* </ul>
*
* @param image a UINT8 {@link TensorImage} object that represents an RGB or YUV image
* @throws AssertionError if error occurs when processing the image from the native code
* @throws IllegalArgumentException if the color space type of image is unsupported
*/
public List<Detection> detect(TensorImage image) {
return detect(image, ImageProcessingOptions.builder().build());
Expand All @@ -404,7 +414,16 @@ public List<Detection> detect(TensorImage image) {
/**
* Performs actual detection on the provided image.
*
* @param image a {@link TensorImage} object that represents a RGB image
* <p>{@link ObjectDetector} supports the following {@link TensorImage} color space types:
*
* <ul>
* <li>{@link ColorSpaceType#RGB}
* <li>{@link ColorSpaceType#NV21}
* <li>{@link ColorSpaceType#YV12}
* <li>{@link ColorSpaceType#YV21}
* </ul>
*
* @param image a UINT8 {@link TensorImage} object that represents an RGB or YUV image
* @param options {@link ObjectDetector} only supports image rotation (through {@link
* ImageProcessingOptions#Builder#setOrientation}) currently. The orientation of an image
* defaults to {@link ImageProcessingOptions#Orientation#TOP_LEFT}.
Expand Down
27 changes: 23 additions & 4 deletions ..._lite_support/java/src/java/org/tensorflow/lite/task/vision/segmenter/ImageSegmenter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,7 +109,7 @@ public static ImageSegmenter createFromFile(File modelFile) throws IOException {
* ImageSegmenterOptions}.
*
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the
* classification model
* segmentation model
* @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
* code
* @throws IllegalArgumentException if the model buffer is not a direct {@link ByteBuffer} or a
Expand Down Expand Up @@ -163,7 +163,7 @@ public static ImageSegmenter createFromFileAndOptions(
* ImageSegmenterOptions}.
*
* @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the
* classification model
* segmentation model
* @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
* code
* @throws IllegalArgumentException if the model buffer is not a direct {@link ByteBuffer} or a
Expand Down Expand Up @@ -253,12 +253,22 @@ public abstract static class Builder {
/**
* Performs actual segmentation on the provided image.
*
* @param image a {@link TensorImage} object that represents an RGB image
* <p>{@link ImageSegmenter} supports the following {@link TensorImage} color space types:
*
* <ul>
* <li>{@link ColorSpaceType#RGB}
* <li>{@link ColorSpaceType#NV21}
* <li>{@link ColorSpaceType#YV12}
* <li>{@link ColorSpaceType#YV21}
* </ul>
*
* @param image a UINT8 {@link TensorImage} object that represents an RGB or YUV image
* @return results of performing image segmentation. Note that at the time, a single {@link
* Segmentation} element is expected to be returned. The result is stored in a {@link List}
* for later extension to e.g. instance segmentation models, which may return one segmentation
* per object.
* @throws AssertionError if error occurs when segmenting the image from the native code
* @throws IllegalArgumentException if the color space type of image is unsupported
*/
public List<Segmentation> segment(TensorImage image) {
return segment(image, ImageProcessingOptions.builder().build());
Expand All @@ -267,7 +277,16 @@ public List<Segmentation> segment(TensorImage image) {
/**
* Performs actual segmentation on the provided image with {@link ImageProcessingOptions}.
*
* @param image a {@link TensorImage} object that represents an RGB image
* <p>{@link ImageSegmenter} supports the following {@link TensorImage} color space types:
*
* <ul>
* <li>{@link ColorSpaceType#RGB}
* <li>{@link ColorSpaceType#NV21}
* <li>{@link ColorSpaceType#YV12}
* <li>{@link ColorSpaceType#YV21}
* </ul>
*
* @param image a UINT8 {@link TensorImage} object that represents an RGB or YUV image
* @param options {@link ImageSegmenter} only supports image rotation (through {@link
* ImageProcessingOptions#Builder#setOrientation}) currently. The orientation of an image
* defaults to {@link ImageProcessingOptions#Orientation#TOP_LEFT}.
Expand Down

0 comments on commit efac3d1

Please sign in to comment.