LoginSignup
6

More than 5 years have passed since last update.

[iOS] ダウンロードしたCore ML modelを使用してVision Frameworkで画像解析する

Last updated at Posted at 2017-11-13

Core ML modelは、アプリのリソースとして組み込む以外に、ダウンロードした.mlmodelファイルをcompileして使うことができる。

Core ML model(.mlmodel)をダウンロードする

var urlSession: URLSession!

func downloadMlModel() {
    var url = FileManager.default.urls(for: .libraryDirectory, in: .userDomainMask).first!
    url.appendPathComponent("MobileNet.mlmodel", isDirectory: false)

    if FileManager.default.fileExists(atPath: url.path) {
        setupVison(url: url)
        setupAVCapture()
        return
    }

    urlSession = URLSession(configuration: URLSessionConfiguration.default,
                            delegate: self,
                            delegateQueue: nil)
    let request = URLRequest(url: URL(string: "https://docs-assets.developer.apple.com/coreml/models/MobileNet.mlmodel")!)
    let task = urlSession.downloadTask(with: request)
    task.resume()
}

//MARK: URLSessionDelegate, URLSessionTaskDelegate, URLSessionDownloadDelegate

public func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
    if let error = error {
        print("Failed to download CoreML model file, error=\(error.localizedDescription)")
    }
}

public func urlSession(_ session: URLSession, downloadTask: URLSessionDownloadTask, didFinishDownloadingTo location: URL) {
    var url = FileManager.default.urls(for: .libraryDirectory, in: .userDomainMask).first!
    url.appendPathComponent("MobileNet.mlmodel", isDirectory: false)

    do {
        try FileManager.default.moveItem(at: location, to: url)
    } catch {
        print("Failed to move CoreML model file, error=\(error.localizedDescription)")
        return
    }
    setupVison(url: url)
}

Core ML model(.mlmodel)をCompileしてVNCoreMLModelを生成する

var mlModel: VNCoreMLModel?
func setupVison(url: URL) {
    let compiledUrl: URL
    do {
        compiledUrl = try MLModel.compileModel(at: url)
    } catch {
        print("Failed to comlile Core ML model, error=\(error.localizedDescription)")
        return
    }

    let fileManager = FileManager.default
    let directory = try! fileManager.url(for: .applicationSupportDirectory,
                                                   in: .userDomainMask,
                                                   appropriateFor: compiledUrl,
                                                   create: true)
    let permanentUrl = directory.appendingPathComponent(compiledUrl.lastPathComponent)
    do {
        if fileManager.fileExists(atPath: permanentUrl.absoluteString) {
            _ = try fileManager.replaceItemAt(permanentUrl, withItemAt: compiledUrl)
        } else {
            try fileManager.copyItem(at: compiledUrl, to: permanentUrl)
        }
    } catch {
        print("Failed to copy compiled Core ML model, error=\(error.localizedDescription)")
        return
    }

    do {
        let model = try MLModel(contentsOf: permanentUrl)
        self.mlModel = try VNCoreMLModel(for: model)
    } catch {
        print("Failed to copy compiled Core ML model, error=\(error.localizedDescription)")
        return

    }
    setupAVCapture()
}

AVFoundationでカメラの入力をCaptureをする

func setupAVCapture() {
    captureSession = AVCaptureSession()

    if UIDevice.current.userInterfaceIdiom == .phone {
        captureSession.sessionPreset = .vga640x480
    } else {
        captureSession.sessionPreset = .photo
    }

    do {
        guard let device = AVCaptureDevice.default(for: .video) else {
            return
        }
        let deviceInput = try AVCaptureDeviceInput(device: device)
        if captureSession.canAddInput(deviceInput) {
            captureSession.addInput(deviceInput)
        }
    } catch {
        print("Failed to create capture device input, error=\(error.localizedDescription)")

        return
    }

    let videoDataOutput = AVCaptureVideoDataOutput()
    videoDataOutput.videoSettings = [kCVPixelBufferPixelFormatTypeKey as String: kCMPixelFormat_32BGRA]
    videoDataOutput.alwaysDiscardsLateVideoFrames = true
    videoDataOutput.setSampleBufferDelegate(self, queue: DispatchQueue(label: "VideoDataOutputQueue"))

    if captureSession.canAddOutput(videoDataOutput) {
        captureSession.addOutput(videoDataOutput)
    }

    let connection = videoDataOutput.connection(with: .video)
    connection?.isEnabled = true

    let previewLayer = AVCaptureVideoPreviewLayer(session: captureSession)
    previewLayer.backgroundColor = UIColor.black.cgColor
    previewLayer.videoGravity = .resizeAspect

    let rootLayer = previewView.layer
    rootLayer.masksToBounds = true
    previewLayer.frame = rootLayer.bounds
    rootLayer.addSublayer(previewLayer)

    captureSession.startRunning()
}

//MARK: AVCaptureVideoDataOutputSampleBufferDelegate

public func captureOutput(_ output: AVCaptureOutput, didOutput sampleBuffer: CMSampleBuffer, from connection: AVCaptureConnection) {
    let pixelBuffer: CVPixelBuffer = CMSampleBufferGetImageBuffer(sampleBuffer)!
    analyze(pixelBuffer: pixelBuffer)
}

Vision Frameworkで画像解析する

func analyze(pixelBuffer: CVPixelBuffer) {
    guard let mlModel = mlModel else {
        return
    }

    let request = VNCoreMLRequest(model: mlModel) { (request, error) in
        if let result: VNClassificationObservation = request.results?.first as? VNClassificationObservation {
            DispatchQueue.main.async {
                print("Result=\(result.identifier)")
                self.resultLabel.text = result.identifier
            }
        }
    }

    let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer, options: [:])
    do {
        try handler.perform([request])
    } catch {
        print("Failed to perform handler, error=\(error.localizedDescription)")
    }
}

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
6