[CoreML] CoreMLのサンプルモデルを使ってみる(YOLOv3)

class YOLOv3Input : MLFeatureProvider {

    /// 416x416 RGB image as color (kCVPixelFormatType_32BGRA) image buffer, 416 pixels wide by 416 pixels high
    var image: CVPixelBuffer

    /// This defines the radius of suppression. as optional double value
    var iouThreshold: Double? = nil

    /// Remove bounding boxes below this threshold (confidences should be nonnegative). as optional double value
    var confidenceThreshold: Double? = nil

class YOLOv3Output : MLFeatureProvider {

    /// Source provided by CoreML

    private let provider : MLFeatureProvider


    /// Confidence derived for each of the bounding boxes.  as multidimensional array of doubles
    lazy var confidence: MLMultiArray = {
        [unowned self] in return self.provider.featureValue(for: "confidence")!.multiArrayValue
    }()!

    /// Normalised coordiantes (relative to the image size) for each of the bounding boxes (x,y,w,h).  as multidimensional array of doubles
    lazy var coordinates: MLMultiArray = {
        [unowned self] in return self.provider.featureValue(for: "coordinates")!.multiArrayValue
    }()!

func CVPixelBufferCreate(_ allocator: CFAllocator?, 
                       _ width: Int, 
                       _ height: Int, 
                       _ pixelFormatType: OSType, 
                       _ pixelBufferAttributes: CFDictionary?, 
                       _ pixelBufferOut: UnsafeMutablePointer) 

struct ContentView: View {
    @State private var isDropTargeted = false
    @State var nsImage:NSImage = (NSImage(contentsOf: Bundle.main.url(forResource: "View", withExtension: "jpg")! )?.copy(size: NSSize(width: 416, height: 416))!)!
    
    // alert
    @State private var showingAlert = false
    @State private var alertTitle = ""
    @State private var alertMessage = ""
    var body: some View {
        VStack {
            Image(nsImage: self.nsImage)
                .resizable()
                .scaledToFit()
                .onDrop(of: [kUTTypeFileURL as String], isTargeted: $isDropTargeted, perform: self.processDroppedFile(provideres:))
            Button(action: {
                self.applyYOLOv3()
            }, label: {
                Text("apply YOLOv3")
            })
            .padding()
        }
        .alert(isPresented: $showingAlert) {
            Alert(title: Text(alertTitle), message: Text(alertMessage), dismissButton: .default(Text("OK")))
        }
        
    }
    
    func applyYOLOv3() {
        let model = YOLOv3()
        
        guard let pixelBuffer = self.nsImage.pixelBuffer() else {
            alertTitle = "Error!"
            alertMessage = "failed to create PixelBuffer"
            showingAlert = true
            return
        }
        
        let input = YOLOv3Input(image: pixelBuffer, iouThreshold: nil, confidenceThreshold: nil)

        do {
            let output = try model.prediction(input: input)
            alertTitle = "Success"
            alertMessage = ""
            print("confidence: \(output.confidence)")
            print("coordinate: \(output.coordinates)")
        } catch {
            alertTitle = "Error!"
            alertMessage = error.localizedDescription
        }
        showingAlert = true
    }
    
    func processDroppedFile(provideres:[NSItemProvider]) -> Bool {
        guard let provider = provideres.first else { return false }
        provider.loadItem(forTypeIdentifier: (kUTTypeFileURL as String), options: nil) { (urlData, error) in
            DispatchQueue.main.async {
                if let urlData = urlData as? Data {
                    let imageURL = NSURL(absoluteURLWithDataRepresentation: urlData, relativeTo: nil) as URL
                    if let localImage = NSImage(contentsOf: imageURL)?.copy(size: NSSize(width: 416, height: 416)) {
                        self.nsImage = localImage
                    }
                }
            }
        }
        return true
    }
}