Merge pull request #13 from roboflow/bugfix-incorrect-seg-model-output-ordering

Update RFInstanceSegmentationModel.swift

Author: MadeWithStone

Sources/Roboflow/Classes/RFInstanceSegmentationModel.swift

@@ -67,19 +67,20 @@ public class RFInstanceSegmentationModel: RFObjectDetectionModel {
         do {
             try handler.perform([coreMLRequest])
             guard let detectResults = coreMLRequest.results else { return }
+
+            let castDetectResults0 = (detectResults[0] as! VNCoreMLFeatureValueObservation).featureValue.multiArrayValue!
+            let castDetectResults1 = (detectResults[1] as! VNCoreMLFeatureValueObservation).featureValue.multiArrayValue!
 
-            let predictions = detectResults[1] as! VNCoreMLFeatureValueObservation
-            let protos = detectResults[0] as! VNCoreMLFeatureValueObservation
-            
-            let pred = predictions.featureValue.multiArrayValue!
-            let proto = protos.featureValue.multiArrayValue!
+            let pred = castDetectResults0.shape.count == 3 ? castDetectResults0 : castDetectResults1
+            let proto = castDetectResults1.shape.count == 4 ? castDetectResults1 : castDetectResults0
 
             let numMasks = 32
             let numCls = self.colors.count
 
             // --- flatten MLMultiArray to Swift [Float] for speed
             let p = pred.dataPointer.bindMemory(to: Float.self,
                                                 capacity: pred.count)
+                                                
             let preds = UnsafeBufferPointer(start: p, count: pred.count)
             let protoShape = (c:Int(truncating: proto.shape[1]),
                               h:Int(truncating: proto.shape[2]),