422 lines
18 KiB
C++
422 lines
18 KiB
C++
/* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
#include <memory>
|
|
#include <type_traits>
|
|
#include <vector>
|
|
|
|
#include "absl/status/status.h"
|
|
#include "absl/status/statusor.h"
|
|
#include "absl/strings/str_format.h"
|
|
#include "mediapipe/calculators/image/image_clone_calculator.pb.h"
|
|
#include "mediapipe/calculators/image/image_transformation_calculator.pb.h"
|
|
#include "mediapipe/calculators/tensor/tensor_converter_calculator.pb.h"
|
|
#include "mediapipe/framework/api2/builder.h"
|
|
#include "mediapipe/framework/api2/port.h"
|
|
#include "mediapipe/framework/formats/image.h"
|
|
#include "mediapipe/framework/formats/rect.pb.h"
|
|
#include "mediapipe/framework/formats/tensor.h"
|
|
#include "mediapipe/framework/port/status_macros.h"
|
|
#include "mediapipe/tasks/cc/common.h"
|
|
#include "mediapipe/tasks/cc/components/processors/image_preprocessing_graph.h"
|
|
#include "mediapipe/tasks/cc/components/processors/proto/image_preprocessing_graph_options.pb.h"
|
|
#include "mediapipe/tasks/cc/core/model_resources.h"
|
|
#include "mediapipe/tasks/cc/core/model_task_graph.h"
|
|
#include "mediapipe/tasks/cc/core/proto/inference_subgraph.pb.h"
|
|
#include "mediapipe/tasks/cc/metadata/metadata_extractor.h"
|
|
#include "mediapipe/tasks/cc/vision/image_segmenter/calculators/tensors_to_segmentation_calculator.pb.h"
|
|
#include "mediapipe/tasks/cc/vision/image_segmenter/proto/image_segmenter_graph_options.pb.h"
|
|
#include "mediapipe/tasks/cc/vision/image_segmenter/proto/segmenter_options.pb.h"
|
|
#include "mediapipe/tasks/cc/vision/utils/image_tensor_specs.h"
|
|
#include "mediapipe/tasks/metadata/metadata_schema_generated.h"
|
|
#include "mediapipe/util/label_map.pb.h"
|
|
#include "mediapipe/util/label_map_util.h"
|
|
#include "tensorflow/lite/schema/schema_generated.h"
|
|
|
|
namespace mediapipe {
|
|
namespace tasks {
|
|
namespace vision {
|
|
namespace image_segmenter {
|
|
|
|
namespace {
|
|
|
|
using ::mediapipe::Image;
|
|
using ::mediapipe::NormalizedRect;
|
|
using ::mediapipe::api2::Input;
|
|
using ::mediapipe::api2::Output;
|
|
using ::mediapipe::api2::builder::Graph;
|
|
using ::mediapipe::api2::builder::MultiSource;
|
|
using ::mediapipe::api2::builder::Source;
|
|
using ::mediapipe::tasks::metadata::ModelMetadataExtractor;
|
|
using ::mediapipe::tasks::vision::image_segmenter::proto::
|
|
ImageSegmenterGraphOptions;
|
|
using ::mediapipe::tasks::vision::image_segmenter::proto::SegmenterOptions;
|
|
using ::tflite::TensorMetadata;
|
|
using LabelItems = mediapipe::proto_ns::Map<int64, ::mediapipe::LabelMapItem>;
|
|
|
|
constexpr char kSegmentationTag[] = "SEGMENTATION";
|
|
constexpr char kGroupedSegmentationTag[] = "GROUPED_SEGMENTATION";
|
|
constexpr char kImageTag[] = "IMAGE";
|
|
constexpr char kImageCpuTag[] = "IMAGE_CPU";
|
|
constexpr char kImageGpuTag[] = "IMAGE_GPU";
|
|
constexpr char kNormRectTag[] = "NORM_RECT";
|
|
constexpr char kTensorsTag[] = "TENSORS";
|
|
constexpr char kOutputSizeTag[] = "OUTPUT_SIZE";
|
|
|
|
// Struct holding the different output streams produced by the image segmenter
|
|
// subgraph.
|
|
struct ImageSegmenterOutputs {
|
|
std::vector<Source<Image>> segmented_masks;
|
|
// The same as the input image, mainly used for live stream mode.
|
|
Source<Image> image;
|
|
};
|
|
|
|
// Struct holding the image and input tensors after image preprocessing and
|
|
// transferred to the requested device.
|
|
struct ImageAndTensorsOnDevice {
|
|
Source<Image> image;
|
|
Source<std::vector<Tensor>> tensors;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
absl::Status SanityCheckOptions(const ImageSegmenterGraphOptions& options) {
|
|
if (options.segmenter_options().output_type() ==
|
|
SegmenterOptions::UNSPECIFIED) {
|
|
return CreateStatusWithPayload(absl::StatusCode::kInvalidArgument,
|
|
"`output_type` must not be UNSPECIFIED",
|
|
MediaPipeTasksStatus::kInvalidArgumentError);
|
|
}
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
absl::StatusOr<LabelItems> GetLabelItemsIfAny(
|
|
const ModelMetadataExtractor& metadata_extractor,
|
|
const TensorMetadata& tensor_metadata, absl::string_view locale) {
|
|
const std::string labels_filename =
|
|
ModelMetadataExtractor::FindFirstAssociatedFileName(
|
|
tensor_metadata, tflite::AssociatedFileType_TENSOR_AXIS_LABELS);
|
|
if (labels_filename.empty()) {
|
|
LabelItems empty_label_items;
|
|
return empty_label_items;
|
|
}
|
|
ASSIGN_OR_RETURN(absl::string_view labels_file,
|
|
metadata_extractor.GetAssociatedFile(labels_filename));
|
|
const std::string display_names_filename =
|
|
ModelMetadataExtractor::FindFirstAssociatedFileName(
|
|
tensor_metadata, tflite::AssociatedFileType_TENSOR_AXIS_LABELS,
|
|
locale);
|
|
absl::string_view display_names_file;
|
|
if (!display_names_filename.empty()) {
|
|
ASSIGN_OR_RETURN(display_names_file, metadata_extractor.GetAssociatedFile(
|
|
display_names_filename));
|
|
}
|
|
return mediapipe::BuildLabelMapFromFiles(labels_file, display_names_file);
|
|
}
|
|
|
|
absl::Status ConfigureTensorsToSegmentationCalculator(
|
|
const ImageSegmenterGraphOptions& segmenter_option,
|
|
const core::ModelResources& model_resources,
|
|
TensorsToSegmentationCalculatorOptions* options) {
|
|
*options->mutable_segmenter_options() = segmenter_option.segmenter_options();
|
|
const tflite::Model& model = *model_resources.GetTfLiteModel();
|
|
if (model.subgraphs()->size() != 1) {
|
|
return CreateStatusWithPayload(
|
|
absl::StatusCode::kInvalidArgument,
|
|
"Segmentation tflite models are assumed to have a single subgraph.",
|
|
MediaPipeTasksStatus::kInvalidArgumentError);
|
|
}
|
|
const auto* primary_subgraph = (*model.subgraphs())[0];
|
|
if (primary_subgraph->outputs()->size() != 1) {
|
|
return CreateStatusWithPayload(
|
|
absl::StatusCode::kInvalidArgument,
|
|
"Segmentation tflite models are assumed to have a single output.",
|
|
MediaPipeTasksStatus::kInvalidArgumentError);
|
|
}
|
|
|
|
const ModelMetadataExtractor* metadata_extractor =
|
|
model_resources.GetMetadataExtractor();
|
|
ASSIGN_OR_RETURN(
|
|
*options->mutable_label_items(),
|
|
GetLabelItemsIfAny(*metadata_extractor,
|
|
*metadata_extractor->GetOutputTensorMetadata()->Get(0),
|
|
segmenter_option.display_names_locale()));
|
|
return absl::OkStatus();
|
|
}
|
|
|
|
// Get the output tensor from the tflite model of given model resources.
|
|
absl::StatusOr<const tflite::Tensor*> GetOutputTensor(
|
|
const core::ModelResources& model_resources) {
|
|
const tflite::Model& model = *model_resources.GetTfLiteModel();
|
|
const auto* primary_subgraph = (*model.subgraphs())[0];
|
|
const auto* output_tensor =
|
|
(*primary_subgraph->tensors())[(*primary_subgraph->outputs())[0]];
|
|
return output_tensor;
|
|
}
|
|
|
|
// Get the input tensor from the tflite model of given model resources.
|
|
absl::StatusOr<const tflite::Tensor*> GetInputTensor(
|
|
const core::ModelResources& model_resources) {
|
|
const tflite::Model& model = *model_resources.GetTfLiteModel();
|
|
const auto* primary_subgraph = (*model.subgraphs())[0];
|
|
const auto* input_tensor =
|
|
(*primary_subgraph->tensors())[(*primary_subgraph->inputs())[0]];
|
|
return input_tensor;
|
|
}
|
|
|
|
// Configure the ImageTransformationCalculator according to the input tensor.
|
|
void ConfigureImageTransformationCalculator(
|
|
const tflite::Tensor& tflite_input_tensor,
|
|
mediapipe::ImageTransformationCalculatorOptions& options) {
|
|
options.set_output_height(tflite_input_tensor.shape()->data()[1]);
|
|
options.set_output_width(tflite_input_tensor.shape()->data()[2]);
|
|
}
|
|
|
|
// Configure the TensorConverterCalculator to convert the image to tensor.
|
|
void ConfigureTensorConverterCalculator(
|
|
const ImageTensorSpecs& image_tensor_specs,
|
|
mediapipe::TensorConverterCalculatorOptions& options) {
|
|
float mean = image_tensor_specs.normalization_options->mean_values[0];
|
|
float std = image_tensor_specs.normalization_options->std_values[0];
|
|
options.set_max_num_channels(4);
|
|
options.mutable_output_tensor_float_range()->set_min((0.0f - mean) / std);
|
|
options.mutable_output_tensor_float_range()->set_max((255.0f - mean) / std);
|
|
}
|
|
|
|
// Image preprocessing step to convert the given image to the input tensors for
|
|
// the tflite model.
|
|
absl::StatusOr<ImageAndTensorsOnDevice> ConvertImageToTensors(
|
|
Source<Image> image_in, Source<NormalizedRect> norm_rect_in, bool use_gpu,
|
|
const core::ModelResources& model_resources, Graph& graph) {
|
|
ASSIGN_OR_RETURN(const tflite::Tensor* tflite_input_tensor,
|
|
GetInputTensor(model_resources));
|
|
if (tflite_input_tensor->shape()->size() != 4) {
|
|
return absl::InvalidArgumentError(
|
|
absl::StrFormat("Expect segmentation model has input image tensor to "
|
|
"be 4 dims. Got input tensor with "
|
|
"dims: %d",
|
|
tflite_input_tensor->shape()->size()));
|
|
}
|
|
const int input_tensor_channel = tflite_input_tensor->shape()->data()[3];
|
|
if (input_tensor_channel != 3 && input_tensor_channel != 4) {
|
|
return absl::InvalidArgumentError(absl::StrFormat(
|
|
"Expect segmentation model has input image tensor with channels = 3 or "
|
|
"4. Get "
|
|
"channel = %d",
|
|
tflite_input_tensor->shape()->data()[3]));
|
|
} else if (input_tensor_channel == 3) {
|
|
// ImagePreprocessingGraph is backed by ImageToTensorCalculator which only
|
|
// supports Tensor with channel = 3.
|
|
auto& preprocessing = graph.AddNode(
|
|
"mediapipe.tasks.components.processors.ImagePreprocessingGraph");
|
|
MP_RETURN_IF_ERROR(components::processors::ConfigureImagePreprocessingGraph(
|
|
model_resources, use_gpu,
|
|
&preprocessing.GetOptions<tasks::components::processors::proto::
|
|
ImagePreprocessingGraphOptions>()));
|
|
image_in >> preprocessing.In(kImageTag);
|
|
norm_rect_in >> preprocessing.In(kNormRectTag);
|
|
return {{preprocessing.Out(kImageTag).Cast<Image>(),
|
|
preprocessing.Out(kTensorsTag).Cast<std::vector<Tensor>>()}};
|
|
} else {
|
|
// TODO Remove legacy preprocessing calculators.
|
|
// For segmentation model with input Tensor with channel = 4, use legacy
|
|
// TfLite preprocessing calculators
|
|
|
|
// Upload image to GPU if requested to use gpu.
|
|
auto& image_clone = graph.AddNode("ImageCloneCalculator");
|
|
image_clone.GetOptions<mediapipe::ImageCloneCalculatorOptions>()
|
|
.set_output_on_gpu(use_gpu);
|
|
image_in >> image_clone.In("");
|
|
Source<Image> image_on_device = image_clone.Out("").Cast<Image>();
|
|
|
|
// Convert from Image to legacy ImageFrame or GpuBuffer.
|
|
auto& from_image = graph.AddNode("FromImageCalculator");
|
|
image_on_device >> from_image.In(kImageTag);
|
|
auto image_cpu_or_gpu =
|
|
from_image.Out(use_gpu ? kImageGpuTag : kImageCpuTag);
|
|
|
|
// Resize the input image to the model input size.
|
|
auto& image_transformation = graph.AddNode("ImageTransformationCalculator");
|
|
ConfigureImageTransformationCalculator(
|
|
*tflite_input_tensor,
|
|
image_transformation
|
|
.GetOptions<mediapipe::ImageTransformationCalculatorOptions>());
|
|
const absl::string_view image_or_image_gpu_tag =
|
|
use_gpu ? kImageGpuTag : kImageTag;
|
|
image_cpu_or_gpu >> image_transformation.In(image_or_image_gpu_tag);
|
|
auto transformed_image = image_transformation.Out(image_or_image_gpu_tag);
|
|
|
|
// Convert image to mediapipe tensor.
|
|
auto& tensor_converter = graph.AddNode("TensorConverterCalculator");
|
|
ASSIGN_OR_RETURN(auto image_tensor_specs,
|
|
vision::BuildInputImageTensorSpecs(model_resources));
|
|
ConfigureTensorConverterCalculator(
|
|
image_tensor_specs,
|
|
tensor_converter
|
|
.GetOptions<mediapipe::TensorConverterCalculatorOptions>());
|
|
|
|
transformed_image >> tensor_converter.In(image_or_image_gpu_tag);
|
|
auto tensors =
|
|
tensor_converter.Out(kTensorsTag).Cast<std::vector<Tensor>>();
|
|
|
|
return {{image_on_device, tensors}};
|
|
}
|
|
}
|
|
|
|
// An "mediapipe.tasks.vision.ImageSegmenterGraph" performs semantic
|
|
// segmentation.
|
|
// Two kinds of outputs are provided: SEGMENTATION and GROUPED_SEGMENTATION.
|
|
// Users can retrieve segmented mask of only particular category/channel from
|
|
// SEGMENTATION, and users can also get all segmented masks from
|
|
// GROUPED_SEGMENTATION.
|
|
// - Accepts CPU input images and outputs segmented masks on CPU.
|
|
//
|
|
// Inputs:
|
|
// IMAGE - Image
|
|
// Image to perform segmentation on.
|
|
// NORM_RECT - NormalizedRect @Optional
|
|
// Describes image rotation and region of image to perform detection
|
|
// on.
|
|
// @Optional: rect covering the whole image is used if not specified.
|
|
//
|
|
// Outputs:
|
|
// SEGMENTATION - mediapipe::Image @Multiple
|
|
// Segmented masks for individual category. Segmented mask of single
|
|
// category can be accessed by index based output stream.
|
|
// GROUPED_SEGMENTATION - std::vector<mediapipe::Image>
|
|
// The output segmented masks grouped in a vector.
|
|
// IMAGE - mediapipe::Image
|
|
// The image that image segmenter runs on.
|
|
//
|
|
// Example:
|
|
// node {
|
|
// calculator: "mediapipe.tasks.vision.ImageSegmenterGraph"
|
|
// input_stream: "IMAGE:image"
|
|
// output_stream: "SEGMENTATION:segmented_masks"
|
|
// options {
|
|
// [mediapipe.tasks.vision.image_segmenter.proto.ImageSegmenterGraphOptions.ext]
|
|
// {
|
|
// base_options {
|
|
// model_asset {
|
|
// file_name: "/path/to/model.tflite"
|
|
// }
|
|
// }
|
|
// segmenter_options {
|
|
// output_type: CONFIDENCE_MASK
|
|
// activation: SOFTMAX
|
|
// }
|
|
// }
|
|
// }
|
|
// }
|
|
class ImageSegmenterGraph : public core::ModelTaskGraph {
|
|
public:
|
|
absl::StatusOr<mediapipe::CalculatorGraphConfig> GetConfig(
|
|
mediapipe::SubgraphContext* sc) override {
|
|
ASSIGN_OR_RETURN(const auto* model_resources,
|
|
CreateModelResources<ImageSegmenterGraphOptions>(sc));
|
|
Graph graph;
|
|
ASSIGN_OR_RETURN(
|
|
auto output_streams,
|
|
BuildSegmentationTask(
|
|
sc->Options<ImageSegmenterGraphOptions>(), *model_resources,
|
|
graph[Input<Image>(kImageTag)],
|
|
graph[Input<NormalizedRect>::Optional(kNormRectTag)], graph));
|
|
|
|
auto& merge_images_to_vector =
|
|
graph.AddNode("MergeImagesToVectorCalculator");
|
|
for (int i = 0; i < output_streams.segmented_masks.size(); ++i) {
|
|
output_streams.segmented_masks[i] >>
|
|
merge_images_to_vector[Input<Image>::Multiple("")][i];
|
|
output_streams.segmented_masks[i] >>
|
|
graph[Output<Image>::Multiple(kSegmentationTag)][i];
|
|
}
|
|
merge_images_to_vector.Out("") >>
|
|
graph[Output<std::vector<Image>>(kGroupedSegmentationTag)];
|
|
output_streams.image >> graph[Output<Image>(kImageTag)];
|
|
return graph.GetConfig();
|
|
}
|
|
|
|
private:
|
|
// Adds a mediapipe image segmentation task pipeline graph into the provided
|
|
// builder::Graph instance. The segmentation pipeline takes images
|
|
// (mediapipe::Image) as the input and returns segmented image mask as output.
|
|
//
|
|
// task_options: the mediapipe tasks ImageSegmenterGraphOptions proto.
|
|
// model_resources: the ModelSources object initialized from a segmentation
|
|
// model file with model metadata.
|
|
// image_in: (mediapipe::Image) stream to run segmentation on.
|
|
// graph: the mediapipe builder::Graph instance to be updated.
|
|
absl::StatusOr<ImageSegmenterOutputs> BuildSegmentationTask(
|
|
const ImageSegmenterGraphOptions& task_options,
|
|
const core::ModelResources& model_resources, Source<Image> image_in,
|
|
Source<NormalizedRect> norm_rect_in, Graph& graph) {
|
|
MP_RETURN_IF_ERROR(SanityCheckOptions(task_options));
|
|
|
|
// Adds preprocessing calculators and connects them to the graph input image
|
|
// stream.
|
|
bool use_gpu =
|
|
components::processors::DetermineImagePreprocessingGpuBackend(
|
|
task_options.base_options().acceleration());
|
|
ASSIGN_OR_RETURN(auto image_and_tensors,
|
|
ConvertImageToTensors(image_in, norm_rect_in, use_gpu,
|
|
model_resources, graph));
|
|
// Adds inference subgraph and connects its input stream to the output
|
|
// tensors produced by the ImageToTensorCalculator.
|
|
auto& inference = AddInference(
|
|
model_resources, task_options.base_options().acceleration(), graph);
|
|
image_and_tensors.tensors >> inference.In(kTensorsTag);
|
|
|
|
// Adds segmentation calculators for output streams.
|
|
auto& tensor_to_images =
|
|
graph.AddNode("mediapipe.tasks.TensorsToSegmentationCalculator");
|
|
RET_CHECK_OK(ConfigureTensorsToSegmentationCalculator(
|
|
task_options, model_resources,
|
|
&tensor_to_images
|
|
.GetOptions<TensorsToSegmentationCalculatorOptions>()));
|
|
inference.Out(kTensorsTag) >> tensor_to_images.In(kTensorsTag);
|
|
|
|
// Adds image property calculator for output size.
|
|
auto& image_properties = graph.AddNode("ImagePropertiesCalculator");
|
|
image_in >> image_properties.In("IMAGE");
|
|
image_properties.Out("SIZE") >> tensor_to_images.In(kOutputSizeTag);
|
|
|
|
// Exports multiple segmented masks.
|
|
std::vector<Source<Image>> segmented_masks;
|
|
if (task_options.segmenter_options().output_type() ==
|
|
SegmenterOptions::CATEGORY_MASK) {
|
|
segmented_masks.push_back(
|
|
Source<Image>(tensor_to_images[Output<Image>(kSegmentationTag)]));
|
|
} else {
|
|
ASSIGN_OR_RETURN(const tflite::Tensor* output_tensor,
|
|
GetOutputTensor(model_resources));
|
|
const int segmentation_streams_num = *output_tensor->shape()->rbegin();
|
|
for (int i = 0; i < segmentation_streams_num; ++i) {
|
|
segmented_masks.push_back(Source<Image>(
|
|
tensor_to_images[Output<Image>::Multiple(kSegmentationTag)][i]));
|
|
}
|
|
}
|
|
return ImageSegmenterOutputs{/*segmented_masks=*/segmented_masks,
|
|
/*image=*/image_and_tensors.image};
|
|
}
|
|
};
|
|
|
|
REGISTER_MEDIAPIPE_GRAPH(
|
|
::mediapipe::tasks::vision::image_segmenter::ImageSegmenterGraph);
|
|
|
|
} // namespace image_segmenter
|
|
} // namespace vision
|
|
} // namespace tasks
|
|
} // namespace mediapipe
|