example.swift

import UIKit

class ViewController: UIViewController {

    @IBOutlet weak var imgView: UIImageView!

    override func viewDidLoad() {
        super.viewDidLoad()
        loadSample()
    }

    @IBAction func onClickInput(_ sender: Any) {
        loadSample()
    }

    private func loadSample() {
        let image = UIImage(named: "sample")
        imgView.image = image
    }

    @IBAction func onClickMask32(_ sender: Any) {
        pred(predFn: { pixelBuffer in
            let model = unet()
            let input = unetInput(_0: pixelBuffer)
            let pred = try! model.prediction(input: input)
            let resized = resizePixelBuffer(pred._595, width: 600, height: 800)
            return resized!
        })
    }

    @IBAction func onClickMask16(_ sender: Any) {
        pred(predFn: { pixelBuffer in
            let model = unet_16_1()
            let input = unet_16_1Input(_0: pixelBuffer)
            let pred = try! model.prediction(input: input)
            let resized = resizePixelBuffer(pred._595, width: 600, height: 800)
            return resized!
        })
    }

    @IBAction func onClickMask8(_ sender: Any) {
        pred(predFn: { pixelBuffer in
            let model = unet_8_1()
            let input = unet_8_1Input(_0: pixelBuffer)
            let pred = try! model.prediction(input: input)
            let resized = resizePixelBuffer(pred._595, width: 600, height: 800)
            return resized!
        })
    }

    private func pred(predFn: (CVPixelBuffer) -> CVPixelBuffer) {
        let image = UIImage(named: "sample")
        let inBuff = image!.pixelBuffer(width: 224, height: 224)
        let outBuff = predFn(inBuff!)
        imgView.image = UIImage(pixelBuffer: outBuff)
    }

}