update ImageSegmenterGraph to always output confidence mask and optionally output category mask
PiperOrigin-RevId: 521679910
This commit is contained in:
MediaPipe Team
committed by
Copybara-Service
Copybara-Service
parent
c31a4681e5
commit
367ccbfdf3
@@ -26,6 +26,7 @@ limitations under the License.
|
||||
#include "mediapipe/tasks/cc/core/base_options.h"
|
||||
#include "mediapipe/tasks/cc/vision/core/base_vision_task_api.h"
|
||||
#include "mediapipe/tasks/cc/vision/core/image_processing_options.h"
|
||||
#include "mediapipe/tasks/cc/vision/image_segmenter/image_segmenter_result.h"
|
||||
#include "tensorflow/lite/kernels/register.h"
|
||||
|
||||
namespace mediapipe {
|
||||
@@ -52,23 +53,14 @@ struct ImageSegmenterOptions {
|
||||
// Metadata, if any. Defaults to English.
|
||||
std::string display_names_locale = "en";
|
||||
|
||||
// The output type of segmentation results.
|
||||
enum OutputType {
|
||||
// Gives a single output mask where each pixel represents the class which
|
||||
// the pixel in the original image was predicted to belong to.
|
||||
CATEGORY_MASK = 0,
|
||||
// Gives a list of output masks where, for each mask, each pixel represents
|
||||
// the prediction confidence, usually in the [0, 1] range.
|
||||
CONFIDENCE_MASK = 1,
|
||||
};
|
||||
|
||||
OutputType output_type = OutputType::CATEGORY_MASK;
|
||||
// Whether to output category mask.
|
||||
bool output_category_mask = false;
|
||||
|
||||
// The user-defined result callback for processing live stream data.
|
||||
// The result callback should only be specified when the running mode is set
|
||||
// to RunningMode::LIVE_STREAM.
|
||||
std::function<void(absl::StatusOr<std::vector<mediapipe::Image>>,
|
||||
const Image&, int64)>
|
||||
std::function<void(absl::StatusOr<ImageSegmenterResult>, const Image&,
|
||||
int64_t)>
|
||||
result_callback = nullptr;
|
||||
};
|
||||
|
||||
@@ -84,13 +76,9 @@ struct ImageSegmenterOptions {
|
||||
// 1 or 3).
|
||||
// - if type is kTfLiteFloat32, NormalizationOptions are required to be
|
||||
// attached to the metadata for input normalization.
|
||||
// Output tensors:
|
||||
// (kTfLiteUInt8/kTfLiteFloat32)
|
||||
// - list of segmented masks.
|
||||
// - if `output_type` is CATEGORY_MASK, uint8 Image, Image vector of size 1.
|
||||
// - if `output_type` is CONFIDENCE_MASK, float32 Image list of size
|
||||
// `channels`.
|
||||
// - batch is always 1
|
||||
// Output ImageSegmenterResult:
|
||||
// Provides confidence masks and an optional category mask if
|
||||
// `output_category_mask` is set true.
|
||||
// An example of such model can be found at:
|
||||
// https://tfhub.dev/tensorflow/lite-model/deeplabv3/1/metadata/2
|
||||
class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
@@ -114,12 +102,8 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
// setting its 'rotation_degrees' field. Note that specifying a
|
||||
// region-of-interest using the 'region_of_interest' field is NOT supported
|
||||
// and will result in an invalid argument error being returned.
|
||||
//
|
||||
// If the output_type is CATEGORY_MASK, the returned vector of images is
|
||||
// per-category segmented image mask.
|
||||
// If the output_type is CONFIDENCE_MASK, the returned vector of images
|
||||
// contains only one confidence image mask.
|
||||
absl::StatusOr<std::vector<mediapipe::Image>> Segment(
|
||||
|
||||
absl::StatusOr<ImageSegmenterResult> Segment(
|
||||
mediapipe::Image image,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options =
|
||||
std::nullopt);
|
||||
@@ -137,13 +121,8 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
// setting its 'rotation_degrees' field. Note that specifying a
|
||||
// region-of-interest using the 'region_of_interest' field is NOT supported
|
||||
// and will result in an invalid argument error being returned.
|
||||
//
|
||||
// If the output_type is CATEGORY_MASK, the returned vector of images is
|
||||
// per-category segmented image mask.
|
||||
// If the output_type is CONFIDENCE_MASK, the returned vector of images
|
||||
// contains only one confidence image mask.
|
||||
absl::StatusOr<std::vector<mediapipe::Image>> SegmentForVideo(
|
||||
mediapipe::Image image, int64 timestamp_ms,
|
||||
absl::StatusOr<ImageSegmenterResult> SegmentForVideo(
|
||||
mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions> image_processing_options =
|
||||
std::nullopt);
|
||||
|
||||
@@ -164,17 +143,13 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
// and will result in an invalid argument error being returned.
|
||||
//
|
||||
// The "result_callback" prvoides
|
||||
// - A vector of segmented image masks.
|
||||
// If the output_type is CATEGORY_MASK, the returned vector of images is
|
||||
// per-category segmented image mask.
|
||||
// If the output_type is CONFIDENCE_MASK, the returned vector of images
|
||||
// contains only one confidence image mask.
|
||||
// - An ImageSegmenterResult.
|
||||
// - The const reference to the corresponding input image that the image
|
||||
// segmentation runs on. Note that the const reference to the image will
|
||||
// no longer be valid when the callback returns. To access the image data
|
||||
// outside of the callback, callers need to make a copy of the image.
|
||||
// - The input timestamp in milliseconds.
|
||||
absl::Status SegmentAsync(mediapipe::Image image, int64 timestamp_ms,
|
||||
absl::Status SegmentAsync(mediapipe::Image image, int64_t timestamp_ms,
|
||||
std::optional<core::ImageProcessingOptions>
|
||||
image_processing_options = std::nullopt);
|
||||
|
||||
@@ -182,9 +157,9 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
absl::Status Close() { return runner_->Close(); }
|
||||
|
||||
// Get the category label list of the ImageSegmenter can recognize. For
|
||||
// CATEGORY_MASK type, the index in the category mask corresponds to the
|
||||
// category in the label list. For CONFIDENCE_MASK type, the output mask list
|
||||
// at index corresponds to the category in the label list.
|
||||
// CATEGORY_MASK, the index in the category mask corresponds to the category
|
||||
// in the label list. For CONFIDENCE_MASK, the output mask list at index
|
||||
// corresponds to the category in the label list.
|
||||
//
|
||||
// If there is no labelmap provided in the model file, empty label list is
|
||||
// returned.
|
||||
@@ -192,6 +167,7 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi {
|
||||
|
||||
private:
|
||||
std::vector<std::string> labels_;
|
||||
bool output_category_mask_;
|
||||
};
|
||||
|
||||
} // namespace image_segmenter
|
||||
|
||||
Reference in New Issue
Block a user