Как решить проблему с тем, что платформа Vision в Swift показывает слишком много фона ⇐ IOS

[Anonymous]

Я работаю над функцией для приложения Photo Booth на Swift, позволяющей выполнять замену фона с использованием платформы Vision и AVFoundation. Однако у меня возникла проблема: исходный фон отображается больше, чем должен, или, возможно, неправильно выровнен.
Вот изображение проблемы, которую я вижу:


Вот мой код для обработки окончательного предварительного просмотра:

Код: Выделить всё

func processFrameForFinalImage(
pixelBuffer: CVPixelBuffer,
orientation: CGImagePropertyOrientation,
quality: VNGeneratePersonSegmentationRequest.QualityLevel,
showInPreview: Bool,
completion: @escaping (UIImage?) ->  Void
) {

let request = VNGeneratePersonSegmentationRequest()
request.qualityLevel = quality
request.outputPixelFormat = kCVPixelFormatType_OneComponent8

let handler = VNImageRequestHandler(cvPixelBuffer: pixelBuffer,
orientation: .up,
options: [:])

DispatchQueue.global(qos: .userInitiated).async {
do {
try handler.perform([request])
guard let maskBuffer = request.results?.first?.pixelBuffer else {
DispatchQueue.main.async { completion(nil) }
return
}

// Convert buffers to CIImage and apply consistent orientation
let cameraImage = CIImage(cvPixelBuffer: pixelBuffer).oriented(orientation)
var maskImage = CIImage(cvPixelBuffer: maskBuffer).oriented(orientation)

maskImage = maskImage.applyingFilter("CIGaussianBlur", parameters: [
kCIInputRadiusKey: 2.0
])

if let erode = CIFilter(name: "CIMorphologyMinimum",
parameters: [kCIInputImageKey: maskImage,
"inputRadius": 1.0]),
let output = erode.outputImage {
maskImage = output
}

// Increase contrast to push mask values toward 0 or 1
if let controls = CIFilter(name: "CIColorControls") {
controls.setValue(maskImage, forKey: kCIInputImageKey)
controls.setValue(1.2, forKey: kCIInputContrastKey)
controls.setValue(0.0, forKey: kCIInputSaturationKey)
if let output = controls.outputImage {
maskImage = output
}
}

// Clamp to [0, 1]
if let clamp = CIFilter(name: "CIColorClamp") {
clamp.setValue(maskImage, forKey: kCIInputImageKey)
clamp.setValue(CIVector(x: 0, y: 0, z: 0, w: 0), forKey: "inputMinComponents")
clamp.setValue(CIVector(x: 1, y: 1, z: 1, w: 1), forKey: "inputMaxComponents")
if let output = clamp.outputImage {
maskImage = output
}
}

var composited = cameraImage

if self.isBackgroundReplacementEnabled,
let bgImage = self.backgroundImage {

// Scale background to match camera image size
var bg = bgImage.transformed(by: CGAffineTransform(
scaleX: cameraImage.extent.width / bgImage.extent.width,
y: cameraImage.extent.height / bgImage.extent.height
))

// Align background origin
let dx = cameraImage.extent.origin.x - bg.extent.origin.x
let dy = cameraImage.extent.origin.y - bg.extent.origin.y
bg = bg.transformed(by: CGAffineTransform(translationX: dx, y: dy))

// Blend the person (camera) over background using mask
composited = cameraImage.applyingFilter("CIBlendWithMask", parameters: [
kCIInputBackgroundImageKey: bg,
kCIInputMaskImageKey: maskImage
])
}

// Prepare for preview output
let finalImage = self.aspectFillImage(composited, targetSize: self.previewLayer.bounds.size)
guard let cgImage = self.ciContext.createCGImage(finalImage, from: finalImage.extent) else {
DispatchQueue.main.async { completion(nil) }
return
}

let uiImage = UIImage(cgImage: cgImage)

DispatchQueue.main.async {
if showInPreview {
self.previewLayer.contents = cgImage
}
completion(uiImage)
}

} catch {
print("Segmentation error: \(error)")
DispatchQueue.main.async { completion(nil) }
}
}
}

А это код для настройки камеры, я использую только переднюю камеру:

Код: Выделить всё

func setupCamera() {
session.beginConfiguration()
session.sessionPreset = .medium
guard let device = AVCaptureDevice.default(.builtInWideAngleCamera, for:  .video, position: .front),
let input = try? AVCaptureDeviceInput(device: device) else {
print("⚠️ Could not create camera input")
return
}

if session.canAddInput(input) {
session.addInput(input)
}

let output = AVCaptureVideoDataOutput()
output.setSampleBufferDelegate(self, queue: DispatchQueue(label: "camera.frame.queue"))
if session.canAddOutput(output) {
session.addOutput(output)
}

session.commitConfiguration()
}

Я не уверен, что делаю неправильно. Я думаю, это может быть связано с ориентацией.
Одна вещь, которую я заметил, это то, что я могу получить хорошие результаты только с помощью VNImageRequestHandler с ориентацией вверх. Первоначально я пытался передать то же значение, что и то, которое я называю ProcessFrameForFinalImage, то есть leftMirrored, но это привело к неправильному положению маски.

Подробнее здесь: https://stackoverflow.com/questions/797 ... background