Всем привет! Однажды передо мной встала задача сделать фон картинки прозрачным, без мам, пап и бекндов... Задача есть, надо её решать! В этой статье вы узнаете, с чем мне пришлось столкнуться и как я реализовал вырезание заднего фона у фотографий в приложении на iOS.

Зачем вообще вот это вот всё ?

В приложении, над которым я работаю, есть редактор создания коллажа. Коллаж создаётся из фотографий одежды, аксессуаров, обуви и прочего. Все фотографии нашего каталога, имеют белый фон, но для реализации качественного результата этот фон нужно вырезать, а иначе получится мозаика из фотографий. Исторически, обрезание фона происходило на бекенде на устройстве силами библиотеки GPUImage, но результат оставлял желать лучшего...

Как можно заметить, присутствует очень много "снега" и артефактов.
Решено, внедряем машинное обучение !

Ищем способ удаления заднего фона

На просторах интернета есть хорошие туториалы, описывающие удаление фона с использованием модели машинного обучения DeeplabV3, представленной на официальном сайте Аpple. Только есть одно но. Данная моделька распознаёт ограниченный набор натренированных объектов и если скормить ей картинку с кроссовками или сумкой, она просто не распознает объекта на ней.

Продолжительный гуглинг открыл для меня прекрасную модель под названием u2Net. Данная модель очень качественно сегментирует фотографию на объекты, расположенные на ней. Но на страничке гитхаба модельки лежат в неведомом формате .pth, а ведь наш родной CoreML принимает только формат .mlModel, что делать ?
Есть два пути:

1. Используем официальную python тулзу для конвертации моделей в .mlModel формат

2. Заходим на страничку гитхаба, где сконвертировали всё за вас

Выбор очевиден ????

Внедряем U2Net в своё приложение

И так, для того, что бы магия свершилась, нам нужно, чтобы моделька оказалась внутри нашего приложения. Самый простой путь - перетащить её в директорию. На выбор вам предоставляется две модели, u2Net и u2Netp, пока выбирайте любую, в чём их отличие помимо размера мы поговорим далее.

После того как модель окажется внутри вашего файлового дерева, будет сгенерирован класс для импортированной модели, название его вы увидите на страничке модельки.

Наконец-то настало время попрограммировать. Разминаем пальчики !

Давайте, для удобства сделаем функцию, которая будет присутствовать у каждого экземпляра UIImage класса, что бы, если нам вздумается, мы могли вырезать фон где угодно !! И забегая немного вперёд, уточню, что модель машинного обучения принимают изображения строго определённого размера, который описан на страничке модели во вкладке Predictions.

import UIKit

extension UIImage {
  func removeBackgroudIfPosible() -> UIImage? {
      let resizedImage: UIImage = resize(size: .init(width: 320, height: 320))
      ...
  }
}

extension UIImage {
    func resize(size: CGSize? = nil, insets: UIEdgeInsets = .zero, fill: UIColor = .white) -> UIImage {
      var size: CGSize = size ?? self.size
      let widthRatio  = size.width / self.size.width
      let heightRatio = size.height / self.size.height
    
      if widthRatio > heightRatio {
        size = CGSize(width: floor(self.size.width * heightRatio), height: floor(self.size.height * heightRatio))
      } else if heightRatio > widthRatio {
        size = CGSize(width: floor(self.size.width * widthRatio), height: floor(self.size.height * widthRatio))
      }
    
      let rect = CGRect(x: 0,
                        y: 0,
                        width: size.width + insets.left + insets.right,
                        height: size.height + insets.top + insets.bottom)
    
      UIGraphicsBeginImageContextWithOptions(rect.size, false, scale)
    
      fill.setFill()
      UIGraphicsGetCurrentContext()?.fill(rect)
    
      draw(in: CGRect(x: insets.left,
                      y: insets.top,
                      width: size.width,
                      height: size.height))
      let newImage = UIGraphicsGetImageFromCurrentImageContext()
    
      UIGraphicsEndImageContext()
    
      return newImage!
    }
}

Так, картинку мы отресайзили под нужные нам размеры, что дальше, спросите вы ?
Дальше - магия машинного обучения и предсказание !!

import UIKit
extension UIImage {
    func removeBackgroudIfPosible() -> UIImage? {
        let resizedImage: UIImage = resize(size: .init(width: 320, height: 320))
        
        guard
            let resizedCGImage = resizedImage.cgImage,
            let originalCGImage = cgImage,
            let mlModel = try? u2netp(),
            let resultMask = try? mlModel.prediction(input: u2netpInput(in_0With: resizedCGImage)).out_p1 else {
            return nil
        }
        ...
    }
}

Давайте я опишу, получившийся, жирный guard и что там происходит.

• resizedCGImage - У u2netpInput, есть несколько инициализаторов, мы будем работать с тем, что принимает cgImage;

• originalCGImage - cgImage оригинального изображения, будем использовать далее для наложения маски;

• mlModel - инстанс модели машинного обучения, инициализированный на основе сгенерированного класса;

• resultMask - та самая магия! Функция prediction, возвращает нам u2netpOutput, где out_p1 можно использовать в качестве маски нашего изображения!

Если преобразовать получившуюся маску в картинку, то мы сможем увидеть то, как моделька обработала входящие картинки. Так как возвращаемая картинка имеет размер 320x320, не забудьте заресайзить её под исходные размеры.

let originalImage = CIImage(cgImage: originalCGImage)
var maskImage = CIImage(cvPixelBuffer: resultMask)
        
let scaleX = originalImage.extent.width / maskImage.extent.width
let scaleY = originalImage.extent.height / maskImage.extent.height
maskImage = maskImage.transformed(by: .init(scaleX: scaleX, y: scaleY))

return UIImage(ciImage: maskImage)

Вау! Выглядит очень аккуратно, нет снега и каки-то сильных вкраплений, давайте накладывать маску на оригинальную картинку !

Применяем получившуюся маску

У Apple есть хорошая статья, демонстрирующая подход по применения сегментирующей маски к исходному изображению. Её то мы и возьмём за основу с небольшими доработками.

let context = CIContext(options: nil)
        
guard let inputCGImage = context.createCGImage(originalImage, from: originalImage.extent) else {
    return nil
}
        
let blendFilter = CIFilter.blendWithRedMask()
        
blendFilter.inputImage = CIImage(cgImage: inputCGImage)
blendFilter.maskImage = maskImage
        
guard let outputCIImage = blendFilter.outputImage?.oriented(.up),
      let outputCGImage = context.createCGImage(outputCIImage, from: outputCIImage.extent) else {
        return nil
}
        
return UIImage(cgImage: outputCGImage)

Итак, поскольку мы оперируем CIImage и применяем к ней фильтр, я бы хотел сохранить исходное качество, с использованием CIContext удалось этого добиться.

Используя эти не хитрые манипуляции, мы можем взглянуть на результат!

Хочется сказать только одно - ВАУ! О такой качественном обрезании мы могли только мечтать, а количество кода, которое пришлось написать, просто мизирное!

import UIKit
import CoreML
import CoreImage.CIFilterBuiltins

extension UIImage {
    func removeBackgroudIfPosible() -> UIImage? {
        let resizedImage: UIImage = resize(size: .init(width: 320, height: 320))
        
        guard
            let resizedCGImage = resizedImage.cgImage,
            let originalCGImage = cgImage,
            let mlModel = try? u2netp(),
            let resultMask = try? mlModel.prediction(input: u2netpInput(in_0With: resizedCGImage)).out_p1 else {
            return nil
        }
        
        let originalImage = CIImage(cgImage: originalCGImage)
        var maskImage = CIImage(cvPixelBuffer: resultMask)
        
        let scaleX = originalImage.extent.width / maskImage.extent.width
        let scaleY = originalImage.extent.height / maskImage.extent.height
        maskImage = maskImage.transformed(by: .init(scaleX: scaleX, y: scaleY))
        
        let context = CIContext(options: nil)
        
        guard let inputCGImage = context.createCGImage(originalImage, from: originalImage.extent) else {
            return nil
        }
        
        let blendFilter = CIFilter.blendWithRedMask()
        
        blendFilter.inputImage = CIImage(cgImage: inputCGImage)
        blendFilter.maskImage = maskImage
        
        guard let outputCIImage = blendFilter.outputImage?.oriented(.up),
              let outputCGImage = context.createCGImage(outputCIImage, from: outputCIImage.extent) else {
            return nil
        }
        
        return UIImage(cgImage: outputCGImage)
    }
}

Всё круто, но как ты собираешься хранить такую тяжёлую модель ?

Модель машинного обучения u2Net, распространяется в двух версиях, u2Netp - light версия 4.6 мб и u2Net - full версия 175.9 мб . Разница между ними очень ощутима, если обратить внимание на их вес. Полную версию u2Net даже не получится хранить в гите, так что не советую вам комитить код с полной моделькой, если планируете его пушить на github. На результат работы light модели можно взглянуть ниже.

Мы же приняли решение хранить в приложении light версию модели, а большую, полную версию модели загружать в фоне, что позволило сохранить функциональность редактора, пока модель не загружена. Если вы захотите сделать тоже самое, Apple и тут о нас подумали и написали туториал !

Итог

Теперь мы умеем пользоваться моделями машинного обучения в iOS приложении и применять их на практике. Машинное обучение открывает море возможностей для распознания, сегментации, классификации, оценки глубины, обнаружения объектов и позы человека и это всё только про фото, а ведь можно ещё и работать с текстом !

Помимо сторонних моделей машинного обучения, iOS имеет встроенный фраемворк Vision, который, оперируя алгоритмами машинного зрения, поможет вам обнаружить лицо, кисть, штрихкод или человека!

Спасибо за прочтение! Надеюсь, что эта статья помогла вам решить аналогичную задачу или просто помогла узнать, как работать с ML моделями в iOS.

Используемые материалы

Список ML моделей от Apple

Тулза по конвертации моделей в mlModel

Конвертированные популярные ML модели

Статья Apple по применению сегментированной маски

Статья Apple по загрузки ML модели из сети

Статья по удалению фона с использованием DeeplabV3