Skip to content
Snippets Groups Projects
ObjectDetection.swift 2.18 KiB
Newer Older
  • Learn to ignore specific revisions
  • ufuktugay's avatar
    ufuktugay committed
    //
    //  ObjectDetection.swift
    //  Objekterkennung in AR mit ARKit und CoreML
    //
    //  Created by Ufuk Kaan on 18.06.21.
    //
    
    import CoreML
    import Vision
    import SceneKit
    
    class ObjectDetection{
        
        struct Detection_Object {
            let classification: String
            let confidence: Float
            let boundingBox: CGRect
        }
        
    
    ufuktugay's avatar
    ufuktugay committed
        //Object Detection Model initialisieren
    
    ufuktugay's avatar
    ufuktugay committed
        var mblock_model = try! VNCoreMLModel(for: mblockv1().model)
    
    ufuktugay's avatar
    ufuktugay committed
        
        //Detection Request
    
    ufuktugay's avatar
    ufuktugay committed
        lazy var detection_request: VNCoreMLRequest = { return VNCoreMLRequest(model: mblock_model,completionHandler: self.completionHandler)}()
       
    
    ufuktugay's avatar
    ufuktugay committed
        //Array aus erkannten Objekten für Ausgabe/Controller
    
    ufuktugay's avatar
    ufuktugay committed
        var detection_object_array =  [Detection_Object]()
        
    
    ufuktugay's avatar
    ufuktugay committed
        
    
    ufuktugay's avatar
    ufuktugay committed
        func detectObjects(on request: CVPixelBuffer) {
            do {
                try VNImageRequestHandler(cvPixelBuffer: request).perform([detection_request])
            } catch {
                print("VNImageRequestHandler failed")
            }
        }
        
        func completionHandler(_ request: VNRequest?, error: Error?) {
            
            guard let results = request?.results as? [VNRecognizedObjectObservation] else {
                return
            }
            for observation in results {
                if observation.confidence > 0.5 {
    //                print("Observation: \(observation.confidence) \nLabel: \(observation.labels.first!.identifier) \nLocation Center: (\(observation.boundingBox.midX), \(observation.boundingBox.midY))")
                    
                    let responseItem =  Detection_Object.init(classification: observation.labels.first!.identifier, confidence: observation.confidence,boundingBox: observation.boundingBox)
                    
                    
    //                print("Observation: \(responseItem.confidence) \nLabel: \(responseItem.classification) \nLocation Center:
    //                (\(responseItem.boundingBox.midX), \(responseItem.boundingBox.midY))")
                    
                    
                    detection_object_array.append(responseItem)
                }
                else {
                    print("Kein Objekt im Bild erkannt")
                }
            }
            print("\n------------------------------------------------------------------------------------------\n")
    
        }
    }