问ARKit:跟踪VisonCoreML检测到的对象

EN

@IBOutlet var sceneView: ARSCNView!
private var viewportSize: CGSize!
private var previousAnchor: ARAnchor?
private var trackingNode: SCNNode!
lazy var objectDetectionRequest: VNCoreMLRequest = {
    do {
        let model = try VNCoreMLModel(for: yolov5s(configuration: MLModelConfiguration()).model)
        let request = VNCoreMLRequest(model: model) { [weak self] request, error in
            self?.processDetections(for: request, error: error)
        }
        return request
    } catch {
        fatalError("Failed to load Vision ML model.")
    }
}()

func renderer(_ renderer: SCNSceneRenderer, willRenderScene scene: SCNScene, atTime time: TimeInterval) {
    guard let capturedImage = sceneView.session.currentFrame?.capturedImage
    else { return }
    
    let imageRequestHandler = VNImageRequestHandler(cvPixelBuffer: capturedImage, orientation: .leftMirrored, options: [:])
    
    do {
        try imageRequestHandler.perform([objectDetectionRequest])
    } catch {
        print("Failed to perform image request.")
    }
}

func processDetections(for request: VNRequest, error: Error?) {
    guard error == nil else {
        print("Object detection error: \(error!.localizedDescription)")
        return
    }
    
    guard let results = request.results else { return }
    
    for observation in results where observation is VNRecognizedObjectObservation {
        
        let objectObservation = observation as! VNRecognizedObjectObservation
        let topLabelObservation = objectObservation.labels.first
        print(topLabelObservation!.identifier + " " + "\(Int(topLabelObservation!.confidence * 100))%")
        guard recognisedObject(topLabelObservation!.identifier) && topLabelObservation!.confidence > 0.9
            else { continue }

        let rect = VNImageRectForNormalizedRect(
            objectObservation.boundingBox,
            Int(self.sceneView.bounds.width),
            Int(self.sceneView.bounds.height))
        
        let midPoint = CGPoint(x: rect.midX, y: rect.midY)
        
        let raycastQuery = self.sceneView.raycastQuery(from: midPoint,
                                                  allowing: .estimatedPlane,
                                                  alignment: .any)
        let raycastArray = self.sceneView.session.raycast(raycastQuery!)
        guard let raycastResult = raycastArray.first else { return }
        
        let position = SCNVector3(raycastResult.worldTransform.columns.3.x,
                                raycastResult.worldTransform.columns.3.y,
                                raycastResult.worldTransform.columns.3.z)
        
        if let _ = trackingNode {
            trackingNode!.worldPosition = position
        } else {
            trackingNode = createNode()
            trackingNode!.worldPosition = position
            self.sceneView.scene.rootNode.addChildNode(trackingNode!)
        }
    }
}

private func recognisedObject(_ identifier: String) -> Bool {
    return identifier == "remote" || identifier == "mouse"
}

private func createNode() -> SCNNode {
    let sphereNode = SCNNode(geometry: SCNSphere(radius: 0.01))
    sphereNode.geometry?.firstMaterial?.diffuse.contents = UIColor.purple
    
    return sphereNode
}

private func loadSession() {
    let configuration = ARWorldTrackingConfiguration()
    configuration.planeDetection = []
    sceneView.session.run(configuration)
}


override func viewDidLoad() {
    super.viewDidLoad()
    sceneView.delegate = self
    viewportSize = sceneView.frame.size
}

override func viewWillAppear(_ animated: Bool) {
    super.viewWillAppear(animated)
    loadSession()
}

override func viewWillDisappear(_ animated: Bool) {
    super.viewWillDisappear(animated)
    sceneView.session.pause()
}