diff --git a/YOLO.xcodeproj/project.pbxproj b/YOLO.xcodeproj/project.pbxproj
index 68916e9..c2507f5 100644
--- a/YOLO.xcodeproj/project.pbxproj
+++ b/YOLO.xcodeproj/project.pbxproj
@@ -13,16 +13,24 @@
 		636EFCAF21E62DD300DE43BC /* VideoCapture.swift in Sources */ = {isa = PBXBuildFile; fileRef = 636EFCA221E62DD300DE43BC /* VideoCapture.swift */; };
 		636EFCB321E62DD300DE43BC /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 636EFCA721E62DD300DE43BC /* AppDelegate.swift */; };
 		636EFCB921E62E3900DE43BC /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 636EFCB821E62E3900DE43BC /* Assets.xcassets */; };
-		6381D2182B7817C200ABA4E8 /* yolov8l.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2132B7817C200ABA4E8 /* yolov8l.mlpackage */; };
-		6381D2192B7817C200ABA4E8 /* yolov8x.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2142B7817C200ABA4E8 /* yolov8x.mlpackage */; };
-		6381D21A2B7817C200ABA4E8 /* yolov8s.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2152B7817C200ABA4E8 /* yolov8s.mlpackage */; };
-		6381D21B2B7817C200ABA4E8 /* yolov8m.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2162B7817C200ABA4E8 /* yolov8m.mlpackage */; };
-		6381D21C2B7817C200ABA4E8 /* yolov8n.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2172B7817C200ABA4E8 /* yolov8n.mlpackage */; };
+	    	6381D2182B7817C200ABA4E8 /* yolov8l.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2132B7817C200ABA4E8 /* yolov8l.mlpackage */; };
+	    	6381D2192B7817C200ABA4E8 /* yolov8x.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2142B7817C200ABA4E8 /* yolov8x.mlpackage */; };
+	    	6381D21A2B7817C200ABA4E8 /* yolov8s.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2152B7817C200ABA4E8 /* yolov8s.mlpackage */; };
+	    	6381D21B2B7817C200ABA4E8 /* yolov8m.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2162B7817C200ABA4E8 /* yolov8m.mlpackage */; };
+	    	6381D21C2B7817C200ABA4E8 /* yolov8n.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 6381D2172B7817C200ABA4E8 /* yolov8n.mlpackage */; };
 		63CF371F2514455300E2DEA1 /* LaunchScreen.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 6323C44D22186177008AE681 /* LaunchScreen.storyboard */; };
 		63CF37202514455300E2DEA1 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 6323C44F22186177008AE681 /* Main.storyboard */; };
 		63CF37212514455300E2DEA1 /* ultralytics_yolo_logotype.png in Resources */ = {isa = PBXBuildFile; fileRef = 6323C45122186177008AE681 /* ultralytics_yolo_logotype.png */; };
+		733310CD2C71BEB0001D647B /* yolov8x-seg.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 733310C82C71BEAB001D647B /* yolov8x-seg.mlpackage */; };
+		733310CE2C71BEB0001D647B /* yolov8n-seg.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 733310C92C71BEAC001D647B /* yolov8n-seg.mlpackage */; };
+		733310CF2C71BEB0001D647B /* yolov8m-seg.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 733310CA2C71BEAD001D647B /* yolov8m-seg.mlpackage */; };
+		733310D02C71BEB0001D647B /* yolov8s-seg.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 733310CB2C71BEAE001D647B /* yolov8s-seg.mlpackage */; };
+		733310D12C71BEB0001D647B /* yolov8l-seg.mlpackage in Sources */ = {isa = PBXBuildFile; fileRef = 733310CC2C71BEB0001D647B /* yolov8l-seg.mlpackage */; };
+		733310D42C71C5CE001D647B /* Colors.swift in Sources */ = {isa = PBXBuildFile; fileRef = 733310D22C71C5CE001D647B /* Colors.swift */; };
+		733310D52C71C5CE001D647B /* PostProcessSegment.swift in Sources */ = {isa = PBXBuildFile; fileRef = 733310D32C71C5CE001D647B /* PostProcessSegment.swift */; };
 		8EDAA33950796844333D60A7 /* BoundingBoxView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8EDAA633C1F2B50286D16008 /* BoundingBoxView.swift */; };
 /* End PBXBuildFile section */
+
 /* Begin PBXFileReference section */
 		6323C44D22186177008AE681 /* LaunchScreen.storyboard */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; path = LaunchScreen.storyboard; sourceTree = "<group>"; };
 		6323C44F22186177008AE681 /* Main.storyboard */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; path = Main.storyboard; sourceTree = "<group>"; };
@@ -40,6 +48,13 @@
 		6381D2162B7817C200ABA4E8 /* yolov8m.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = yolov8m.mlpackage; sourceTree = "<group>"; };
 		6381D2172B7817C200ABA4E8 /* yolov8n.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = yolov8n.mlpackage; sourceTree = "<group>"; };
 		63B8B0A821E62A890026FBC3 /* .gitignore */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = .gitignore; sourceTree = "<group>"; };
+		733310C82C71BEAB001D647B /* yolov8x-seg.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = "yolov8x-seg.mlpackage"; sourceTree = "<group>"; };
+		733310C92C71BEAC001D647B /* yolov8n-seg.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = "yolov8n-seg.mlpackage"; sourceTree = "<group>"; };
+		733310CA2C71BEAD001D647B /* yolov8m-seg.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = "yolov8m-seg.mlpackage"; sourceTree = "<group>"; };
+		733310CB2C71BEAE001D647B /* yolov8s-seg.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = "yolov8s-seg.mlpackage"; sourceTree = "<group>"; };
+		733310CC2C71BEB0001D647B /* yolov8l-seg.mlpackage */ = {isa = PBXFileReference; lastKnownFileType = folder.mlpackage; path = "yolov8l-seg.mlpackage"; sourceTree = "<group>"; };
+		733310D22C71C5CE001D647B /* Colors.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = Colors.swift; sourceTree = "<group>"; };
+		733310D32C71C5CE001D647B /* PostProcessSegment.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = PostProcessSegment.swift; sourceTree = "<group>"; };
 		7BCB411721C3096100BFC4D0 /* YOLO.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = YOLO.app; sourceTree = BUILT_PRODUCTS_DIR; };
 		8EDAA633C1F2B50286D16008 /* BoundingBoxView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = BoundingBoxView.swift; sourceTree = "<group>"; };
 		8EDAAA4507D2D23D7FAB827F /* README.md */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = net.daringfireball.markdown; path = README.md; sourceTree = "<group>"; };
@@ -61,6 +76,8 @@
 			children = (
 				636166E9251443B20054FA7E /* ThresholdProvider.swift */,
 				8EDAA633C1F2B50286D16008 /* BoundingBoxView.swift */,
+				733310D32C71C5CE001D647B /* PostProcessSegment.swift */,
+				733310D22C71C5CE001D647B /* Colors.swift */,
 			);
 			path = Utilities;
 			sourceTree = "<group>";
@@ -91,6 +108,11 @@
 				6381D2172B7817C200ABA4E8 /* yolov8n.mlpackage */,
 				6381D2152B7817C200ABA4E8 /* yolov8s.mlpackage */,
 				6381D2142B7817C200ABA4E8 /* yolov8x.mlpackage */,
+				733310CC2C71BEB0001D647B /* yolov8l-seg.mlpackage */,
+				733310CA2C71BEAD001D647B /* yolov8m-seg.mlpackage */,
+				733310C92C71BEAC001D647B /* yolov8n-seg.mlpackage */,
+				733310CB2C71BEAE001D647B /* yolov8s-seg.mlpackage */,
+				733310C82C71BEAB001D647B /* yolov8x-seg.mlpackage */,
 			);
 			path = Models;
 			sourceTree = "<group>";
@@ -209,13 +231,20 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
-				6381D21B2B7817C200ABA4E8 /* yolov8m.mlpackage in Sources */,
 				6381D21C2B7817C200ABA4E8 /* yolov8n.mlpackage in Sources */,
-				636EFCAF21E62DD300DE43BC /* VideoCapture.swift in Sources */,
-				636166EA251443B20054FA7E /* ThresholdProvider.swift in Sources */,
-				6381D2182B7817C200ABA4E8 /* yolov8l.mlpackage in Sources */,
 				6381D21A2B7817C200ABA4E8 /* yolov8s.mlpackage in Sources */,
+				6381D21B2B7817C200ABA4E8 /* yolov8m.mlpackage in Sources */,
+				6381D2182B7817C200ABA4E8 /* yolov8l.mlpackage in Sources */,
 				6381D2192B7817C200ABA4E8 /* yolov8x.mlpackage in Sources */,
+				733310CE2C71BEB0001D647B /* yolov8n-seg.mlpackage in Sources */,
+				733310D02C71BEB0001D647B /* yolov8s-seg.mlpackage in Sources */,
+				733310CF2C71BEB0001D647B /* yolov8m-seg.mlpackage in Sources */,
+				733310D12C71BEB0001D647B /* yolov8l-seg.mlpackage in Sources */,
+				733310CD2C71BEB0001D647B /* yolov8x-seg.mlpackage in Sources */,
+				733310D52C71C5CE001D647B /* PostProcessSegment.swift in Sources */,
+				636EFCAF21E62DD300DE43BC /* VideoCapture.swift in Sources */,
+				733310D42C71C5CE001D647B /* Colors.swift in Sources */,
+				636166EA251443B20054FA7E /* ThresholdProvider.swift in Sources */,
 				636EFCB321E62DD300DE43BC /* AppDelegate.swift in Sources */,
 				636EFCAA21E62DD300DE43BC /* ViewController.swift in Sources */,
 				8EDAA33950796844333D60A7 /* BoundingBoxView.swift in Sources */,
diff --git a/YOLO/Main.storyboard b/YOLO/Main.storyboard
index 549bc72..4f34283 100644
--- a/YOLO/Main.storyboard
+++ b/YOLO/Main.storyboard
@@ -55,8 +55,8 @@
                                 <rect key="frame" x="178" y="405" width="37" height="37"/>
                                 <autoresizingMask key="autoresizingMask" flexibleMinX="YES" flexibleMaxX="YES" flexibleMinY="YES" flexibleMaxY="YES"/>
                             </activityIndicatorView>
-                            <label opaque="NO" userInteractionEnabled="NO" contentMode="center" horizontalHuggingPriority="251" verticalHuggingPriority="251" misplaced="YES" insetsLayoutMarginsFromSafeArea="NO" text="1.00x" textAlignment="center" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" translatesAutoresizingMaskIntoConstraints="NO" id="vlL-Ye-WgQ" userLabel="zoom">
-                                <rect key="frame" x="166" y="736" width="60" height="20"/>
+                            <label opaque="NO" userInteractionEnabled="NO" contentMode="center" horizontalHuggingPriority="251" verticalHuggingPriority="251" insetsLayoutMarginsFromSafeArea="NO" text="1.00x" textAlignment="center" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" translatesAutoresizingMaskIntoConstraints="NO" id="vlL-Ye-WgQ" userLabel="zoom">
+                                <rect key="frame" x="176.66666666666666" y="736" width="40" height="20"/>
                                 <accessibility key="accessibilityConfiguration">
                                     <accessibilityTraits key="traits" staticText="YES" updatesFrequently="YES"/>
                                 </accessibility>
@@ -139,7 +139,7 @@
                                 </connections>
                             </slider>
                             <toolbar opaque="NO" clearsContextBeforeDrawing="NO" contentMode="scaleToFill" horizontalHuggingPriority="1" verticalHuggingPriority="1" insetsLayoutMarginsFromSafeArea="NO" translatesAutoresizingMaskIntoConstraints="NO" id="3Rg-UC-ogI">
-                                <rect key="frame" x="0.0" y="787" width="393" height="44"/>
+                                <rect key="frame" x="0.0" y="782" width="393" height="49"/>
                                 <color key="backgroundColor" white="0.0" alpha="0.0" colorSpace="custom" customColorSpace="genericGamma22GrayColorSpace"/>
                                 <constraints>
                                     <constraint firstAttribute="height" constant="24" id="wAH-Ui-IPj"/>
@@ -229,6 +229,16 @@
                                 <color key="textColor" systemColor="groupTableViewBackgroundColor"/>
                                 <nil key="highlightedColor"/>
                             </label>
+                            <segmentedControl opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="left" contentVerticalAlignment="top" segmentControlStyle="plain" selectedSegmentIndex="0" translatesAutoresizingMaskIntoConstraints="NO" id="BbV-RI-qx8">
+                                <rect key="frame" x="267" y="171" width="121" height="32"/>
+                                <segments>
+                                    <segment title="Detect"/>
+                                    <segment title="Seg"/>
+                                </segments>
+                                <connections>
+                                    <action selector="taskSegmentControlChanged:" destination="BYZ-38-t0r" eventType="valueChanged" id="iTM-XR-4Wt"/>
+                                </connections>
+                            </segmentedControl>
                         </subviews>
                         <viewLayoutGuide key="safeArea" id="6Tk-OE-BBY"/>
                         <color key="backgroundColor" white="0.0" alpha="1" colorSpace="custom" customColorSpace="genericGamma22GrayColorSpace"/>
@@ -239,6 +249,7 @@
                             <constraint firstItem="mCE-j4-QAA" firstAttribute="leading" secondItem="8bC-Xf-vdC" secondAttribute="leading" id="2Hf-C7-PKF"/>
                             <constraint firstItem="tZo-QP-q2z" firstAttribute="top" secondItem="Gax-hR-qdY" secondAttribute="bottom" constant="1" id="3qP-Re-vUy"/>
                             <constraint firstItem="Hif-Vv-5nn" firstAttribute="top" secondItem="6Tk-OE-BBY" secondAttribute="top" id="4VU-KS-qnX"/>
+                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="BbV-RI-qx8" secondAttribute="trailing" constant="5" id="5wx-2K-np3"/>
                             <constraint firstItem="uD3-7n-e7T" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="6" id="8ED-Ty-itL"/>
                             <constraint firstItem="ewl-Fj-awn" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="5" id="CCH-hl-rSK"/>
                             <constraint firstItem="9cm-Nq-b3N" firstAttribute="top" secondItem="Hif-Vv-5nn" secondAttribute="bottom" constant="9" id="CsQ-eU-xjO"/>
@@ -263,6 +274,7 @@
                             <constraint firstItem="mCE-j4-QAA" firstAttribute="top" secondItem="8bC-Xf-vdC" secondAttribute="top" id="fJk-bN-O88"/>
                             <constraint firstItem="Ytf-30-5R8" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="8" id="fio-88-KON"/>
                             <constraint firstItem="3Rg-UC-ogI" firstAttribute="bottom" secondItem="8bC-Xf-vdC" secondAttribute="bottom" constant="-21" id="gd1-rI-r2l"/>
+                            <constraint firstItem="BbV-RI-qx8" firstAttribute="centerY" secondItem="Ytf-30-5R8" secondAttribute="centerY" id="giv-9H-E32"/>
                             <constraint firstItem="q4C-U4-wS4" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="5" id="hDd-mh-Ihd"/>
                             <constraint firstItem="ewl-Fj-awn" firstAttribute="top" secondItem="uD3-7n-e7T" secondAttribute="bottom" constant="4" id="iOx-Xn-wjC"/>
                             <constraint firstItem="tZo-QP-q2z" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="6" id="lkb-bl-PkB"/>
diff --git a/YOLO/Models/README.md b/YOLO/Models/README.md
index c43cffe..6a46404 100644
--- a/YOLO/Models/README.md
+++ b/YOLO/Models/README.md
@@ -27,15 +27,16 @@ If you prefer to use specific model versions or need to customize the models, yo
    pip install ultralytics
    ```
 
-2. **Export Models:** Use the following Python script to export YOLOv8 models to the CoreML format, optimized for INT8 quantization for better performance on iOS devices. The script exports all YOLOv8 model sizes (`n`, `s`, `m`, `l`, `x`) as CoreML models.
+2. **Export Models:** Use the following Python script to export YOLOv8 models to the CoreML format, optimized for INT8 quantization for better performance on iOS devices. The script exports all YOLOv8 model sizes (n, s, m, l, x) as CoreML models for both instance segmentation and object detection.
 
-   ```python
-   from ultralytics import YOLO
+```python
+from ultralytics import YOLO
 
-   # Export all YOLOv8 models to CoreML INT8
-   for size in ("n", "s", "m", "l", "x"):  # all YOLOv8 model sizes
-       YOLO(f"yolov8{size}.pt").export(format="coreml", int8=True, nms=True, imgsz=[640, 384])
-   ```
+# Export all YOLOv8 models to CoreML INT8
+for size in ("n", "s", "m", "l", "x"):  # all YOLOv8 model sizes
+    YOLO(f"yolov8{size}.pt").export(format="coreml", int8=True, nms=True, imgsz=[640, 384])
+    YOLO(f"yolov8{size}-seg.pt").export(format="coreml", int8=True, imgsz=[640, 384])
+```
 
 3. **Place Models in Project:** After exporting, locate the CoreML model files and place them in the `YOLO/Models` directory of your project.
 
diff --git a/YOLO/Utilities/Colors.swift b/YOLO/Utilities/Colors.swift
new file mode 100644
index 0000000..7edb817
--- /dev/null
+++ b/YOLO/Utilities/Colors.swift
@@ -0,0 +1,31 @@
+//  Ultralytics YOLO 🚀 - AGPL-3.0 License
+//
+//  Colors for Ultralytics YOLO App
+
+// These colors are set to use the same Ultralytics color palette as the Python version.
+
+import Foundation
+import UIKit
+
+let ultralyticsColorsolors: [UIColor] = [
+  UIColor(red: 4 / 255, green: 42 / 255, blue: 255 / 255, alpha: 0.6),  // #042AFF
+  UIColor(red: 11 / 255, green: 219 / 255, blue: 235 / 255, alpha: 0.6),  // #0BDBEB
+  UIColor(red: 243 / 255, green: 243 / 255, blue: 243 / 255, alpha: 0.6),  // #F3F3F3
+  UIColor(red: 0 / 255, green: 223 / 255, blue: 183 / 255, alpha: 0.6),  // #00DFB7
+  UIColor(red: 17 / 255, green: 31 / 255, blue: 104 / 255, alpha: 0.6),  // #111F68
+  UIColor(red: 255 / 255, green: 111 / 255, blue: 221 / 255, alpha: 0.6),  // #FF6FDD
+  UIColor(red: 255 / 255, green: 68 / 255, blue: 79 / 255, alpha: 0.6),  // #FF444F
+  UIColor(red: 204 / 255, green: 237 / 255, blue: 0 / 255, alpha: 0.6),  // #CCED00
+  UIColor(red: 0 / 255, green: 243 / 255, blue: 68 / 255, alpha: 0.6),  // #00F344
+  UIColor(red: 189 / 255, green: 0 / 255, blue: 255 / 255, alpha: 0.6),  // #BD00FF
+  UIColor(red: 0 / 255, green: 180 / 255, blue: 255 / 255, alpha: 0.6),  // #00B4FF
+  UIColor(red: 221 / 255, green: 0 / 255, blue: 186 / 255, alpha: 0.6),  // #DD00BA
+  UIColor(red: 0 / 255, green: 255 / 255, blue: 255 / 255, alpha: 0.6),  // #00FFFF
+  UIColor(red: 38 / 255, green: 192 / 255, blue: 0 / 255, alpha: 0.6),  // #26C000
+  UIColor(red: 1 / 255, green: 255 / 255, blue: 179 / 255, alpha: 0.6),  // #01FFB3
+  UIColor(red: 125 / 255, green: 36 / 255, blue: 255 / 255, alpha: 0.6),  // #7D24FF
+  UIColor(red: 123 / 255, green: 0 / 255, blue: 104 / 255, alpha: 0.6),  // #7B0068
+  UIColor(red: 255 / 255, green: 27 / 255, blue: 108 / 255, alpha: 0.6),  // #FF1B6C
+  UIColor(red: 252 / 255, green: 109 / 255, blue: 47 / 255, alpha: 0.6),  // #FC6D2F
+  UIColor(red: 162 / 255, green: 255 / 255, blue: 11 / 255, alpha: 0.6),  // #A2FF0B
+]
diff --git a/YOLO/Utilities/PostProcessSegment.swift b/YOLO/Utilities/PostProcessSegment.swift
new file mode 100644
index 0000000..b891a34
--- /dev/null
+++ b/YOLO/Utilities/PostProcessSegment.swift
@@ -0,0 +1,317 @@
+//  Ultralytics YOLO 🚀 - AGPL-3.0 License
+//
+//  PostProcessSegment for Ultralytics YOLO App
+
+// These functions are designed to post-process inference results from the YOLOv8-Segment model in the Ultralytics YOLO app to display segment masks.
+
+//  Access the source code: https://github.com/ultralytics/yolo-ios-app
+
+import Accelerate
+import MetalPerformanceShaders
+import UIKit
+import Vision
+
+extension ViewController {
+  func setupMaskLayer() {
+    let width = videoPreview.bounds.width
+    let height = videoPreview.bounds.height
+
+    var ratio: CGFloat = 1.0
+    if videoCapture.captureSession.sessionPreset == .photo {
+      ratio = (4.0 / 3.0)
+    } else {
+      ratio = (16.0 / 9.0)
+    }
+    var offSet = CGFloat.zero
+    var margin = CGFloat.zero
+    if view.bounds.width < view.bounds.height {
+      offSet = height / ratio
+      margin = (offSet - self.videoPreview.bounds.width) / 2
+      self.maskLayer.frame = CGRect(
+        x: -margin, y: 0, width: offSet, height: self.videoPreview.bounds.height)
+    } else {
+      offSet = width / ratio
+      margin = (offSet - self.videoPreview.bounds.height) / 2
+      self.maskLayer.frame = CGRect(
+        x: 0, y: -margin, width: self.videoPreview.bounds.width, height: offSet)
+
+    }
+  }
+
+  func getBoundingBoxesAndMasks(
+    feature: MLMultiArray, confidenceThreshold: Float, iouThreshold: Float
+  ) -> [(CGRect, Int, Float, MLMultiArray)] {
+    let numAnchors = feature.shape[2].intValue
+    let numFeatures = feature.shape[1].intValue
+    let boxFeatureLength = 4
+    let maskConfidenceLength = 32
+    let numClasses = numFeatures - boxFeatureLength - maskConfidenceLength
+
+    var results = [(CGRect, Int, Float, MLMultiArray)]()
+    let featurePointer = feature.dataPointer.assumingMemoryBound(to: Float.self)
+
+    let resultsQueue = DispatchQueue(label: "resultsQueue", attributes: .concurrent)
+
+    DispatchQueue.concurrentPerform(iterations: numAnchors) { j in
+      let baseOffset = j
+      let x = featurePointer[baseOffset]
+      let y = featurePointer[numAnchors + baseOffset]
+      let width = featurePointer[2 * numAnchors + baseOffset]
+      let height = featurePointer[3 * numAnchors + baseOffset]
+
+      let boxWidth = CGFloat(width)
+      let boxHeight = CGFloat(height)
+      let boxX = CGFloat(x - width / 2)
+      let boxY = CGFloat(y - height / 2)
+
+      let boundingBox = CGRect(x: boxX, y: boxY, width: boxWidth, height: boxHeight)
+
+      var classProbs = [Float](repeating: 0, count: numClasses)
+      classProbs.withUnsafeMutableBufferPointer { classProbsPointer in
+        vDSP_mtrans(
+          featurePointer + 4 * numAnchors + baseOffset, numAnchors, classProbsPointer.baseAddress!,
+          1, 1, vDSP_Length(numClasses))
+      }
+      var maxClassValue: Float = 0
+      var maxClassIndex: vDSP_Length = 0
+      vDSP_maxvi(classProbs, 1, &maxClassValue, &maxClassIndex, vDSP_Length(numClasses))
+
+      if maxClassValue > confidenceThreshold {
+        let maskProbsPointer = featurePointer + (4 + numClasses) * numAnchors + baseOffset
+        let maskProbs = try! MLMultiArray(
+          shape: [NSNumber(value: maskConfidenceLength)], dataType: .float32)
+        for i in 0..<maskConfidenceLength {
+          maskProbs[i] = NSNumber(value: maskProbsPointer[i * numAnchors])
+        }
+
+        let result = (boundingBox, Int(maxClassIndex), maxClassValue, maskProbs)
+
+        // Using resultsQueue to synchronize access to results
+        resultsQueue.async(flags: .barrier) {
+          results.append(result)
+        }
+      }
+    }
+
+    // Ensuring all updates to results are complete
+    resultsQueue.sync(flags: .barrier) {}
+
+    var selectedBoxesAndFeatures = [(CGRect, Int, Float, MLMultiArray)]()
+
+    // Perform NMS class by class
+    for classIndex in 0..<numClasses {
+      let classResults = results.filter { $0.1 == classIndex }
+      if !classResults.isEmpty {
+        let boxesOnly = classResults.map { $0.0 }
+        let scoresOnly = classResults.map { $0.2 }
+        let selectedIndices = nonMaxSuppression(
+          boxes: boxesOnly, scores: scoresOnly, threshold: iouThreshold)
+        for idx in selectedIndices {
+          selectedBoxesAndFeatures.append(
+            (classResults[idx].0, classResults[idx].1, classResults[idx].2, classResults[idx].3))
+        }
+      }
+    }
+
+    return selectedBoxesAndFeatures
+  }
+
+  func updateMaskAndBoxes(
+    detectedObjects: [(CGRect, Int, Float, MLMultiArray)], maskArray: MLMultiArray
+  ) {
+    if detectedObjects.isEmpty {
+      removeAllMaskSubLayers()
+      return
+    }
+
+    let startTime = Date()
+    let group = DispatchGroup()
+
+    let sortedObjects = detectedObjects.sorted {
+      $0.0.size.width * $0.0.size.height > $1.0.size.width * $1.0.size.height
+    }
+
+    var newLayers: [CALayer] = []
+
+    for (box, classIndex, _, masksIn) in sortedObjects {
+      group.enter()
+      DispatchQueue.global(qos: .userInitiated).async {
+        defer { group.leave() }
+        if let maskImage = self.generateColoredMaskImage(
+          from: masksIn, protos: maskArray, in: self.maskLayer.bounds.size, colorIndex: classIndex,
+          boundingBox: box)
+        {
+          DispatchQueue.main.async {
+            let adjustedBox = self.adjustBox(box, toFitIn: self.maskLayer.bounds.size)
+
+            let maskImageLayer = CALayer()
+            maskImageLayer.frame = adjustedBox
+            maskImageLayer.contents = maskImage
+            maskImageLayer.opacity = 0.5
+            newLayers.append(maskImageLayer)
+          }
+        }
+      }
+    }
+
+    group.notify(queue: .main) {
+      self.removeAllMaskSubLayers()
+      newLayers.forEach { self.maskLayer.addSublayer($0) }
+
+      //      print("Processing Time: \(Date().timeIntervalSince(startTime)) seconds")
+    }
+  }
+
+  func generateColoredMaskImage(
+    from masksIn: MLMultiArray, protos: MLMultiArray, in size: CGSize, colorIndex: Int,
+    boundingBox: CGRect
+  ) -> CGImage? {
+    let maskWidth = protos.shape[3].intValue
+    let maskHeight = protos.shape[2].intValue
+    let maskChannels = protos.shape[1].intValue
+
+    guard protos.shape.count == 4, protos.shape[0].intValue == 1, masksIn.shape.count == 1,
+      masksIn.shape[0].intValue == maskChannels
+    else {
+      print("Invalid shapes for protos or masksIn")
+      return nil
+    }
+
+    let masksPointer = masksIn.dataPointer.assumingMemoryBound(to: Float.self)
+    let protosPointer = protos.dataPointer.assumingMemoryBound(to: Float.self)
+
+    let masksPointerOutput = UnsafeMutablePointer<Float>.allocate(capacity: maskHeight * maskWidth)
+    vDSP_mmul(
+      masksPointer, 1, protosPointer, 1, masksPointerOutput, 1, vDSP_Length(1),
+      vDSP_Length(maskHeight * maskWidth), vDSP_Length(maskChannels))
+
+    let threshold: Float = 0.5
+    let color = colorsForMask[colorIndex]
+    let red = UInt8(color.red)
+    let green = UInt8(color.green)
+    let blue = UInt8(color.blue)
+
+    var maskPixels = [UInt8](repeating: 0, count: maskHeight * maskWidth * 4)
+    for y in 0..<maskHeight {
+      for x in 0..<maskWidth {
+        let index = y * maskWidth + x
+        let maskValue = masksPointerOutput[index]
+        if maskValue > threshold {
+          let pixelIndex = index * 4
+          maskPixels[pixelIndex] = red
+          maskPixels[pixelIndex + 1] = green
+          maskPixels[pixelIndex + 2] = blue
+          maskPixels[pixelIndex + 3] = 255
+        }
+      }
+    }
+
+    let maskDataPointer = UnsafeMutablePointer<UInt8>.allocate(capacity: maskPixels.count)
+    maskDataPointer.initialize(from: maskPixels, count: maskPixels.count)
+
+    let bitmapInfo = CGBitmapInfo(rawValue: CGImageAlphaInfo.premultipliedLast.rawValue)
+    let colorSpace = CGColorSpaceCreateDeviceRGB()
+
+    let maskDataProvider = CGDataProvider(
+      dataInfo: nil, data: maskDataPointer, size: maskPixels.count
+    ) { _, data, _ in
+      data.deallocate()
+    }
+
+    guard
+      let maskCGImage = CGImage(
+        width: maskWidth, height: maskHeight, bitsPerComponent: 8, bitsPerPixel: 32,
+        bytesPerRow: maskWidth * 4, space: colorSpace, bitmapInfo: bitmapInfo,
+        provider: maskDataProvider!, decode: nil, shouldInterpolate: true, intent: .defaultIntent)
+    else {
+      masksPointerOutput.deallocate()
+      return nil
+    }
+
+    let maskCIImage = CIImage(cgImage: maskCGImage)
+    let scaledCIImage = maskCIImage.transformed(
+      by: CGAffineTransform(
+        scaleX: size.width / CGFloat(maskWidth), y: size.height / CGFloat(maskHeight)))
+    let invertedY = size.height - (boundingBox.origin.y + boundingBox.height) * size.height / 640.0
+    let cropRect = CGRect(
+      x: boundingBox.origin.x * size.width / 640.0, y: invertedY,
+      width: boundingBox.width * size.width / 640.0,
+      height: boundingBox.height * size.height / 640.0)
+
+    let croppedCIImage = scaledCIImage.cropped(to: cropRect)
+
+    let ciContext = CIContext()
+    guard let cgImage = ciContext.createCGImage(croppedCIImage, from: cropRect) else {
+      masksPointerOutput.deallocate()
+      return nil
+    }
+
+    masksPointerOutput.deallocate()
+
+    return cgImage
+  }
+
+  func removeAllMaskSubLayers() {
+    self.maskLayer.sublayers?.forEach { layer in
+      layer.removeFromSuperlayer()
+    }
+    self.maskLayer.sublayers = nil
+  }
+
+  func adjustBox(_ box: CGRect, toFitIn containerSize: CGSize) -> CGRect {
+    let xScale = containerSize.width / 640.0
+    let yScale = containerSize.height / 640.0
+    return CGRect(
+      x: box.origin.x * xScale, y: box.origin.y * yScale, width: box.size.width * xScale,
+      height: box.size.height * yScale)
+  }
+}
+
+extension UIColor {
+  func toRGBComponents() -> (red: UInt8, green: UInt8, blue: UInt8)? {
+    var red: CGFloat = 0
+    var green: CGFloat = 0
+    var blue: CGFloat = 0
+    var alpha: CGFloat = 0
+
+    let success = self.getRed(&red, green: &green, blue: &blue, alpha: &alpha)
+
+    if success {
+      let redUInt8 = UInt8(red * 255.0)
+      let greenUInt8 = UInt8(green * 255.0)
+      let blueUInt8 = UInt8(blue * 255.0)
+      return (red: redUInt8, green: greenUInt8, blue: blueUInt8)
+    } else {
+      return nil
+    }
+  }
+}
+
+func nonMaxSuppression(boxes: [CGRect], scores: [Float], threshold: Float) -> [Int] {
+  let sortedIndices = scores.enumerated().sorted { $0.element > $1.element }.map { $0.offset }
+  var selectedIndices = [Int]()
+  var activeIndices = [Bool](repeating: true, count: boxes.count)
+
+  for i in 0..<sortedIndices.count {
+    let idx = sortedIndices[i]
+    if activeIndices[idx] {
+      selectedIndices.append(idx)
+      for j in i + 1..<sortedIndices.count {
+        let otherIdx = sortedIndices[j]
+        if activeIndices[otherIdx] {
+          let intersection = boxes[idx].intersection(boxes[otherIdx])
+          if intersection.area > CGFloat(threshold) * min(boxes[idx].area, boxes[otherIdx].area) {
+            activeIndices[otherIdx] = false
+          }
+        }
+      }
+    }
+  }
+  return selectedIndices
+}
+
+extension CGRect {
+  var area: CGFloat {
+    return width * height
+  }
+}
diff --git a/YOLO/ViewController.swift b/YOLO/ViewController.swift
index 840d91d..17cf0c5 100644
--- a/YOLO/ViewController.swift
+++ b/YOLO/ViewController.swift
@@ -19,6 +19,11 @@ import Vision
 
 var mlModel = try! yolov8m(configuration: .init()).model
 
+enum Task {
+  case detect
+  case seg
+}
+
 class ViewController: UIViewController {
   @IBOutlet var videoPreview: UIView!
   @IBOutlet var View0: UIView!
@@ -42,6 +47,7 @@ class ViewController: UIViewController {
   @IBOutlet weak var activityIndicator: UIActivityIndicatorView!
   @IBOutlet weak var forcus: UIImageView!
   @IBOutlet weak var toolBar: UIToolbar!
+  var maskLayer: CALayer = CALayer()
 
   let selection = UISelectionFeedbackGenerator()
   var detector = try! VNCoreMLModel(for: mlModel)
@@ -73,6 +79,11 @@ class ViewController: UIViewController {
     return request
   }()
 
+  var task: Task = .detect
+  var colors: [String: UIColor] = [:]
+  var colorsForMask: [(red: UInt8, green: UInt8, blue: UInt8)] = []
+  var classes: [String] = []
+
   override func viewDidLoad() {
     super.viewDidLoad()
     slider.value = 30
@@ -80,7 +91,6 @@ class ViewController: UIViewController {
     setUpBoundingBoxViews()
     setUpOrientationChangeNotification()
     startVideo()
-    // setModel()
   }
 
   override func viewWillTransition(
@@ -116,7 +126,13 @@ class ViewController: UIViewController {
     }
     self.videoCapture.previewLayer?.frame = CGRect(
       x: 0, y: 0, width: size.width, height: size.height)
-
+    coordinator.animate(
+      alongsideTransition: { context in
+      },
+      completion: { context in
+        self.setupMaskLayer()
+      }
+    )
   }
 
   private func setUpOrientationChangeNotification() {
@@ -134,29 +150,9 @@ class ViewController: UIViewController {
   }
 
   @IBAction func indexChanged(_ sender: Any) {
+    self.removeAllMaskSubLayers()
     selection.selectionChanged()
     activityIndicator.startAnimating()
-
-    /// Switch model
-    switch segmentedControl.selectedSegmentIndex {
-    case 0:
-      self.labelName.text = "YOLOv8n"
-      mlModel = try! yolov8n(configuration: .init()).model
-    case 1:
-      self.labelName.text = "YOLOv8s"
-      mlModel = try! yolov8s(configuration: .init()).model
-    case 2:
-      self.labelName.text = "YOLOv8m"
-      mlModel = try! yolov8m(configuration: .init()).model
-    case 3:
-      self.labelName.text = "YOLOv8l"
-      mlModel = try! yolov8l(configuration: .init()).model
-    case 4:
-      self.labelName.text = "YOLOv8x"
-      mlModel = try! yolov8x(configuration: .init()).model
-    default:
-      break
-    }
     setModel()
     setUpBoundingBoxViews()
     activityIndicator.stopAnimating()
@@ -164,21 +160,87 @@ class ViewController: UIViewController {
 
   func setModel() {
 
-    /// VNCoreMLModel
-    detector = try! VNCoreMLModel(for: mlModel)
-    detector.featureProvider = ThresholdProvider()
+    /// Switch model
+    switch task {
+    case .detect:
+      switch segmentedControl.selectedSegmentIndex {
+      case 0:
+        self.labelName.text = "YOLOv8n"
+        mlModel = try! yolov8n(configuration: .init()).model
+      case 1:
+        self.labelName.text = "YOLOv8s"
+        mlModel = try! yolov8s(configuration: .init()).model
+      case 2:
+        self.labelName.text = "YOLOv8m"
+        mlModel = try! yolov8m(configuration: .init()).model
+      case 3:
+        self.labelName.text = "YOLOv8l"
+        mlModel = try! yolov8l(configuration: .init()).model
+      case 4:
+        self.labelName.text = "YOLOv8x"
+        mlModel = try! yolov8x(configuration: .init()).model
+      default:
+        break
+      }
+    case .seg:
+      switch segmentedControl.selectedSegmentIndex {
+      case 0:
+        self.labelName.text = "YOLOv8n"
+        if #available(iOS 15.0, *) {
+          mlModel = try! yolov8n_seg(configuration: .init()).model
+        } else {
+          // Fallback on earlier versions
+        }
+      case 1:
+        self.labelName.text = "YOLOv8s"
+        if #available(iOS 15.0, *) {
+          mlModel = try! yolov8s_seg(configuration: .init()).model
+        } else {
+          // Fallback on earlier versions
+        }
+      case 2:
+        self.labelName.text = "YOLOv8m"
+        if #available(iOS 15.0, *) {
+          mlModel = try! yolov8m_seg(configuration: .init()).model
+        } else {
+          // Fallback on earlier versions
+        }
+      case 3:
+        self.labelName.text = "YOLOv8l"
+        if #available(iOS 15.0, *) {
+          mlModel = try! yolov8l_seg(configuration: .init()).model
+        } else {
+          // Fallback on earlier versions
+        }
+      case 4:
+        self.labelName.text = "YOLOv8x"
+        if #available(iOS 15.0, *) {
+          mlModel = try! yolov8x_seg(configuration: .init()).model
+        } else {
+          // Fallback on earlier versions
+        }
+      default: break
+      }
 
-    /// VNCoreMLRequest
-    let request = VNCoreMLRequest(
-      model: detector,
-      completionHandler: { [weak self] request, error in
-        self?.processObservations(for: request, error: error)
-      })
-    request.imageCropAndScaleOption = .scaleFill  // .scaleFit, .scaleFill, .centerCrop
-    visionRequest = request
-    t2 = 0.0  // inference dt smoothed
-    t3 = CACurrentMediaTime()  // FPS start
-    t4 = 0.0  // FPS dt smoothed
+    }
+    DispatchQueue.global(qos: .userInitiated).async { [self] in
+
+      /// VNCoreMLModel
+      detector = try! VNCoreMLModel(for: mlModel)
+      detector.featureProvider = ThresholdProvider()
+
+      /// VNCoreMLRequest
+      let request = VNCoreMLRequest(
+        model: detector,
+        completionHandler: { [weak self] request, error in
+          self?.processObservations(for: request, error: error)
+        })
+      request.imageCropAndScaleOption = .scaleFill  // .scaleFit, .scaleFill, .centerCrop
+      visionRequest = request
+      t2 = 0.0  // inference dt smoothed
+      t3 = CACurrentMediaTime()  // FPS start
+      t4 = 0.0  // FPS dt smoothed
+    }
   }
 
   /// Update thresholds from slider values
@@ -190,6 +252,28 @@ class ViewController: UIViewController {
     detector.featureProvider = ThresholdProvider(iouThreshold: iou, confidenceThreshold: conf)
   }
 
+  @IBAction func taskSegmentControlChanged(_ sender: UISegmentedControl) {
+    self.removeAllMaskSubLayers()
+
+    switch sender.selectedSegmentIndex {
+    case 0:
+      if self.task != .detect {
+        self.task = .detect
+        self.setModel()
+      }
+    case 1:
+      if self.task != .seg {
+        self.task = .seg
+        for i in 0..<self.boundingBoxViews.count {
+          self.boundingBoxViews[i].hide()
+        }
+        self.setModel()
+      }
+    default:
+      break
+    }
+  }
+
   @IBAction func takePhoto(_ sender: Any?) {
     let t0 = DispatchTime.now().uptimeNanoseconds
 
@@ -278,7 +362,6 @@ class ViewController: UIViewController {
 
   let maxBoundingBoxViews = 100
   var boundingBoxViews = [BoundingBoxView]()
-  var colors: [String: UIColor] = [:]
 
   func setUpBoundingBoxViews() {
     // Ensure all bounding box views are initialized up to the maximum allowed.
@@ -290,16 +373,30 @@ class ViewController: UIViewController {
     guard let classLabels = mlModel.modelDescription.classLabels as? [String] else {
       fatalError("Class labels are missing from the model description")
     }
+    classes = classLabels
 
     // Assign random colors to the classes.
+    var count = 0
     for label in classLabels {
+      let color = ultralyticsColorsolors[count]
+      count += 1
+      if count > 19 {
+        count = 0
+      }
       if colors[label] == nil {  // if key not in dict
-        colors[label] = UIColor(
-          red: CGFloat.random(in: 0...1),
-          green: CGFloat.random(in: 0...1),
-          blue: CGFloat.random(in: 0...1),
-          alpha: 0.6)
+        colors[label] = color
+      }
+    }
+
+    count = 0
+    for (key, color) in colors {
+      let color = ultralyticsColorsolors[count]
+      count += 1
+      if count > 19 {
+        count = 0
       }
+      guard let colorForMask = color.toRGBComponents() else { fatalError() }
+      colorsForMask.append(colorForMask)
     }
   }
 
@@ -315,6 +412,8 @@ class ViewController: UIViewController {
           self.videoPreview.layer.addSublayer(previewLayer)
           self.videoCapture.previewLayer?.frame = self.videoPreview.bounds  // resize preview layer
         }
+        self.setupMaskLayer()
+        self.videoPreview.layer.addSublayer(self.maskLayer)
 
         // Add the bounding box layers to the UI, on top of the video preview.
         for box in self.boundingBoxViews {
@@ -368,21 +467,49 @@ class ViewController: UIViewController {
   }
 
   func processObservations(for request: VNRequest, error: Error?) {
-    DispatchQueue.main.async {
-      if let results = request.results as? [VNRecognizedObjectObservation] {
-        self.show(predictions: results)
-      } else {
-        self.show(predictions: [])
+    switch task {
+    case .detect:
+      DispatchQueue.main.async {
+        if let results = request.results as? [VNRecognizedObjectObservation] {
+          self.show(predictions: results, processedBoxAndMasks: [])
+        } else {
+          self.show(predictions: [], processedBoxAndMasks: [])
+        }
+
+        // Measure FPS
+        if self.t1 < 10.0 {  // valid dt
+          self.t2 = self.t1 * 0.05 + self.t2 * 0.95  // smoothed inference time
+        }
+        self.t4 = (CACurrentMediaTime() - self.t3) * 0.05 + self.t4 * 0.95  // smoothed delivered FPS
+        self.labelFPS.text = String(format: "%.1f FPS - %.1f ms", 1 / self.t4, self.t2 * 1000)  // t2 seconds to ms
+        self.t3 = CACurrentMediaTime()
       }
+    case .seg:
+      if let results = request.results as? [VNCoreMLFeatureValueObservation] {
+        DispatchQueue.main.async { [self] in
+          guard results.count == 2 else { return }
+          let masks = results[0].featureValue.multiArrayValue
+          let pred = results[1].featureValue.multiArrayValue
+          let processed = getBoundingBoxesAndMasks(
+            feature: pred!, confidenceThreshold: 0.25, iouThreshold: 0.4)
+
+          self.show(predictions: [], processedBoxAndMasks: processed)
+          DispatchQueue.main.async {
+            let a = Date()
+            self.updateMaskAndBoxes(detectedObjects: processed, maskArray: masks!)
+            print(Date().timeIntervalSince(a))
+          }
 
-      // Measure FPS
-      if self.t1 < 10.0 {  // valid dt
-        self.t2 = self.t1 * 0.05 + self.t2 * 0.95  // smoothed inference time
+          if self.t1 < 10.0 {  // valid dt
+            self.t2 = self.t1 * 0.05 + self.t2 * 0.95  // smoothed inference time
+          }
+          self.t4 = (CACurrentMediaTime() - self.t3) * 0.05 + self.t4 * 0.95  // smoothed delivered FPS
+          self.labelFPS.text = String(format: "%.1f FPS - %.1f ms", 1 / self.t4, self.t2 * 1000)  // t2 seconds to ms
+          self.t3 = CACurrentMediaTime()
+        }
       }
-      self.t4 = (CACurrentMediaTime() - self.t3) * 0.05 + self.t4 * 0.95  // smoothed delivered FPS
-      self.labelFPS.text = String(format: "%.1f FPS - %.1f ms", 1 / self.t4, self.t2 * 1000)  // t2 seconds to ms
-      self.t3 = CACurrentMediaTime()
     }
+
   }
 
   // Save text file
@@ -448,20 +575,22 @@ class ViewController: UIViewController {
     }
   }
 
-  func show(predictions: [VNRecognizedObjectObservation]) {
-    let width = videoPreview.bounds.width  // 375 pix
-    let height = videoPreview.bounds.height  // 812 pix
+  func show(
+    predictions: [VNRecognizedObjectObservation],
+    processedBoxAndMasks: [(CGRect, Int, Float, MLMultiArray)]
+  ) {
+    let width = videoPreview.bounds.width
+    let height = videoPreview.bounds.height
     var str = ""
 
-    // ratio = videoPreview AR divided by sessionPreset AR
     var ratio: CGFloat = 1.0
+
     if videoCapture.captureSession.sessionPreset == .photo {
-      ratio = (height / width) / (4.0 / 3.0)  // .photo
+      ratio = (height / width) / (4.0 / 3.0)
     } else {
-      ratio = (height / width) / (16.0 / 9.0)  // .hd4K3840x2160, .hd1920x1080, .hd1280x720 etc.
+      ratio = (height / width) / (16.0 / 9.0)
     }
 
-    // date
     let date = Date()
     let calendar = Calendar.current
     let hour = calendar.component(.hour, from: date)
@@ -469,30 +598,61 @@ class ViewController: UIViewController {
     let seconds = calendar.component(.second, from: date)
     let nanoseconds = calendar.component(.nanosecond, from: date)
     let sec_day =
-      Double(hour) * 3600.0 + Double(minutes) * 60.0 + Double(seconds) + Double(nanoseconds) / 1E9  // seconds in the day
+      Double(hour) * 3600.0 + Double(minutes) * 60.0 + Double(seconds) + Double(nanoseconds) / 1E9
 
-    self.labelSlider.text =
-      String(predictions.count) + " items (max " + String(Int(slider.value)) + ")"
-    for i in 0..<boundingBoxViews.count {
-      if i < predictions.count && i < Int(slider.value) {
-        let prediction = predictions[i]
+    var resultCount = 0
 
-        var rect = prediction.boundingBox  // normalized xywh, origin lower left
+    switch task {
+    case .detect:
+      resultCount = predictions.count
+    case .seg:
+      resultCount = processedBoxAndMasks.count
+    }
+    self.labelSlider.text = String(resultCount) + " items (max " + String(Int(slider.value)) + ")"
+    for i in 0..<boundingBoxViews.count {
+      if i < (resultCount) && i < Int(slider.value) {
+        var rect = CGRect.zero
+        var label = ""
+        var boxColor: UIColor = .white
+        var confidence: CGFloat = 0
+        var alpha: CGFloat = 0.9
+        var bestClass = ""
+        switch task {
+        case .detect:
+          let prediction = predictions[i]
+          rect = prediction.boundingBox
+          bestClass = prediction.labels[0].identifier
+          confidence = CGFloat(prediction.labels[0].confidence)
+          label = String(format: "%@ %.1f", bestClass, confidence * 100)
+          boxColor = colors[bestClass] ?? UIColor.white
+          alpha = CGFloat((confidence - 0.2) / (1.0 - 0.2) * 0.9)
+        case .seg:
+          let processed = processedBoxAndMasks[i]
+          let box = processed.0
+          rect = CGRect(
+            x: box.minX / 640, y: box.minY / 640, width: box.width / 640, height: box.height / 640)
+          confidence = CGFloat(processed.2)
+          bestClass = classes[processed.1]
+          label = String(format: "%@ %.1f", bestClass, confidence * 100)
+          boxColor = colors[bestClass] ?? UIColor.white
+          alpha = CGFloat((confidence - 0.2) / (1.0 - 0.2) * 0.9)
+        }
+        var displayRect = rect
         switch UIDevice.current.orientation {
         case .portraitUpsideDown:
-          rect = CGRect(
+          displayRect = CGRect(
             x: 1.0 - rect.origin.x - rect.width,
             y: 1.0 - rect.origin.y - rect.height,
             width: rect.width,
             height: rect.height)
         case .landscapeLeft:
-          rect = CGRect(
+          displayRect = CGRect(
             x: rect.origin.x,
             y: rect.origin.y,
             width: rect.width,
             height: rect.height)
         case .landscapeRight:
-          rect = CGRect(
+          displayRect = CGRect(
             x: rect.origin.x,
             y: rect.origin.y,
             width: rect.width,
@@ -502,55 +662,43 @@ class ViewController: UIViewController {
           fallthrough
         default: break
         }
-
-        if ratio >= 1 {  // iPhone ratio = 1.218
-          let offset = (1 - ratio) * (0.5 - rect.minX)
-          let transform = CGAffineTransform(scaleX: 1, y: -1).translatedBy(x: offset, y: -1)
-          rect = rect.applying(transform)
-          rect.size.width *= ratio
-        } else {  // iPad ratio = 0.75
-          let offset = (ratio - 1) * (0.5 - rect.maxY)
-          let transform = CGAffineTransform(scaleX: 1, y: -1).translatedBy(x: 0, y: offset - 1)
-          rect = rect.applying(transform)
+        if ratio >= 1 {
+          let offset = (1 - ratio) * (0.5 - displayRect.minX)
+          if task == .detect {
+            let transform = CGAffineTransform(scaleX: 1, y: -1).translatedBy(x: offset, y: -1)
+            displayRect = displayRect.applying(transform)
+          } else {
+            let transform = CGAffineTransform(translationX: offset, y: 0)
+            displayRect = displayRect.applying(transform)
+          }
+          displayRect.size.width *= ratio
+        } else {
+          if task == .detect {
+            let offset = (ratio - 1) * (0.5 - displayRect.maxY)
+            let transform = CGAffineTransform(scaleX: 1, y: -1).translatedBy(x: 0, y: offset - 1)
+            displayRect = displayRect.applying(transform)
+          } else {
+            let offset = (ratio - 1) * (0.5 - displayRect.minY)
+            let transform = CGAffineTransform(translationX: 0, y: offset)
+            displayRect = displayRect.applying(transform)
+          }
           ratio = (height / width) / (3.0 / 4.0)
-          rect.size.height /= ratio
+          displayRect.size.height /= ratio
         }
+        displayRect = VNImageRectForNormalizedRect(displayRect, Int(width), Int(height))
 
-        // Scale normalized to pixels [375, 812] [width, height]
-        rect = VNImageRectForNormalizedRect(rect, Int(width), Int(height))
-
-        // The labels array is a list of VNClassificationObservation objects,
-        // with the highest scoring class first in the list.
-        let bestClass = prediction.labels[0].identifier
-        let confidence = prediction.labels[0].confidence
-        // print(confidence, rect)  // debug (confidence, xywh) with xywh origin top left (pixels)
-        let label = String(format: "%@ %.1f", bestClass, confidence * 100)
-        let alpha = CGFloat((confidence - 0.2) / (1.0 - 0.2) * 0.9)
-        // Show the bounding box.
         boundingBoxViews[i].show(
-          frame: rect,
-          label: label,
-          color: colors[bestClass] ?? UIColor.white,
-          alpha: alpha)  // alpha 0 (transparent) to 1 (opaque) for conf threshold 0.2 to 1.0)
+          frame: displayRect, label: label, color: boxColor, alpha: alpha)
 
         if developerMode {
-          // Write
           if save_detections {
             str += String(
               format: "%.3f %.3f %.3f %@ %.2f %.1f %.1f %.1f %.1f\n",
               sec_day, freeSpace(), UIDevice.current.batteryLevel, bestClass, confidence,
               rect.origin.x, rect.origin.y, rect.size.width, rect.size.height)
           }
-
-          // Action trigger upon detection
-          // if false {
-          //     if (bestClass == "car") {  // "cell phone", "car", "person"
-          //         self.takePhoto(nil)
-          //         // self.pauseButton(nil)
-          //         sleep(2)
-          //     }
-          // }
         }
+
       } else {
         boundingBoxViews[i].hide()
       }
@@ -569,11 +717,6 @@ class ViewController: UIViewController {
         saveText(text: str, file: "frames.txt")  // Write stats for each image
       }
     }
-
-    // Debug
-    // print(str)
-    // print(UIDevice.current.identifierForVendor!)
-    // saveImage()
   }
 
   // Pinch to Zoom Start ---------------------------------------------------------------------------------------------