Added Image Segmenter Result Helpers

mediapipe/tasks/ios/vision/image_segmenter/utils/BUILD

@@ -30,3 +30,13 @@ objc_library(
     ],
 )
 
+objc_library(
+    name = "MPPImageSegmenterResultHelpers",
+    srcs = ["sources/MPPImageSegmenterResult+Helpers.mm"],
+    hdrs = ["sources/MPPImageSegmenterResult+Helpers.h"],
+    deps = [
+        "//mediapipe/framework:packet",
+        "//mediapipe/framework/formats:image",
+        "//mediapipe/tasks/ios/vision/image_segmenter:MPPImageSegmenterResult",
+    ],
+)

mediapipe/tasks/ios/vision/image_segmenter/utils/sources/MPPImageSegmenterResult+Helpers.h

@@ -0,0 +1,48 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// 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.
+
+#import "mediapipe/tasks/ios/vision/image_segmenter/sources/MPPImageSegmenterResult.h"
+
+#include "mediapipe/framework/packet.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface MPPImageSegmenterResult (Helpers)
+
+/**
+ * Creates an `MPPImageSegmenterResult` from confidence masks, category mask and quality scores
+ * packets.
+ *
+ * If `shouldCopyMaskPacketData` is set to `YES`, the confidence and catergory masks of the newly
+ * created `MPPImageSegmenterResult` holds references to deep copied pixel data of the output
+ * respective masks.
+ *
+ * @param confidenceMasksPacket A MediaPipe packet wrapping a `std::vector<mediapipe::Image>`.
+ * @param categoryMaskPacket A MediaPipe packet wrapping a `<mediapipe::Image>`.
+ * @param qualityScoresPacket a MediaPipe packet wrapping a `std::vector<float>`.
+ * @param shouldCopyMaskPacketData A `BOOL` which indicates if the pixel data of the output masks
+ * must be deep copied to the newly created `MPPImageSegmenterResult`.
+ *
+ * @return  An `MPPImageSegmenterResult` object that contains the image segmentation results.
+ */
++ (MPPImageSegmenterResult *)
+    imageSegmenterResultWithConfidenceMasksPacket:(const mediapipe::Packet &)confidenceMasksPacket
+                               categoryMaskPacket:(const mediapipe::Packet &)categoryMaskPacket
+                              qualityScoresPacket:(const mediapipe::Packet &)qualityScoresPacket
+                          timestampInMilliseconds:(NSInteger)timestampInMilliseconds
+                         shouldCopyMaskPacketData:(BOOL)shouldCopyMaskPacketData;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/vision/image_segmenter/utils/sources/MPPImageSegmenterResult+Helpers.mm

@@ -0,0 +1,78 @@
+// Copyright 2023 The MediaPipe Authors.
+//
+// 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.
+
+#import "mediapipe/tasks/ios/vision/image_segmenter/utils/sources/MPPImageSegmenterResult+Helpers.h"
+
+#include "mediapipe/framework/formats/image.h"
+#include "mediapipe/framework/packet.h"
+
+namespace {
+using ::mediapipe::Image;
+using ::mediapipe::ImageFrameSharedPtr;
+using ::mediapipe::Packet;
+}  // namespace
+
+@implementation MPPImageSegmenterResult (Helpers)
+
++ (MPPImageSegmenterResult *)
+    imageSegmenterResultWithConfidenceMasksPacket:(const Packet &)confidenceMasksPacket
+                               categoryMaskPacket:(const Packet &)categoryMaskPacket
+                              qualityScoresPacket:(const Packet &)qualityScoresPacket
+                          timestampInMilliseconds:(NSInteger)timestampInMilliseconds
+                         shouldCopyMaskPacketData:(BOOL)shouldCopyMaskPacketData {
+  NSMutableArray<MPPMask *> *confidenceMasks;
+  MPPMask *categoryMask;
+  NSMutableArray<NSNumber *> *qualityScores;
+
+  if (confidenceMasksPacket.ValidateAsType<std::vector<Image>>().ok()) {
+    std::vector<Image> cppConfidenceMasks = confidenceMasksPacket.Get<std::vector<Image>>();
+    confidenceMasks = [NSMutableArray arrayWithCapacity:(NSUInteger)cppConfidenceMasks.size()];
+
+    for (const auto &confidenceMask : cppConfidenceMasks) {
+      [confidenceMasks
+          addObject:[[MPPMask alloc]
+                        initWithFloat32Data:(float *)confidenceMask.GetImageFrameSharedPtr()
+                                                .get()
+                                                ->PixelData()
+                                      width:confidenceMask.width()
+                                     height:confidenceMask.height()
+                                 shouldCopy:shouldCopyMaskPacketData ? YES : NO]];
+    }
+  }
+
+  if (categoryMaskPacket.ValidateAsType<Image>().ok()) {
+    const Image &cppCategoryMask = confidenceMasksPacket.Get<Image>();
+    categoryMask = [[MPPMask alloc]
+        initWithUInt8Data:(UInt8 *)cppCategoryMask.GetImageFrameSharedPtr().get()->PixelData()
+                    width:cppCategoryMask.width()
+                   height:cppCategoryMask.height()
+               shouldCopy:shouldCopyMaskPacketData ? YES : NO];
+  }
+
+  if (qualityScoresPacket.ValidateAsType<std::vector<float>>().ok()) {
+    std::vector<float> cppQualityScores = qualityScoresPacket.Get<std::vector<float>>();
+    qualityScores = [NSMutableArray arrayWithCapacity:(NSUInteger)cppQualityScores.size()];
+
+    for (const auto &qualityScore : cppQualityScores) {
+      [qualityScores addObject:[NSNumber numberWithFloat:qualityScore]];
+    }
+  }
+
+  return [[MPPImageSegmenterResult alloc] initWithConfidenceMasks:confidenceMasks
+                                                     categoryMask:categoryMask
+                                                    qualityScores:qualityScores
+                                          timestampInMilliseconds:timestampInMilliseconds];
+}
+
+@end