DeepLabV3を使って写真の背景を抽出し白く塗り潰したい

前提・実現したいこと

取り込んだ写真の背景を白くしたいです。
下記が目標です。（photoroomというアプリを使用しました）

色々調べたところ、Core MLモデルのDeeplabV3を使用して背景透過処理をされている方の記事を発見しました。
https://tech.basicinc.jp/articles/206

この方はmacOSのアプリのために使っていたため、私の方ではNSImageの部分をUIImageに書き換えました。
また、UIImageをCVPixelBufferに変換する関数（Buffer）は
https://stackoverrun.com/ja/q/12215731
の記事を参考にしました。

自分なりにコードを書いてみましたが、コメントアウト部分の背景を白くする記述が分からない状況です。

該当のソースコード

swift
1//略
2
3private func backgroundToWhite(brightnessValue: Float, contrastValue: Float, saturationValue: Float, isUpdatedSegmentation: Bool = false) {
4        guard let img = sourceImageView.image?.cropping(to: selectedRect!) else {
5            assertionFailure()
6            return
7        }
8        //cropImageViewにimgを格納しておく（比較するため）
9        cropImageView.image = img
10        //写真サイズの表記をトリミングされたimgのサイズに変更
11        imageSizeLabel.text = "(Int(img.size.width)) x (Int(img.size.height))"
12        
13        //ここから本処理です。
14        let image = buffer(from:img)
15
16        let model = DeepLabV3()
17        guard let output = try? model.prediction(image: image!) else {
18            fatalError("Error")
19        }
20
21        for y in 0..<Int(img.size.height) {
22            for x in 0..<Int(img.size.width) {
23                let getY = Int(img.size.width) - y - 1
24                let index: [NSNumber] = [NSNumber(value: getY), NSNumber(value: x)]
25                let flag = output.semanticPredictions[index]
26                if flag == 0 {
27                    // 背景部分を白く変える処理を書きたいです。
28                    
29                }
30            }
31        }
32        //newImage.addRepresentation(rep)
33        
34    }
35    
36    private func buffer(from image: UIImage) -> CVPixelBuffer? {
37        let attrs = [kCVPixelBufferCGImageCompatibilityKey: kCFBooleanTrue, kCVPixelBufferCGBitmapContextCompatibilityKey: kCFBooleanTrue] as CFDictionary
38        var pixelBuffer : CVPixelBuffer?
39        let status = CVPixelBufferCreate(kCFAllocatorDefault, Int(image.size.width), Int(image.size.height), kCVPixelFormatType_32ARGB, attrs, &pixelBuffer)
40        guard (status == kCVReturnSuccess) else {
41        return nil
42        }
43
44        CVPixelBufferLockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
45        let pixelData = CVPixelBufferGetBaseAddress(pixelBuffer!)
46
47        let rgbColorSpace = CGColorSpaceCreateDeviceRGB()
48        let context = CGContext(data: pixelData, width: Int(image.size.width), height: Int(image.size.height), bitsPerComponent: 8, bytesPerRow: CVPixelBufferGetBytesPerRow(pixelBuffer!), space: rgbColorSpace, bitmapInfo: CGImageAlphaInfo.noneSkipFirst.rawValue)
49
50        context?.translateBy(x: 0, y: image.size.height)
51        context?.scaleBy(x: 1.0, y: -1.0)
52
53        UIGraphicsPushContext(context!)
54        image.draw(in: CGRect(x: 0, y: 0, width: image.size.width, height: image.size.height))
55        UIGraphicsPopContext()
56        CVPixelBufferUnlockBaseAddress(pixelBuffer!, CVPixelBufferLockFlags(rawValue: 0))
57
58        return pixelBuffer
59    }
60//略