prediction.swift

private let classifier = Resnet50()

func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    let imageBuffer = CMSampleBufferGetImageBuffer(sampleBuffer)
    let rescaled = rescaledImageRectangle(createUIImage(from: imageBuffer), dimension: inputImageScale)
    let rescaledBuffer = createCVPixelBuffer(from: rescaled)
        
    guard let prediction = try? classifier.prediction(image: rescaledBuffer!) else { return }
    DispatchQueue.main.async {
        self.predictionLabel.text = prediction.classLabel
    }
}