Skip to content
This repository was archived by the owner on Aug 28, 2024. It is now read-only.

Commit 0d8774f

Browse files
husthychusthyc
committed
Issue fixed: dangling pointer that 'UnsafeMutablePointer()' casused made invalid model outputs
Issue fixed: dangling pointer that 'UnsafeMutablePointer()' casused made invalid model outputs
1 parent 0678fd7 commit 0d8774f

File tree

9 files changed

+87
-58
lines changed

9 files changed

+87
-58
lines changed

‎D2Go/D2Go/Live/LiveObjectDetectionViewController.swift

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,43 +17,47 @@ class LiveObjectDetectionViewController: ViewController {
1717
private var cameraController = CameraController()
1818
private var imageViewLive = UIImageView()
1919
private var inferencer = ObjectDetector()
20-
20+
2121
override func viewDidLoad() {
2222
super.viewDidLoad()
2323
cameraController.configPreviewLayer(cameraView)
2424
imageViewLive.frame = CGRect(x: 0, y: 0, width: cameraView.frame.size.width, height: cameraView.frame.size.height)
2525
cameraView.addSubview(imageViewLive)
26-
26+
2727
cameraController.videoCaptureCompletionBlock = { [weak self] buffer, error in
2828
guard let strongSelf = self else { return }
2929
if error != nil {
3030
return
3131
}
32-
guard var pixelBuffer = buffer else { return }
33-
32+
guard let pixelBuffer = buffer else { return }
33+
3434
let currentTimestamp = CACurrentMediaTime()
3535
if (currentTimestamp - strongSelf.prevTimestampMs) * 1000 <= strongSelf.delayMs { return }
3636
strongSelf.prevTimestampMs = currentTimestamp
3737
let startTime = CACurrentMediaTime()
38-
guard let outputs = self?.inferencer.module.detect(image: &pixelBuffer) else {
38+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
39+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
40+
guard let outputs = self?.inferencer.module.detect(image: copiedBufferPtr) else {
41+
copiedBufferPtr.deallocate()
3942
return
4043
}
44+
copiedBufferPtr.deallocate()
4145
let inferenceTime = CACurrentMediaTime() - startTime
42-
46+
4347
DispatchQueue.main.async {
4448
let ivScaleX : Double = Double(strongSelf.imageViewLive.frame.size.width / CGFloat(PrePostProcessor.inputWidth))
4549
let ivScaleY : Double = Double(strongSelf.imageViewLive.frame.size.height / CGFloat(PrePostProcessor.inputHeight))
4650

4751
let startX = Double((strongSelf.imageViewLive.frame.size.width - CGFloat(ivScaleX) * CGFloat(PrePostProcessor.inputWidth))/2)
4852
let startY = Double((strongSelf.imageViewLive.frame.size.height - CGFloat(ivScaleY) * CGFloat(PrePostProcessor.inputHeight))/2)
49-
53+
5054
let predictions = PrePostProcessor.outputsToPredictions(outputs: outputs, imgScaleX: 1.0, imgScaleY: 1.0, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
5155

5256
PrePostProcessor.cleanDetection(imageView: strongSelf.imageViewLive)
5357
strongSelf.indicator.isHidden = true
5458
strongSelf.benchmarkLabel.isHidden = false
5559
strongSelf.benchmarkLabel.text = String(format: "%.2fms", 1000*inferenceTime)
56-
60+
5761
PrePostProcessor.showDetection(imageView: strongSelf.imageViewLive, nmsPredictions: predictions, classes: strongSelf.inferencer.classes)
5862
}
5963
}

‎D2Go/D2Go/ViewController.swift

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,13 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
1010
@IBOutlet weak var imageView: UIImageView!
1111
@IBOutlet weak var btnRun: UIButton!
1212
@IBOutlet weak var btnNext: UIButton!
13-
13+
1414
private let testImages = ["test1.png", "test2.jpg", "test3.png"]
1515
private var imgIndex = 0
1616

1717
private var image : UIImage?
1818
private var inferencer = ObjectDetector()
19-
19+
2020
override func viewDidLoad() {
2121
super.viewDidLoad()
2222
image = UIImage(named: testImages[imgIndex])!
@@ -31,27 +31,31 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
3131
btnRun.setTitle("Running the model...", for: .normal)
3232

3333
let resizedImage = image!.resized(to: CGSize(width: CGFloat(PrePostProcessor.inputWidth), height: CGFloat(PrePostProcessor.inputHeight)))
34-
34+
3535
let imgScaleX = Double(image!.size.width / CGFloat(PrePostProcessor.inputWidth));
3636
let imgScaleY = Double(image!.size.height / CGFloat(PrePostProcessor.inputHeight));
37-
37+
3838
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))
3939
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))
4040

4141
let startX = Double((imageView.frame.size.width - CGFloat(ivScaleX) * imageView.image!.size.width)/2)
4242
let startY = Double((imageView.frame.size.height - CGFloat(ivScaleY) * imageView.image!.size.height)/2)
4343

44-
guard var pixelBuffer = resizedImage.normalized() else {
44+
guard let pixelBuffer = resizedImage.normalized() else {
4545
return
4646
}
47-
47+
4848
DispatchQueue.global().async {
49-
guard let outputs = self.inferencer.module.detect(image: &pixelBuffer) else {
49+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
50+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
51+
guard let outputs = self.inferencer.module.detect(image: copiedBufferPtr) else {
52+
copiedBufferPtr.deallocate()
5053
return
5154
}
52-
55+
copiedBufferPtr.deallocate()
56+
5357
let predictions = PrePostProcessor.outputsToPredictions(outputs: outputs, imgScaleX: imgScaleX, imgScaleY: imgScaleY, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
54-
58+
5559
DispatchQueue.main.async {
5660
PrePostProcessor.showDetection(imageView: self.imageView, nmsPredictions: predictions, classes: self.inferencer.classes)
5761
self.btnRun.isEnabled = true
@@ -75,7 +79,7 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
7579
imagePickerController.sourceType = .photoLibrary
7680
self.present(imagePickerController, animated: true, completion: nil)
7781
}
78-
82+
7983
@IBAction func cameraTapped(_ sender: Any) {
8084
PrePostProcessor.cleanDetection(imageView: imageView)
8185
if UIImagePickerController.isSourceTypeAvailable(.camera) {
@@ -85,13 +89,11 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
8589
self.present(imagePickerController, animated: true, completion: nil)
8690
}
8791
}
88-
92+
8993
func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
9094
image = info[UIImagePickerController.InfoKey.originalImage] as? UIImage
9195
image = image!.resized(to: CGSize(width: CGFloat(PrePostProcessor.inputWidth), height: CGFloat(PrePostProcessor.inputHeight)*image!.size.height/image!.size.width))
9296
imageView.image = image
9397
self.dismiss(animated: true, completion: nil)
9498
}
9599
}
96-
97-

‎HelloWorld-CoreML/HelloWorld/HelloWorld/ViewController.swift

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,16 @@ class ViewController: UIViewController {
2626
let image = UIImage(named: "image.png")!
2727
imageView.image = image
2828
let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
29-
guard var pixelBuffer = resizedImage.normalized() else {
29+
guard let pixelBuffer = resizedImage.normalized() else {
3030
return
3131
}
32-
guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
32+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
33+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
34+
guard let outputs = module.predict(image: copiedBufferPtr) else {
35+
copiedBufferPtr.deallocate()
3336
return
3437
}
38+
copiedBufferPtr.deallocate()
3539
let zippedResults = zip(labels.indices, outputs)
3640
let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
3741
var text = ""

‎HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,16 @@ class ViewController: UIViewController {
2626
let image = UIImage(named: "image.png")!
2727
imageView.image = image
2828
let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
29-
guard var pixelBuffer = resizedImage.normalized() else {
29+
guard let pixelBuffer = resizedImage.normalized() else {
3030
return
3131
}
32-
guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
32+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
33+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
34+
guard let outputs = module.predict(image: copiedBufferPtr) else {
35+
copiedBufferPtr.deallocate()
3336
return
3437
}
38+
copiedBufferPtr.deallocate()
3539
let zippedResults = zip(labels.indices, outputs)
3640
let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
3741
var text = ""

‎HelloWorld/HelloWorld/HelloWorld/ViewController.swift

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,12 +26,16 @@ class ViewController: UIViewController {
2626
let image = UIImage(named: "image.png")!
2727
imageView.image = image
2828
let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
29-
guard var pixelBuffer = resizedImage.normalized() else {
29+
guard let pixelBuffer = resizedImage.normalized() else {
3030
return
3131
}
32-
guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
32+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
33+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
34+
guard let outputs = module.predict(image: copiedBufferPtr) else {
35+
copiedBufferPtr.deallocate()
3336
return
3437
}
38+
copiedBufferPtr.deallocate()
3539
let zippedResults = zip(labels.indices, outputs)
3640
let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
3741
var text = ""

‎ImageSegmentation/ImageSegmentation/ViewController.swift

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ import UIKit
33
class ViewController: UIViewController {
44
@IBOutlet weak var imageView: UIImageView!
55
@IBOutlet weak var btnSegment: UIButton!
6-
6+
77
private var imageName = "deeplab.jpg"
88
private var image : UIImage?
99
private let imageHelper = UIImageHelper()
@@ -17,7 +17,7 @@ class ViewController: UIViewController {
1717
fatalError("Can't find the model file!")
1818
}
1919
}()
20-
20+
2121
override func viewDidLoad() {
2222
super.viewDidLoad()
2323

@@ -30,22 +30,25 @@ class ViewController: UIViewController {
3030
btnSegment.isEnabled = false
3131
btnSegment.setTitle("Running the model...", for: .normal)
3232
let resizedImage = image!.resized(to: CGSize(width: 250, height: 250))
33-
guard var pixelBuffer = resizedImage.normalized() else {
33+
guard let pixelBuffer = resizedImage.normalized() else {
3434
return
3535
}
36-
36+
3737
let w = Int32(resizedImage.size.width)
3838
let h = Int32(resizedImage.size.height)
3939
DispatchQueue.global().async {
40-
let buffer = self.module.segment(image: UnsafeMutableRawPointer(&pixelBuffer), withWidth:w, withHeight: h)
40+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
41+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
42+
let buffer = self.module.segment(image: copiedBufferPtr, withWidth:w, withHeight: h)
43+
copiedBufferPtr.deallocate()
4144
DispatchQueue.main.async {
4245
self.imageView.image = self.imageHelper.convertRGBBuffer(toUIImage: buffer , withWidth: w, withHeight: h)
4346
self.btnSegment.isEnabled = true
4447
self.btnSegment.setTitle("Segment", for: .normal)
4548
}
4649
}
4750
}
48-
51+
4952
@IBAction func doRestart(_ sender: Any) {
5053
if imageName == "deeplab.jpg" {
5154
imageName = "dog.jpg"
@@ -55,6 +58,5 @@ class ViewController: UIViewController {
5558
}
5659
image = UIImage(named: imageName)!
5760
imageView.image = image
58-
}
61+
}
5962
}
60-

‎ObjectDetection/ObjectDetection/Live/LiveObjectDetectionViewController.swift

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -17,43 +17,47 @@ class LiveObjectDetectionViewController: ViewController {
1717
private var cameraController = CameraController()
1818
private var imageViewLive = UIImageView()
1919
private var inferencer = ObjectDetector()
20-
20+
2121
override func viewDidLoad() {
2222
super.viewDidLoad()
2323
cameraController.configPreviewLayer(cameraView)
2424
imageViewLive.frame = CGRect(x: 0, y: 0, width: cameraView.frame.size.width, height: cameraView.frame.size.height)
2525
cameraView.addSubview(imageViewLive)
26-
26+
2727
cameraController.videoCaptureCompletionBlock = { [weak self] buffer, error in
2828
guard let strongSelf = self else { return }
2929
if error != nil {
3030
return
3131
}
32-
guard var pixelBuffer = buffer else { return }
33-
32+
guard let pixelBuffer = buffer else { return }
33+
3434
let currentTimestamp = CACurrentMediaTime()
3535
if (currentTimestamp - strongSelf.prevTimestampMs) * 1000 <= strongSelf.delayMs { return }
3636
strongSelf.prevTimestampMs = currentTimestamp
3737
let startTime = CACurrentMediaTime()
38-
guard let outputs = self?.inferencer.module.detect(image: &pixelBuffer) else {
38+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
39+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
40+
guard let outputs = self?.inferencer.module.detect(image: copiedBufferPtr) else {
41+
copiedBufferPtr.deallocate()
3942
return
4043
}
44+
copiedBufferPtr.deallocate()
4145
let inferenceTime = CACurrentMediaTime() - startTime
42-
46+
4347
DispatchQueue.main.async {
4448
let ivScaleX : Double = Double(strongSelf.imageViewLive.frame.size.width / CGFloat(PrePostProcessor.inputWidth))
4549
let ivScaleY : Double = Double(strongSelf.imageViewLive.frame.size.height / CGFloat(PrePostProcessor.inputHeight))
4650

4751
let startX = Double((strongSelf.imageViewLive.frame.size.width - CGFloat(ivScaleX) * CGFloat(PrePostProcessor.inputWidth))/2)
4852
let startY = Double((strongSelf.imageViewLive.frame.size.height - CGFloat(ivScaleY) * CGFloat(PrePostProcessor.inputHeight))/2)
49-
53+
5054
let nmsPredictions = PrePostProcessor.outputsToNMSPredictions(outputs: outputs, imgScaleX: 1.0, imgScaleY: 1.0, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
5155

5256
PrePostProcessor.cleanDetection(imageView: strongSelf.imageViewLive)
5357
strongSelf.indicator.isHidden = true
5458
strongSelf.benchmarkLabel.isHidden = false
5559
strongSelf.benchmarkLabel.text = String(format: "%.2fms", 1000*inferenceTime)
56-
60+
5761
PrePostProcessor.showDetection(imageView: strongSelf.imageViewLive, nmsPredictions: nmsPredictions, classes: strongSelf.inferencer.classes)
5862
}
5963
}

‎ObjectDetection/ObjectDetection/ViewController.swift

Lines changed: 15 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,13 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
1010
@IBOutlet weak var imageView: UIImageView!
1111
@IBOutlet weak var btnRun: UIButton!
1212
@IBOutlet weak var btnNext: UIButton!
13-
13+
1414
private let testImages = ["test1.png", "test2.jpg", "test3.png"]
1515
private var imgIndex = 0
1616

1717
private var image : UIImage?
1818
private var inferencer = ObjectDetector()
19-
19+
2020
override func viewDidLoad() {
2121
super.viewDidLoad()
2222
image = UIImage(named: testImages[imgIndex])!
@@ -31,27 +31,31 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
3131
btnRun.setTitle("Running the model...", for: .normal)
3232

3333
let resizedImage = image!.resized(to: CGSize(width: CGFloat(PrePostProcessor.inputWidth), height: CGFloat(PrePostProcessor.inputHeight)))
34-
34+
3535
let imgScaleX = Double(image!.size.width / CGFloat(PrePostProcessor.inputWidth));
3636
let imgScaleY = Double(image!.size.height / CGFloat(PrePostProcessor.inputHeight));
37-
37+
3838
let ivScaleX : Double = (image!.size.width > image!.size.height ? Double(imageView.frame.size.width / image!.size.width) : Double(imageView.frame.size.height / image!.size.height))
3939
let ivScaleY : Double = (image!.size.height > image!.size.width ? Double(imageView.frame.size.height / image!.size.height) : Double(imageView.frame.size.width / image!.size.width))
4040

4141
let startX = Double((imageView.frame.size.width - CGFloat(ivScaleX) * image!.size.width)/2)
4242
let startY = Double((imageView.frame.size.height - CGFloat(ivScaleY) * image!.size.height)/2)
4343

44-
guard var pixelBuffer = resizedImage.normalized() else {
44+
guard let pixelBuffer = resizedImage.normalized() else {
4545
return
4646
}
47-
47+
4848
DispatchQueue.global().async {
49-
guard let outputs = self.inferencer.module.detect(image: &pixelBuffer) else {
49+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: pixelBuffer.count)
50+
copiedBufferPtr.initialize(from: pixelBuffer, count: pixelBuffer.count)
51+
guard let outputs = self.inferencer.module.detect(image: copiedBufferPtr) else {
52+
copiedBufferPtr.deallocate()
5053
return
5154
}
52-
55+
copiedBufferPtr.deallocate()
56+
5357
let nmsPredictions = PrePostProcessor.outputsToNMSPredictions(outputs: outputs, imgScaleX: imgScaleX, imgScaleY: imgScaleY, ivScaleX: ivScaleX, ivScaleY: ivScaleY, startX: startX, startY: startY)
54-
58+
5559
DispatchQueue.main.async {
5660
PrePostProcessor.showDetection(imageView: self.imageView, nmsPredictions: nmsPredictions, classes: self.inferencer.classes)
5761
self.btnRun.isEnabled = true
@@ -75,7 +79,7 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
7579
imagePickerController.sourceType = .photoLibrary
7680
self.present(imagePickerController, animated: true, completion: nil)
7781
}
78-
82+
7983
@IBAction func cameraTapped(_ sender: Any) {
8084
PrePostProcessor.cleanDetection(imageView: imageView)
8185
if UIImagePickerController.isSourceTypeAvailable(.camera) {
@@ -85,13 +89,11 @@ class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavig
8589
self.present(imagePickerController, animated: true, completion: nil)
8690
}
8791
}
88-
92+
8993
func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
9094
image = info[UIImagePickerController.InfoKey.originalImage] as? UIImage
9195
image = image!.resized(to: CGSize(width: CGFloat(PrePostProcessor.inputWidth), height: CGFloat(PrePostProcessor.inputHeight)*image!.size.height/image!.size.width))
9296
imageView.image = image
9397
self.dismiss(animated: true, completion: nil)
9498
}
9599
}
96-
97-

‎PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,10 +26,13 @@ class ImagePredictor: Predictor {
2626
}
2727
isRunning = true
2828
let startTime = CACurrentMediaTime()
29-
var tensorBuffer = buffer;
30-
guard let outputs = module.predict(image: UnsafeMutableRawPointer(&tensorBuffer)) else {
29+
let copiedBufferPtr = UnsafeMutablePointer<Float>.allocate(capacity: buffer.count)
30+
copiedBufferPtr.initialize(from: buffer, count: buffer.count)
31+
guard let outputs = module.predict(image: copiedBufferPtr) else {
32+
copiedBufferPtr.deallocate()
3133
throw PredictorError.invalidInputTensor
3234
}
35+
copiedBufferPtr.deallocate()
3336
isRunning = false
3437
let inferenceTime = (CACurrentMediaTime() - startTime) * 1000
3538
let results = topK(scores: outputs, labels: labels, count: resultCount)

0 commit comments

Comments
 (0)