Add arg max map to the segmentation result

Author: JakubGonera

ios/RnExecutorch/models/image_segmentation/ImageSegmentationModel.h

@@ -6,3 +6,16 @@
 - (NSDictionary *)runModel:(cv::Mat &)input;
 
 @end
+
+template <typename T>
+NSArray* matToNSArray(const cv::Mat& mat) {
+    std::size_t numPixels = mat.rows * mat.cols;
+    NSMutableArray *arr = [[NSMutableArray alloc] initWithCapacity:numPixels];
+
+    for (std::size_t x = 0; x < mat.rows; ++x) {
+        for (std::size_t y = 0; y < mat.cols; ++y) {
+            arr[x * mat.cols + y] = @(mat.at<T>(x, y));
+        }
+    }
+    return arr;
+}

ios/RnExecutorch/models/image_segmentation/ImageSegmentationModel.mm

@@ -46,7 +46,6 @@ - (NSDictionary *)postprocess:(NSArray *)output {
   NSAssert(outputSize ==  numLabels * numModelPixels, 
         @"Model generated unexpected output size.");
 
-
   // For each label extract it's matrix and rescale it to the original size
   std::vector<cv::Mat> resizedLabelScores(numLabels);
   for (std::size_t label = 0; label < numLabels; ++label) {
@@ -61,6 +60,8 @@ - (NSDictionary *)postprocess:(NSArray *)output {
     cv::resize(labelMat, resizedLabelScores[label], originalSize);
   }
 
+  cv::Mat maxArg = cv::Mat(originalSize, CV_32S);
+
   // For each pixel apply softmax across all the labels
   for (std::size_t pixel = 0; pixel < numOriginalPixels; ++pixel) {
     int row = pixel / originalSize.width;
@@ -73,26 +74,30 @@ - (NSDictionary *)postprocess:(NSArray *)output {
 
     std::vector<double> adjustedScores = softmax(scores);
 
+    std::size_t maxArgIndex = 0;
+    double maxArgVal = 0;
     for (std::size_t label = 0; label < numLabels; ++label) {
       resizedLabelScores[label].at<double>(row, col) = adjustedScores[label];
+      if (adjustedScores[label] > maxArgVal) {
+        maxArgIndex = label;
+        maxArgVal = adjustedScores[label];
+      }
     }
+
+    maxArg.at<int>(row, col) = maxArgIndex;
   }
 
   NSMutableDictionary *result = [NSMutableDictionary dictionary];
 
+  // Convert to NSArray and populate the final dictionary
   for (std::size_t label = 0; label < numLabels; ++label) {
     NSString *labelString = @(deeplabv3_resnet50_labels[label].c_str());
-    NSMutableArray *arr = [[NSMutableArray alloc] initWithCapacity:numOriginalPixels];
-
-    for (std::size_t x = 0; x < originalSize.height; ++x) {
-        for (std::size_t y = 0; y < originalSize.width; ++y) {
-            arr[x * originalSize.width + y] = @(resizedLabelScores[label].at<double>(x, y));
-        }
-    }
-
+    NSMutableArray *arr = matToNSArray<double>(resizedLabelScores[label]);
     result[labelString] = arr;
   }
 
+  result[@"argmax"] = matToNSArray<int>(maxArg);
+
   return result;
 }
 

src/constants/image_segmentation/image_segmentation.ts

@@ -0,0 +1,23 @@
+export const classLabels = new Map<number, string>([
+  [0, 'background'],
+  [1, 'aeroplane'],
+  [2, 'bicycle'],
+  [3, 'bird'],
+  [4, 'boat'],
+  [5, 'bottle'],
+  [6, 'bus'],
+  [7, 'car'],
+  [8, 'cat'],
+  [9, 'chair'],
+  [10, 'cow'],
+  [11, 'diningtable'],
+  [12, 'dog'],
+  [13, 'horse'],
+  [14, 'motorbike'],
+  [15, 'person'],
+  [16, 'pottedplant'],
+  [17, 'sheep'],
+  [18, 'sofa'],
+  [19, 'train'],
+  [20, 'tvmonitor'],
+]);