Skip to content
Snippets Groups Projects
ObjectDetection.swift 2.16 KiB
Newer Older
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
    var mblock_model = try! VNCoreMLModel(for: mblockv2().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))")
ufuktugay's avatar
ufuktugay committed
                
                
                detection_object_array.append(responseItem)
            }
            else {
                print("Kein Objekt im Bild erkannt")
            }
        }
        print("\n------------------------------------------------------------------------------------------\n")

    }
}