mutating func run()

in SemanticSegmentationSample/SemanticSegmentationCLI/MainCommand.swift [29:84]


    mutating func run() async throws {
        // Compile and load the model
        let config = MLModelConfiguration()
        config.computeUnits = .cpuAndGPU
        let compiledURL = try await MLModel.compileModel(at: URL(filePath: model))
        let model = try MLModel(contentsOf: compiledURL, configuration: config)
        let postProcessor = try DETRPostProcessor(model: model)

        // Load the input image
        guard let inputImage = CIImage(contentsOf: URL(filePath: input)) else {
            print("Failed to load image.")
            throw ExitCode(EXIT_FAILURE)
        }
        print("Original image size \(inputImage.extent)")

        // Resize the image to match the model's expected input
        let resizedImage = inputImage.resized(to: targetSize)

        // Convert to a pixel buffer
        guard let pixelBuffer = context.render(resizedImage, pixelFormat: kCVPixelFormatType_32ARGB) else {
            print("Failed to create pixel buffer for input image.")
            throw ExitCode(EXIT_FAILURE)
        }

        // Execute the model
        let clock = ContinuousClock()
        let start = clock.now
        let featureProvider = try MLDictionaryFeatureProvider(dictionary: ["image": pixelBuffer])
        let result = try await model.prediction(from: featureProvider)
        guard let semanticPredictions = result.featureValue(for: "semanticPredictions")?.shapedArrayValue(of: Int32.self) else {
            print("The model did not return a 'semanticPredictions' output feature.")
            throw ExitCode(EXIT_FAILURE)
        }
        let duration = clock.now - start
        print("Model inference took \(duration.formatted(.units(allowed: [.seconds, .milliseconds])))")

        guard let semanticImage = try? postProcessor.semanticImage(semanticPredictions: semanticPredictions) else {
            print("Error post-processing semanticPredictions")
            throw ExitCode(EXIT_FAILURE)
        }

        // Undo the scale to match the original image size
        // TODO: Bilinear?
        let outputImage = semanticImage.resized(to: CGSize(width: inputImage.extent.width, height: inputImage.extent.height))
        // Save mask if we need to
        if let mask = mask {
            context.writePNG(outputImage, to: URL(filePath: mask))
        }

        // Display mask over original
        guard let outputImage = outputImage.withAlpha(0.5)?.composited(over: inputImage) else {
            print("Failed to blend mask.")
            throw ExitCode(EXIT_FAILURE)
        }
        context.writePNG(outputImage, to: URL(filePath: output))
    }