Swift + SwiftUI原生iOS开发 开发笔记2 – 使用YOLOv3n模型对画面进行识别

Swift + SwiftUI原生iOS开发 开发笔记1 – 实现摄像头调用和拍摄图片

背景

我的毕业设计是《Development of an AI-Powered Mobile Application for Animal Identification and Information》，开发一个可以识别动物的移动软件，在跟导师沟通后打算用Yolo作为深度学习识别模型，移动端软件使用原生的iOS开发，即在XCode上使用Swift+SwiftUI进行iOS开发，原本就一直对iOS开发很感兴趣，借这次机会好好学习一下原生的iOS开发。

开发

在实现摄像头调用和拍摄图片的功能后，接下来要对获取的每一帧画面使用深度学习模型进行识别，提取出画面中已知的物体，并在画面上标注出来。

导入模型

由于我还没有训练自己的模型，所以我借助苹果官网实现CoreML例子中的YOLOv3模型，作为这次识别的模型。

苹果官网CoreML模型介绍

我找到这个模型的时候才知道原来苹果官网有实时拍摄中识别对象的例子，早知道就不用折腾那么久了，对着学就好了，不过折腾这么久，也确实印象更深刻就是了。

在Xcode中，如果导入mlmodel文件，Xcode会为其生成好配置文件，包括模型对象、模型输入格式、模型输出格式。

导入mlmodel文件后生成的信息

在生成好配置文件后，就可以在需要的地方调用Yolov3定义并使用模型了。

let model3: YOLOv3
do {
    let configuration = MLModelConfiguration()
    configuration.computeUnits = .cpuAndGPU
    self.model3 = try YOLOv3(configuration: configuration)
    print("Model successfully loaded with GPU support.")
} catch {
   fatalError("Failed to load the CoreML model with GPU support: \(error.localizedDescription)")
}

声明并使用模型

在导入模型后，就可以新建一个文件专门管理模型，并定义模型的使用逻辑，接受传入的图片，并传出识别结果。

class PredictionManager {
    private let model3: YOLOv3
    
    init() {
        do {
            let configuration = MLModelConfiguration()
            configuration.computeUnits = .cpuAndGPU
            self.model3 = try YOLOv3(configuration: configuration)
            print("Model successfully loaded with GPU support.")
        } catch {
            fatalError("Failed to load the CoreML model with GPU support: \(error.localizedDescription)")
        }
        
    }
    
    func predict(image: UIImage, completion: @escaping ([PredictObject]?) -> Void) {}
    
    func getObject3(predictResult: YOLOv3Output) -> PredictObject {}
}

为了更好的方便视图对结果进行使用，在这里也定义了一个用于存放结果的结构体。

struct PredictObject {
    var xCenter: Float = 0
    var yCenter: Float = 0
    var width: Float = 0
    var height: Float = 0
    var classId: Int = -1
    var confidence: Float = -1
}

在predict函数中，传入图像并对图像进行调整大小等预处理，并调用模型进行识别。

let targetSize = CGSize(width: 416, height: 416) // 根据模型输入要求设置
if let resizedImage = image.resize(to: targetSize),
let pixelBuffer = resizedImage.toCVPixelBuffer() {
let yolov3Input: YOLOv3Input = YOLOv3Input(image: pixelBuffer)
do {
     let results = try self.model3.prediction(input: yolov3Input)
     if (results.coordinates.count != 0) {
           let detectedObject: PredictObject = self.getObject3(predictResult: results)                            
           if(detectedObject.confidence > 0.5) {
                  detectedObjects.append(detectedObject)
           }
     }
} catch {
     print("Predict ERROR")
}

获取识别结果后，在再getObject3函数中具体提炼出结果并构成一个视图可以直接使用的识别结果。

func getObject3(predictResult: YOLOv3Output) -> PredictObject {
   var detectedObject: PredictObject = PredictObject()
        
   let coordinates = predictResult.coordinates

   detectedObject.xCenter = 1 - coordinates[1].floatValue
   detectedObject.yCenter = coordinates[0].floatValue
   detectedObject.width = coordinates[3].floatValue
   detectedObject.height = coordinates[2].floatValue
        
   var maxClassConfidence: Float = 0
   var maxClassId: Int = -1
                    
   for i in 0..<80 {
            
       if (predictResult.confidence[i].floatValue > maxClassConfidence) {
           maxClassConfidence = predictResult.confidence[i].floatValue
           maxClassId = i
       }
   }
        
    detectedObject.classId = maxClassId
    detectedObject.confidence = Float(maxClassConfidence)
    return detectedObject
}

在这里需要注意的点是，经过测试，模型直接返回的虽然是xywh格式，但是默认的原点位置和SwiftUI默认的原点位置有出入，好像逆时针转的第一个顶点才是模型认为的原点，所以返回的xywh无法直接使用，需要进行处理。

至此，模型的定义结束，接下来只要在获取图片的地方，声明一个模型管理，再调用predict函数对图片进行识别就可以得到结果了。

模型识别

在开发笔记1中实现的摄像头实时调用，我预留了一个图片获取的coordinate，可以在其中对获取的图像进行处理。

func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection){
    autoreleasepool {
        guard let imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer) else {
            return
        }
        let ciImage = CIImage(cvImageBuffer: imageBuffer)
        let context = CIContext()
        if let cgImage = context.createCGImage(ciImage, from: ciImage.extent) {
            let image = UIImage(cgImage: cgImage)
            predictionManager.predict(image: image) { results in
                if let results = results, !results.isEmpty {
                    for result in results {
                        print("Detected: \(result.classId) - Confidence: \(result.confidence)")
                        print("x: \(result.xCenter) - y: \(result.yCenter)")
                        print("w: \(result.width) - h: \(result.height)")
                        self.parent.predictObject = result
                    }
                } else {
                    self.parent.predictObject = PredictObject(xCenter: 0, yCenter: 0, width: 0, height: 0, classId: -1, confidence: 0)
                }
            }
        }
    }
}

在调用predict函数后，对获取的结果进行判断，如果有内容就输出结果，并通过父类定义的Binding对象与视图界面的State变量进行传输数据。

结果展示

在视图文件中定义一个State修饰的变量用于存放识别结果。

@State private var predictObject: PredictObject? = nil
并设计一个视图结构体用于在屏幕上标出对应的方框位置


struct BoundingBoxView: View {

    var predictObject: PredictObject
    
    var width: CGFloat = UIScreen.main.bounds.width
    var height: CGFloat = UIScreen.main.bounds.height
    
    var body: some View {
        
        let w = CGFloat(self.predictObject.width) * width
        let h = CGFloat(self.predictObject.height) * height
        let x = CGFloat(self.predictObject.xCenter) * width
        let y = CGFloat(self.predictObject.yCenter) * height
       
        GeometryReader { geometry in
            ZStack {
                if(self.predictObject.classId != -1) {
                    Text("\(LabelList[self.predictObject.classId]): \(String(format: "%.2f", self.predictObject.confidence))")
                        .position(x: x, y: y - h / 2)
                    // 显示矩形框
                    Rectangle()
                        .stroke(Color.green, lineWidth: 2)  // 绿色边框
                        .frame(width: w, height: h)
                        .position(x: x, y: y)
                }
            }
            .frame(width: geometry.size.width, height: geometry.size.height)
        }
    }
}

这样在视图中调用BoudingBoxView并传入识别的结果，就可以在屏幕上显示出识别的物体的位置和对应的名字了。

至此调用模型对摄像头获取的内容进行识别成功结束。

总结

在加入模型识别后本项目基本功能已实现，接下来就要将侧重点放在自己的模型训练上了，构建一个识别动物的模型，并想办法导出自己训练的模型为mlmodel格式，并且要想办法调整输出格式为我们可以使用的格式。