Issue fixed: dangling pointer that 'UnsafeMutablePointer()' casused made invalid model outputs

husthyc · husthyc · commit edc7dcde10b5 · 2021-11-18T21:16:31.000-08:00
diff --git a/D2Go/D2Go/Live/LiveObjectDetectionViewController.swift b/D2Go/D2Go/Live/LiveObjectDetectionViewController.swift
@@ -17,43 +17,51 @@ class LiveObjectDetectionViewController: ViewController {
     private var cameraController = CameraController()
     private var imageViewLive =  UIImageView()
     private var inferencer = ObjectDetector()
-    
+
     override func viewDidLoad() {
         super.viewDidLoad()
         cameraController.configPreviewLayer(cameraView)
         imageViewLive.frame = CGRect(x: 0, y: 0, width: cameraView.frame.size.width, height: cameraView.frame.size.height)
         cameraView.addSubview(imageViewLive)
-        
+
         cameraController.videoCaptureCompletionBlock = { [weak self] buffer, error in
             guard let strongSelf = self else { return }
             if error != nil {
                 return
             }
-            guard var pixelBuffer = buffer else { return }
-            
+            guard let pixelBuffer = buffer else { return }
+
             let currentTimestamp = CACurrentMediaTime()
             if (currentTimestamp - strongSelf.prevTimestampMs) * 1000 <= strongSelf.delayMs { return }
             strongSelf.prevTimestampMs = currentTimestamp
+            
+            // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+            // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+            let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+            copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+
             let startTime = CACurrentMediaTime()
-            guard let outputs = self?.inferencer.module.detect(image: &pixelBuffer) else {
+            guard let outputs = self?.inferencer.module.detect(image: copiedBufferPtr) else {
+                copiedBufferPtr.deallocate()
                 return
             }
+            copiedBufferPtr.deallocate()
             let inferenceTime = CACurrentMediaTime() - startTime
-                
+
             DispatchQueue.main.async {
                 let ivScaleX : Double =  Double(strongSelf.imageViewLive.frame.size.width / CGFloat(PrePostProcessor.inputWidth))
                 let ivScaleY : Double = Double(strongSelf.imageViewLive.frame.size.height / CGFloat(PrePostProcessor.inputHeight))
 
                 let startX = Double((strongSelf.imageViewLive.frame.size.width - CGFloat(ivScaleX) * CGFloat(PrePostProcessor.inputWidth))/2)
                 let startY = Double((strongSelf.imageViewLive.frame.size.height -  CGFloat(ivScaleY) * CGFloat(PrePostProcessor.inputHeight))/2)
-                
+
                 let predictions = PrePostProcessor.outputsToPredictions(outputs: outputs, imgScaleX: 1.0, imgScaleY: 1.0, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
 
                 PrePostProcessor.cleanDetection(imageView: strongSelf.imageViewLive)
                 strongSelf.indicator.isHidden = true
                 strongSelf.benchmarkLabel.isHidden = false
                 strongSelf.benchmarkLabel.text = String(format: "%.2fms", 1000*inferenceTime)
-                
+
                 PrePostProcessor.showDetection(imageView: strongSelf.imageViewLive, nmsPredictions: predictions, classes: strongSelf.inferencer.classes)
             }
         }
diff --git a/D2Go/D2Go/ViewController.swift b/D2Go/D2Go/ViewController.swift
@@ -10,13 +10,13 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
     @IBOutlet weak var imageView: UIImageView!
     @IBOutlet weak var btnRun: UIButton!
     @IBOutlet weak var btnNext: UIButton!
-    
+
     private let testImages = ["test1.png", "test2.jpg", "test3.png"]
     private var imgIndex = 0
 
     private var image : UIImage?
     private var inferencer = ObjectDetector()
-    
+
     override func viewDidLoad() {
         super.viewDidLoad()
         image = UIImage(named: testImages[imgIndex])!
@@ -31,27 +31,33 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
         btnRun.setTitle("Running the model...", for: .normal)
 
         let resizedImage = image!.resized(to: CGSize(width: CGFloat(PrePostProcessor.inputWidth), height: CGFloat(PrePostProcessor.inputHeight)))
-        
+
         let imgScaleX = Double(image!.size.width / CGFloat(PrePostProcessor.inputWidth));
         let imgScaleY = Double(image!.size.height / CGFloat(PrePostProcessor.inputHeight));
-        
+
         let ivScaleX : Double = (image!.size.width > image!.size.height ? Double(imageView.frame.size.width / imageView.image!.size.width) : Double(imageView.image!.size.width / imageView.image!.size.height))
         let ivScaleY : Double = (image!.size.height > image!.size.width ? Double(imageView.frame.size.height / imageView.image!.size.height) : Double(imageView.image!.size.height / imageView.image!.size.width))
 
         let startX = Double((imageView.frame.size.width - CGFloat(ivScaleX) * imageView.image!.size.width)/2)
         let startY = Double((imageView.frame.size.height -  CGFloat(ivScaleY) * imageView.image!.size.height)/2)
 
-        guard var pixelBuffer = resizedImage.normalized() else {
+        guard let pixelBuffer = resizedImage.normalized() else {
             return
         }
-        
+
         DispatchQueue.global().async {
-            guard let outputs = self.inferencer.module.detect(image: &pixelBuffer) else {
+            // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+            // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+            let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+            copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+            guard let outputs = self.inferencer.module.detect(image: copiedBufferPtr) else {
+                copiedBufferPtr.deallocate()
                 return
             }
-                        
+            copiedBufferPtr.deallocate()
+
             let predictions = PrePostProcessor.outputsToPredictions(outputs: outputs, imgScaleX: imgScaleX, imgScaleY: imgScaleY, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
-                        
+
             DispatchQueue.main.async {
                 PrePostProcessor.showDetection(imageView: self.imageView, nmsPredictions: predictions, classes: self.inferencer.classes)
                 self.btnRun.isEnabled = true
@@ -75,7 +81,7 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
         imagePickerController.sourceType = .photoLibrary
         self.present(imagePickerController, animated: true, completion: nil)
     }
-    
+
     @IBAction func cameraTapped(_ sender: Any) {
         PrePostProcessor.cleanDetection(imageView: imageView)
         if UIImagePickerController.isSourceTypeAvailable(.camera) {
@@ -85,13 +91,11 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
             self.present(imagePickerController, animated: true, completion: nil)
         }
     }
-    
+
     func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
         image = info[UIImagePickerController.InfoKey.originalImage] as? UIImage
         image = image!.resized(to: CGSize(width: CGFloat(PrePostProcessor.inputWidth), height: CGFloat(PrePostProcessor.inputHeight)*image!.size.height/image!.size.width))
         imageView.image = image
         self.dismiss(animated: true, completion: nil)
     }
 }
-
-
diff --git a/HelloWorld-CoreML/HelloWorld/HelloWorld/ViewController.swift b/HelloWorld-CoreML/HelloWorld/HelloWorld/ViewController.swift
@@ -26,12 +26,18 @@ class ViewController: UIViewController {
         let image = UIImage(named: "image.png")!
         imageView.image = image
         let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
-        guard var pixelBuffer = resizedImage.normalized() else {
+        guard let pixelBuffer = resizedImage.normalized() else {
             return
         }
-        guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
+        // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+        // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+        let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+        copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+        guard let outputs = module.predict(image: copiedBufferPtr) else {
+            copiedBufferPtr.deallocate()
             return
         }
+        copiedBufferPtr.deallocate()
         let zippedResults = zip(labels.indices, outputs)
         let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
         var text = ""
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift
@@ -26,12 +26,18 @@ class ViewController: UIViewController {
         let image = UIImage(named: "image.png")!
         imageView.image = image
         let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
-        guard var pixelBuffer = resizedImage.normalized() else {
+        guard let pixelBuffer = resizedImage.normalized() else {
             return
         }
-        guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
+        // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+        // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+        let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+        copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+        guard let outputs = module.predict(image: copiedBufferPtr) else {
+            copiedBufferPtr.deallocate()
             return
         }
+        copiedBufferPtr.deallocate()
         let zippedResults = zip(labels.indices, outputs)
         let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
         var text = ""
diff --git a/HelloWorld/HelloWorld/HelloWorld/ViewController.swift b/HelloWorld/HelloWorld/HelloWorld/ViewController.swift
@@ -26,12 +26,18 @@ class ViewController: UIViewController {
         let image = UIImage(named: "image.png")!
         imageView.image = image
         let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
-        guard var pixelBuffer = resizedImage.normalized() else {
+        guard let pixelBuffer = resizedImage.normalized() else {
             return
         }
-        guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
+        // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+        // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+        let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+        copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+        guard let outputs = module.predict(image: copiedBufferPtr) else {
+            copiedBufferPtr.deallocate()
             return
         }
+        copiedBufferPtr.deallocate()
         let zippedResults = zip(labels.indices, outputs)
         let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
         var text = ""
diff --git a/ImageSegmentation/ImageSegmentation/ViewController.swift b/ImageSegmentation/ImageSegmentation/ViewController.swift
@@ -3,7 +3,7 @@ import UIKit
 class ViewController: UIViewController {
     @IBOutlet weak var imageView: UIImageView!
     @IBOutlet weak var btnSegment: UIButton!
-    
+
     private var imageName = "deeplab.jpg"
     private var image : UIImage?
     private let imageHelper = UIImageHelper()
@@ -17,7 +17,7 @@ class ViewController: UIViewController {
             fatalError("Can't find the model file!")
         }
     }()
-    
+
     override func viewDidLoad() {
         super.viewDidLoad()
 
@@ -30,22 +30,27 @@ class ViewController: UIViewController {
         btnSegment.isEnabled = false
         btnSegment.setTitle("Running the model...", for: .normal)
         let resizedImage = image!.resized(to: CGSize(width: 250, height: 250))
-        guard var pixelBuffer = resizedImage.normalized() else {
+        guard let pixelBuffer = resizedImage.normalized() else {
             return
         }
-                
+
         let w = Int32(resizedImage.size.width)
         let h = Int32(resizedImage.size.height)
         DispatchQueue.global().async {
-            let buffer = self.module.segment(image: UnsafeMutableRawPointer(&pixelBuffer), withWidth:w, withHeight: h)
+            // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+            // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+            let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+            copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+            let buffer = self.module.segment(image: copiedBufferPtr, withWidth:w, withHeight: h)
+            copiedBufferPtr.deallocate()
             DispatchQueue.main.async {
                 self.imageView.image = self.imageHelper.convertRGBBuffer(toUIImage: buffer , withWidth: w, withHeight: h)
                 self.btnSegment.isEnabled = true
                 self.btnSegment.setTitle("Segment", for: .normal)
             }
         }
     }
-    
+
     @IBAction func doRestart(_ sender: Any) {
         if imageName == "deeplab.jpg" {
             imageName = "dog.jpg"
@@ -55,6 +60,5 @@ class ViewController: UIViewController {
         }
         image = UIImage(named: imageName)!
         imageView.image = image
-    }    
+    }
 }
-
diff --git a/ObjectDetection/ObjectDetection/Live/LiveObjectDetectionViewController.swift b/ObjectDetection/ObjectDetection/Live/LiveObjectDetectionViewController.swift
@@ -17,43 +17,51 @@ class LiveObjectDetectionViewController: ViewController {
     private var cameraController = CameraController()
     private var imageViewLive =  UIImageView()
     private var inferencer = ObjectDetector()
-    
+
     override func viewDidLoad() {
         super.viewDidLoad()
         cameraController.configPreviewLayer(cameraView)
         imageViewLive.frame = CGRect(x: 0, y: 0, width: cameraView.frame.size.width, height: cameraView.frame.size.height)
         cameraView.addSubview(imageViewLive)
-        
+
         cameraController.videoCaptureCompletionBlock = { [weak self] buffer, error in
             guard let strongSelf = self else { return }
             if error != nil {
                 return
             }
-            guard var pixelBuffer = buffer else { return }
-            
+            guard let pixelBuffer = buffer else { return }
+
             let currentTimestamp = CACurrentMediaTime()
             if (currentTimestamp - strongSelf.prevTimestampMs) * 1000 <= strongSelf.delayMs { return }
             strongSelf.prevTimestampMs = currentTimestamp
+
+            // UnsafeMutablePointer() doesn't guarantee that the converted pointer points to the memory that is still being allocated
+            // So we create a new pointer and copy the &pixelBuffer's memory to where it points to
+            let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
+            copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
+
             let startTime = CACurrentMediaTime()
-            guard let outputs = self?.inferencer.module.detect(image: &pixelBuffer) else {
+            guard let outputs = self?.inferencer.module.detect(image: copiedBufferPtr) else {
+                copiedBufferPtr.deallocate()
                 return
             }
+            copiedBufferPtr.deallocate()
             let inferenceTime = CACurrentMediaTime() - startTime
-                
+
             DispatchQueue.main.async {
                 let ivScaleX : Double =  Double(strongSelf.imageViewLive.frame.size.width / CGFloat(PrePostProcessor.inputWidth))
                 let ivScaleY : Double = Double(strongSelf.imageViewLive.frame.size.height / CGFloat(PrePostProcessor.inputHeight))
 
                 let startX = Double((strongSelf.imageViewLive.frame.size.width - CGFloat(ivScaleX) * CGFloat(PrePostProcessor.inputWidth))/2)
                 let startY = Double((strongSelf.imageViewLive.frame.size.height -  CGFloat(ivScaleY) * CGFloat(PrePostProcessor.inputHeight))/2)
-                
+
                 let nmsPredictions = PrePostProcessor.outputsToNMSPredictions(outputs: outputs, imgScaleX: 1.0, imgScaleY: 1.0, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
 
                 PrePostProcessor.cleanDetection(imageView: strongSelf.imageViewLive)
                 strongSelf.indicator.isHidden = true
                 strongSelf.benchmarkLabel.isHidden = false
                 strongSelf.benchmarkLabel.text = String(format: "%.2fms", 1000*inferenceTime)
-                
+
                 PrePostProcessor.showDetection(imageView: strongSelf.imageViewLive, nmsPredictions: nmsPredictions, classes: strongSelf.inferencer.classes)
             }
         }
diff --git a/ObjectDetection/ObjectDetection/ViewController.swift b/ObjectDetection/ObjectDetection/ViewController.swift
diff --git a/PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift b/PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift
