Image segmenter output both confidence masks and category mask optionally.
PiperOrigin-RevId: 522227345
This commit is contained in:
		
							parent
							
								
									7fe87936e5
								
							
						
					
					
						commit
						d5def9e24d
					
				|  | @ -32,6 +32,7 @@ limitations under the License. | |||
| #include "mediapipe/framework/formats/image.h" | ||||
| #include "mediapipe/framework/formats/image_frame_opencv.h" | ||||
| #include "mediapipe/framework/formats/tensor.h" | ||||
| #include "mediapipe/framework/port/canonical_errors.h" | ||||
| #include "mediapipe/framework/port/opencv_core_inc.h" | ||||
| #include "mediapipe/framework/port/opencv_imgproc_inc.h" | ||||
| #include "mediapipe/framework/port/status_macros.h" | ||||
|  | @ -210,8 +211,9 @@ std::vector<Image> ProcessForConfidenceMaskCpu(const Shape& input_shape, | |||
| }  // namespace
 | ||||
| 
 | ||||
| // Converts Tensors from a vector of Tensor to Segmentation masks. The
 | ||||
| // calculator always output confidence masks, and an optional category mask if
 | ||||
| // CATEGORY_MASK is connected.
 | ||||
| // calculator can output optional confidence masks if CONFIDENCE_MASK is
 | ||||
| // connected, and an optional category mask if CATEGORY_MASK is connected. At
 | ||||
| // least one of CONFIDENCE_MASK and CATEGORY_MASK must be connected.
 | ||||
| //
 | ||||
| // Performs optional resizing to OUTPUT_SIZE dimension if provided,
 | ||||
| // otherwise the segmented masks is the same size as input tensor.
 | ||||
|  | @ -296,6 +298,13 @@ absl::Status TensorsToSegmentationCalculator::Open( | |||
|                  SegmenterOptions::UNSPECIFIED) | ||||
|         << "Must specify output_type as one of " | ||||
|            "[CONFIDENCE_MASK|CATEGORY_MASK]."; | ||||
|   } else { | ||||
|     if (!cc->Outputs().HasTag("CONFIDENCE_MASK") && | ||||
|         !cc->Outputs().HasTag("CATEGORY_MASK")) { | ||||
|       return absl::InvalidArgumentError( | ||||
|           "At least one of CONFIDENCE_MASK and CATEGORY_MASK must be " | ||||
|           "connected."); | ||||
|     } | ||||
|   } | ||||
| #ifdef __EMSCRIPTEN__ | ||||
|   MP_RETURN_IF_ERROR(postprocessor_.Initialize(cc, options_)); | ||||
|  | @ -366,8 +375,9 @@ absl::Status TensorsToSegmentationCalculator::Process( | |||
|     return absl::OkStatus(); | ||||
|   } | ||||
| 
 | ||||
|   std::vector<Image> confidence_masks = | ||||
|       ProcessForConfidenceMaskCpu(input_shape, | ||||
|   if (cc->Outputs().HasTag("CONFIDENCE_MASK")) { | ||||
|     std::vector<Image> confidence_masks = ProcessForConfidenceMaskCpu( | ||||
|         input_shape, | ||||
|         {/* height= */ output_height, | ||||
|          /* width= */ output_width, | ||||
|          /* channels= */ input_shape.channels}, | ||||
|  | @ -375,6 +385,7 @@ absl::Status TensorsToSegmentationCalculator::Process( | |||
|     for (int i = 0; i < confidence_masks.size(); ++i) { | ||||
|       kConfidenceMaskOut(cc)[i].Send(std::move(confidence_masks[i])); | ||||
|     } | ||||
|   } | ||||
|   if (cc->Outputs().HasTag("CATEGORY_MASK")) { | ||||
|     kCategoryMaskOut(cc).Send(ProcessForCategoryMaskCpu( | ||||
|         input_shape, | ||||
|  |  | |||
|  | @ -60,15 +60,19 @@ using ImageSegmenterGraphOptionsProto = ::mediapipe::tasks::vision:: | |||
| // "mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph".
 | ||||
| CalculatorGraphConfig CreateGraphConfig( | ||||
|     std::unique_ptr<ImageSegmenterGraphOptionsProto> options, | ||||
|     bool output_category_mask, bool enable_flow_limiting) { | ||||
|     bool output_confidence_masks, bool output_category_mask, | ||||
|     bool enable_flow_limiting) { | ||||
|   api2::builder::Graph graph; | ||||
|   auto& task_subgraph = graph.AddNode(kSubgraphTypeName); | ||||
|   task_subgraph.GetOptions<ImageSegmenterGraphOptionsProto>().Swap( | ||||
|       options.get()); | ||||
|   graph.In(kImageTag).SetName(kImageInStreamName); | ||||
|   graph.In(kNormRectTag).SetName(kNormRectStreamName); | ||||
|   task_subgraph.Out(kConfidenceMasksTag).SetName(kConfidenceMasksStreamName) >> | ||||
|   if (output_confidence_masks) { | ||||
|     task_subgraph.Out(kConfidenceMasksTag) | ||||
|             .SetName(kConfidenceMasksStreamName) >> | ||||
|         graph.Out(kConfidenceMasksTag); | ||||
|   } | ||||
|   if (output_category_mask) { | ||||
|     task_subgraph.Out(kCategoryMaskTag).SetName(kCategoryMaskStreamName) >> | ||||
|         graph.Out(kCategoryMaskTag); | ||||
|  | @ -135,11 +139,17 @@ absl::StatusOr<std::vector<std::string>> GetLabelsFromGraphConfig( | |||
| 
 | ||||
| absl::StatusOr<std::unique_ptr<ImageSegmenter>> ImageSegmenter::Create( | ||||
|     std::unique_ptr<ImageSegmenterOptions> options) { | ||||
|   if (!options->output_confidence_masks && !options->output_category_mask) { | ||||
|     return absl::InvalidArgumentError( | ||||
|         "At least one of `output_confidence_masks` and `output_category_mask` " | ||||
|         "must be set."); | ||||
|   } | ||||
|   auto options_proto = ConvertImageSegmenterOptionsToProto(options.get()); | ||||
|   tasks::core::PacketsCallback packets_callback = nullptr; | ||||
|   if (options->result_callback) { | ||||
|     auto result_callback = options->result_callback; | ||||
|     bool output_category_mask = options->output_category_mask; | ||||
|     bool output_confidence_masks = options->output_confidence_masks; | ||||
|     packets_callback = | ||||
|         [=](absl::StatusOr<tasks::core::PacketMap> status_or_packets) { | ||||
|           if (!status_or_packets.ok()) { | ||||
|  | @ -151,8 +161,12 @@ absl::StatusOr<std::unique_ptr<ImageSegmenter>> ImageSegmenter::Create( | |||
|           if (status_or_packets.value()[kImageOutStreamName].IsEmpty()) { | ||||
|             return; | ||||
|           } | ||||
|           Packet confidence_masks = | ||||
|               status_or_packets.value()[kConfidenceMasksStreamName]; | ||||
|           std::optional<std::vector<Image>> confidence_masks; | ||||
|           if (output_confidence_masks) { | ||||
|             confidence_masks = | ||||
|                 status_or_packets.value()[kConfidenceMasksStreamName] | ||||
|                     .Get<std::vector<Image>>(); | ||||
|           } | ||||
|           std::optional<Image> category_mask; | ||||
|           if (output_category_mask) { | ||||
|             category_mask = | ||||
|  | @ -160,23 +174,24 @@ absl::StatusOr<std::unique_ptr<ImageSegmenter>> ImageSegmenter::Create( | |||
|           } | ||||
|           Packet image_packet = status_or_packets.value()[kImageOutStreamName]; | ||||
|           result_callback( | ||||
|               {{confidence_masks.Get<std::vector<Image>>(), category_mask}}, | ||||
|               image_packet.Get<Image>(), | ||||
|               confidence_masks.Timestamp().Value() / | ||||
|                   kMicroSecondsPerMilliSecond); | ||||
|               {{confidence_masks, category_mask}}, image_packet.Get<Image>(), | ||||
|               image_packet.Timestamp().Value() / kMicroSecondsPerMilliSecond); | ||||
|         }; | ||||
|   } | ||||
|   auto image_segmenter = | ||||
|       core::VisionTaskApiFactory::Create<ImageSegmenter, | ||||
|                                          ImageSegmenterGraphOptionsProto>( | ||||
|           CreateGraphConfig( | ||||
|               std::move(options_proto), options->output_category_mask, | ||||
|               std::move(options_proto), options->output_confidence_masks, | ||||
|               options->output_category_mask, | ||||
|               options->running_mode == core::RunningMode::LIVE_STREAM), | ||||
|           std::move(options->base_options.op_resolver), options->running_mode, | ||||
|           std::move(packets_callback)); | ||||
|   if (!image_segmenter.ok()) { | ||||
|     return image_segmenter.status(); | ||||
|   } | ||||
|   image_segmenter.value()->output_confidence_masks_ = | ||||
|       options->output_confidence_masks; | ||||
|   image_segmenter.value()->output_category_mask_ = | ||||
|       options->output_category_mask; | ||||
|   ASSIGN_OR_RETURN( | ||||
|  | @ -203,8 +218,11 @@ absl::StatusOr<ImageSegmenterResult> ImageSegmenter::Segment( | |||
|           {{kImageInStreamName, mediapipe::MakePacket<Image>(std::move(image))}, | ||||
|            {kNormRectStreamName, | ||||
|             MakePacket<NormalizedRect>(std::move(norm_rect))}})); | ||||
|   std::vector<Image> confidence_masks = | ||||
|   std::optional<std::vector<Image>> confidence_masks; | ||||
|   if (output_confidence_masks_) { | ||||
|     confidence_masks = | ||||
|         output_packets[kConfidenceMasksStreamName].Get<std::vector<Image>>(); | ||||
|   } | ||||
|   std::optional<Image> category_mask; | ||||
|   if (output_category_mask_) { | ||||
|     category_mask = output_packets[kCategoryMaskStreamName].Get<Image>(); | ||||
|  | @ -233,8 +251,11 @@ absl::StatusOr<ImageSegmenterResult> ImageSegmenter::SegmentForVideo( | |||
|            {kNormRectStreamName, | ||||
|             MakePacket<NormalizedRect>(std::move(norm_rect)) | ||||
|                 .At(Timestamp(timestamp_ms * kMicroSecondsPerMilliSecond))}})); | ||||
|   std::vector<Image> confidence_masks = | ||||
|   std::optional<std::vector<Image>> confidence_masks; | ||||
|   if (output_confidence_masks_) { | ||||
|     confidence_masks = | ||||
|         output_packets[kConfidenceMasksStreamName].Get<std::vector<Image>>(); | ||||
|   } | ||||
|   std::optional<Image> category_mask; | ||||
|   if (output_category_mask_) { | ||||
|     category_mask = output_packets[kCategoryMaskStreamName].Get<Image>(); | ||||
|  |  | |||
|  | @ -53,6 +53,9 @@ struct ImageSegmenterOptions { | |||
|   // Metadata, if any. Defaults to English.
 | ||||
|   std::string display_names_locale = "en"; | ||||
| 
 | ||||
|   // Whether to output confidence masks.
 | ||||
|   bool output_confidence_masks = true; | ||||
| 
 | ||||
|   // Whether to output category mask.
 | ||||
|   bool output_category_mask = false; | ||||
| 
 | ||||
|  | @ -77,8 +80,10 @@ struct ImageSegmenterOptions { | |||
| //    - if type is kTfLiteFloat32, NormalizationOptions are required to be
 | ||||
| //      attached to the metadata for input normalization.
 | ||||
| // Output ImageSegmenterResult:
 | ||||
| //    Provides confidence masks and an optional category mask if
 | ||||
| //    `output_category_mask` is set true.
 | ||||
| //    Provides optional confidence masks if `output_confidence_masks` is set
 | ||||
| //    true,  and an optional category mask if `output_category_mask` is set
 | ||||
| //    true. At least one of `output_confidence_masks` and `output_category_mask`
 | ||||
| //    must be set to 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 { | ||||
|  | @ -167,6 +172,7 @@ class ImageSegmenter : tasks::vision::core::BaseVisionTaskApi { | |||
| 
 | ||||
|  private: | ||||
|   std::vector<std::string> labels_; | ||||
|   bool output_confidence_masks_; | ||||
|   bool output_category_mask_; | ||||
| }; | ||||
| 
 | ||||
|  |  | |||
|  | @ -326,8 +326,10 @@ absl::StatusOr<ImageAndTensorsOnDevice> ConvertImageToTensors( | |||
| } | ||||
| 
 | ||||
| // An "mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph" performs
 | ||||
| // semantic segmentation. The graph always output confidence masks, and an
 | ||||
| // optional category mask if CATEGORY_MASK is connected.
 | ||||
| // semantic segmentation. The graph can output optional confidence masks if
 | ||||
| // CONFIDENCE_MASKS is connected, and an optional category mask if CATEGORY_MASK
 | ||||
| // is connected. At least one of CONFIDENCE_MASK, CONFIDENCE_MASKS and
 | ||||
| // CATEGORY_MASK must be connected.
 | ||||
| //
 | ||||
| //  Two kinds of outputs for confidence mask are provided: CONFIDENCE_MASK and
 | ||||
| //  CONFIDENCE_MASKS. Users can retrieve segmented mask of only particular
 | ||||
|  | @ -347,7 +349,7 @@ absl::StatusOr<ImageAndTensorsOnDevice> ConvertImageToTensors( | |||
| //   CONFIDENCE_MASK - mediapipe::Image @Multiple
 | ||||
| //     Confidence masks for individual category. Confidence mask of single
 | ||||
| //     category can be accessed by index based output stream.
 | ||||
| //   CONFIDENCE_MASKS - std::vector<mediapipe::Image>
 | ||||
| //   CONFIDENCE_MASKS - std::vector<mediapipe::Image> @Optional
 | ||||
| //     The output confidence masks grouped in a vector.
 | ||||
| //   CATEGORY_MASK - mediapipe::Image @Optional
 | ||||
| //     Optional Category mask.
 | ||||
|  | @ -356,7 +358,7 @@ absl::StatusOr<ImageAndTensorsOnDevice> ConvertImageToTensors( | |||
| //
 | ||||
| // Example:
 | ||||
| // node {
 | ||||
| //   calculator: "mediapipe.tasks.vision.ImageSegmenterGraph"
 | ||||
| //   calculator: "mediapipe.tasks.vision.image_segmenter.ImageSegmenterGraph"
 | ||||
| //   input_stream: "IMAGE:image"
 | ||||
| //   output_stream: "SEGMENTATION:segmented_masks"
 | ||||
| //   options {
 | ||||
|  | @ -382,17 +384,20 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { | |||
|                      CreateModelResources<ImageSegmenterGraphOptions>(sc)); | ||||
|     Graph graph; | ||||
|     const auto& options = sc->Options<ImageSegmenterGraphOptions>(); | ||||
|     // TODO: remove deprecated output type support.
 | ||||
|     if (!options.segmenter_options().has_output_type()) { | ||||
|       MP_RETURN_IF_ERROR(SanityCheck(sc)); | ||||
|     } | ||||
|     ASSIGN_OR_RETURN( | ||||
|         auto output_streams, | ||||
|         BuildSegmentationTask( | ||||
|             options, *model_resources, graph[Input<Image>(kImageTag)], | ||||
|             graph[Input<NormalizedRect>::Optional(kNormRectTag)], | ||||
|             HasOutput(sc->OriginalNode(), kCategoryMaskTag), graph)); | ||||
|             graph[Input<NormalizedRect>::Optional(kNormRectTag)], graph)); | ||||
| 
 | ||||
|     auto& merge_images_to_vector = | ||||
|         graph.AddNode("MergeImagesToVectorCalculator"); | ||||
|     // TODO: remove deprecated output type support.
 | ||||
|     if (options.segmenter_options().has_output_type()) { | ||||
|       auto& merge_images_to_vector = | ||||
|           graph.AddNode("MergeImagesToVectorCalculator"); | ||||
|       for (int i = 0; i < output_streams.segmented_masks->size(); ++i) { | ||||
|         output_streams.segmented_masks->at(i) >> | ||||
|             merge_images_to_vector[Input<Image>::Multiple("")][i]; | ||||
|  | @ -402,6 +407,9 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { | |||
|       merge_images_to_vector.Out("") >> | ||||
|           graph[Output<std::vector<Image>>(kGroupedSegmentationTag)]; | ||||
|     } else { | ||||
|       if (output_streams.confidence_masks) { | ||||
|         auto& merge_images_to_vector = | ||||
|             graph.AddNode("MergeImagesToVectorCalculator"); | ||||
|         for (int i = 0; i < output_streams.confidence_masks->size(); ++i) { | ||||
|           output_streams.confidence_masks->at(i) >> | ||||
|               merge_images_to_vector[Input<Image>::Multiple("")][i]; | ||||
|  | @ -409,7 +417,8 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { | |||
|               graph[Output<Image>::Multiple(kConfidenceMaskTag)][i]; | ||||
|         } | ||||
|         merge_images_to_vector.Out("") >> | ||||
|           graph[Output<std::vector<Image>>(kConfidenceMasksTag)]; | ||||
|             graph[Output<std::vector<Image>>::Optional(kConfidenceMasksTag)]; | ||||
|       } | ||||
|       if (output_streams.category_mask) { | ||||
|         *output_streams.category_mask >> graph[Output<Image>(kCategoryMaskTag)]; | ||||
|       } | ||||
|  | @ -419,6 +428,19 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { | |||
|   } | ||||
| 
 | ||||
|  private: | ||||
|   absl::Status SanityCheck(mediapipe::SubgraphContext* sc) { | ||||
|     const auto& node = sc->OriginalNode(); | ||||
|     output_confidence_masks_ = HasOutput(node, kConfidenceMaskTag) || | ||||
|                                HasOutput(node, kConfidenceMasksTag); | ||||
|     output_category_mask_ = HasOutput(node, kCategoryMaskTag); | ||||
|     if (!output_confidence_masks_ && !output_category_mask_) { | ||||
|       return absl::InvalidArgumentError( | ||||
|           "At least one of CONFIDENCE_MASK, CONFIDENCE_MASKS and CATEGORY_MASK " | ||||
|           "must be connected."); | ||||
|     } | ||||
|     return absl::OkStatus(); | ||||
|   } | ||||
| 
 | ||||
|   // 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.
 | ||||
|  | @ -431,8 +453,7 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { | |||
|   absl::StatusOr<ImageSegmenterOutputs> BuildSegmentationTask( | ||||
|       const ImageSegmenterGraphOptions& task_options, | ||||
|       const core::ModelResources& model_resources, Source<Image> image_in, | ||||
|       Source<NormalizedRect> norm_rect_in, bool output_category_mask, | ||||
|       Graph& graph) { | ||||
|       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
 | ||||
|  | @ -485,26 +506,32 @@ class ImageSegmenterGraph : public core::ModelTaskGraph { | |||
|                                    /*category_mask=*/std::nullopt, | ||||
|                                    /*image=*/image_and_tensors.image}; | ||||
|     } else { | ||||
|       std::optional<std::vector<Source<Image>>> confidence_masks; | ||||
|       if (output_confidence_masks_) { | ||||
|         ASSIGN_OR_RETURN(const tflite::Tensor* output_tensor, | ||||
|                          GetOutputTensor(model_resources)); | ||||
|         int segmentation_streams_num = *output_tensor->shape()->rbegin(); | ||||
|       std::vector<Source<Image>> confidence_masks; | ||||
|       confidence_masks.reserve(segmentation_streams_num); | ||||
|         confidence_masks = std::vector<Source<Image>>(); | ||||
|         confidence_masks->reserve(segmentation_streams_num); | ||||
|         for (int i = 0; i < segmentation_streams_num; ++i) { | ||||
|         confidence_masks.push_back(Source<Image>( | ||||
|             tensor_to_images[Output<Image>::Multiple(kConfidenceMaskTag)][i])); | ||||
|           confidence_masks->push_back(Source<Image>( | ||||
|               tensor_to_images[Output<Image>::Multiple(kConfidenceMaskTag)] | ||||
|                               [i])); | ||||
|         } | ||||
|       return ImageSegmenterOutputs{ | ||||
|           /*segmented_masks=*/std::nullopt, | ||||
|       } | ||||
|       std::optional<Source<Image>> category_mask; | ||||
|       if (output_category_mask_) { | ||||
|         category_mask = tensor_to_images[Output<Image>(kCategoryMaskTag)]; | ||||
|       } | ||||
|       return ImageSegmenterOutputs{/*segmented_masks=*/std::nullopt, | ||||
|                                    /*confidence_masks=*/confidence_masks, | ||||
|           /*category_mask=*/ | ||||
|           output_category_mask | ||||
|               ? std::make_optional( | ||||
|                     tensor_to_images[Output<Image>(kCategoryMaskTag)]) | ||||
|               : std::nullopt, | ||||
|                                    /*category_mask=*/category_mask, | ||||
|                                    /*image=*/image_and_tensors.image}; | ||||
|     } | ||||
|   } | ||||
| 
 | ||||
|   bool output_confidence_masks_ = false; | ||||
|   bool output_category_mask_ = false; | ||||
| }; | ||||
| 
 | ||||
| REGISTER_MEDIAPIPE_GRAPH( | ||||
|  |  | |||
|  | @ -29,7 +29,7 @@ namespace image_segmenter { | |||
| struct ImageSegmenterResult { | ||||
|   // Multiple masks of float image in VEC32F1 format where, for each mask, each
 | ||||
|   // pixel represents the prediction confidence, usually in the [0, 1] range.
 | ||||
|   std::vector<Image> confidence_masks; | ||||
|   std::optional<std::vector<Image>> confidence_masks; | ||||
|   // A category mask of uint8 image in GRAY8 format where each pixel represents
 | ||||
|   // the class which the pixel in the original image was predicted to belong to.
 | ||||
|   std::optional<Image> category_mask; | ||||
|  |  | |||
|  | @ -278,6 +278,7 @@ TEST_F(ImageModeTest, SucceedsWithCategoryMask) { | |||
|   auto options = std::make_unique<ImageSegmenterOptions>(); | ||||
|   options->base_options.model_asset_path = | ||||
|       JoinPath("./", kTestDataDirectory, kDeeplabV3WithMetadata); | ||||
|   options->output_confidence_masks = false; | ||||
|   options->output_category_mask = true; | ||||
|   MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter, | ||||
|                           ImageSegmenter::Create(std::move(options))); | ||||
|  | @ -306,7 +307,7 @@ TEST_F(ImageModeTest, SucceedsWithConfidenceMask) { | |||
|   MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter, | ||||
|                           ImageSegmenter::Create(std::move(options))); | ||||
|   MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image)); | ||||
|   EXPECT_EQ(result.confidence_masks.size(), 21); | ||||
|   EXPECT_EQ(result.confidence_masks->size(), 21); | ||||
| 
 | ||||
|   cv::Mat expected_mask = cv::imread( | ||||
|       JoinPath("./", kTestDataDirectory, "cat_mask.jpg"), cv::IMREAD_GRAYSCALE); | ||||
|  | @ -315,7 +316,7 @@ TEST_F(ImageModeTest, SucceedsWithConfidenceMask) { | |||
| 
 | ||||
|   // Cat category index 8.
 | ||||
|   cv::Mat cat_mask = mediapipe::formats::MatView( | ||||
|       result.confidence_masks[8].GetImageFrameSharedPtr().get()); | ||||
|       result.confidence_masks->at(8).GetImageFrameSharedPtr().get()); | ||||
|   EXPECT_THAT(cat_mask, | ||||
|               SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); | ||||
| } | ||||
|  | @ -336,7 +337,7 @@ TEST_F(ImageModeTest, DISABLED_SucceedsWithRotation) { | |||
|   image_processing_options.rotation_degrees = -90; | ||||
|   MP_ASSERT_OK_AND_ASSIGN(auto result, | ||||
|                           segmenter->Segment(image, image_processing_options)); | ||||
|   EXPECT_EQ(result.confidence_masks.size(), 21); | ||||
|   EXPECT_EQ(result.confidence_masks->size(), 21); | ||||
| 
 | ||||
|   cv::Mat expected_mask = | ||||
|       cv::imread(JoinPath("./", kTestDataDirectory, "cat_rotated_mask.jpg"), | ||||
|  | @ -346,7 +347,7 @@ TEST_F(ImageModeTest, DISABLED_SucceedsWithRotation) { | |||
| 
 | ||||
|   // Cat category index 8.
 | ||||
|   cv::Mat cat_mask = mediapipe::formats::MatView( | ||||
|       result.confidence_masks[8].GetImageFrameSharedPtr().get()); | ||||
|       result.confidence_masks->at(8).GetImageFrameSharedPtr().get()); | ||||
|   EXPECT_THAT(cat_mask, | ||||
|               SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); | ||||
| } | ||||
|  | @ -384,7 +385,7 @@ TEST_F(ImageModeTest, SucceedsSelfie128x128Segmentation) { | |||
|   MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter, | ||||
|                           ImageSegmenter::Create(std::move(options))); | ||||
|   MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image)); | ||||
|   EXPECT_EQ(result.confidence_masks.size(), 2); | ||||
|   EXPECT_EQ(result.confidence_masks->size(), 2); | ||||
| 
 | ||||
|   cv::Mat expected_mask = | ||||
|       cv::imread(JoinPath("./", kTestDataDirectory, | ||||
|  | @ -395,7 +396,7 @@ TEST_F(ImageModeTest, SucceedsSelfie128x128Segmentation) { | |||
| 
 | ||||
|   // Selfie category index 1.
 | ||||
|   cv::Mat selfie_mask = mediapipe::formats::MatView( | ||||
|       result.confidence_masks[1].GetImageFrameSharedPtr().get()); | ||||
|       result.confidence_masks->at(1).GetImageFrameSharedPtr().get()); | ||||
|   EXPECT_THAT(selfie_mask, | ||||
|               SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); | ||||
| } | ||||
|  | @ -409,7 +410,7 @@ TEST_F(ImageModeTest, SucceedsSelfie144x256Segmentations) { | |||
|   MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter, | ||||
|                           ImageSegmenter::Create(std::move(options))); | ||||
|   MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image)); | ||||
|   EXPECT_EQ(result.confidence_masks.size(), 1); | ||||
|   EXPECT_EQ(result.confidence_masks->size(), 1); | ||||
| 
 | ||||
|   cv::Mat expected_mask = | ||||
|       cv::imread(JoinPath("./", kTestDataDirectory, | ||||
|  | @ -419,7 +420,7 @@ TEST_F(ImageModeTest, SucceedsSelfie144x256Segmentations) { | |||
|   expected_mask.convertTo(expected_mask_float, CV_32FC1, 1 / 255.f); | ||||
| 
 | ||||
|   cv::Mat selfie_mask = mediapipe::formats::MatView( | ||||
|       result.confidence_masks[0].GetImageFrameSharedPtr().get()); | ||||
|       result.confidence_masks->at(0).GetImageFrameSharedPtr().get()); | ||||
|   EXPECT_THAT(selfie_mask, | ||||
|               SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); | ||||
| } | ||||
|  | @ -434,7 +435,7 @@ TEST_F(ImageModeTest, SucceedsPortraitSelfieSegmentationConfidenceMask) { | |||
|   MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter, | ||||
|                           ImageSegmenter::Create(std::move(options))); | ||||
|   MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image)); | ||||
|   EXPECT_EQ(result.confidence_masks.size(), 1); | ||||
|   EXPECT_EQ(result.confidence_masks->size(), 1); | ||||
|   MP_ASSERT_OK(segmenter->Close()); | ||||
| 
 | ||||
|   cv::Mat expected_mask = cv::imread( | ||||
|  | @ -445,7 +446,7 @@ TEST_F(ImageModeTest, SucceedsPortraitSelfieSegmentationConfidenceMask) { | |||
|   expected_mask.convertTo(expected_mask_float, CV_32FC1, 1 / 255.f); | ||||
| 
 | ||||
|   cv::Mat selfie_mask = mediapipe::formats::MatView( | ||||
|       result.confidence_masks[0].GetImageFrameSharedPtr().get()); | ||||
|       result.confidence_masks->at(0).GetImageFrameSharedPtr().get()); | ||||
|   EXPECT_THAT(selfie_mask, | ||||
|               SimilarToFloatMask(expected_mask_float, kGoldenMaskSimilarity)); | ||||
| } | ||||
|  | @ -506,10 +507,10 @@ TEST_F(ImageModeTest, SucceedsHairSegmentation) { | |||
|   MP_ASSERT_OK_AND_ASSIGN(std::unique_ptr<ImageSegmenter> segmenter, | ||||
|                           ImageSegmenter::Create(std::move(options))); | ||||
|   MP_ASSERT_OK_AND_ASSIGN(auto result, segmenter->Segment(image)); | ||||
|   EXPECT_EQ(result.confidence_masks.size(), 2); | ||||
|   EXPECT_EQ(result.confidence_masks->size(), 2); | ||||
| 
 | ||||
|   cv::Mat hair_mask = mediapipe::formats::MatView( | ||||
|       result.confidence_masks[1].GetImageFrameSharedPtr().get()); | ||||
|       result.confidence_masks->at(1).GetImageFrameSharedPtr().get()); | ||||
|   MP_ASSERT_OK(segmenter->Close()); | ||||
|   cv::Mat expected_mask = cv::imread( | ||||
|       JoinPath("./", kTestDataDirectory, "portrait_hair_expected_mask.jpg"), | ||||
|  |  | |||
		Loading…
	
		Reference in New Issue
	
	Block a user