Всем привет! Однажды передо мной встала задача сделать фон картинки прозрачным, без мам, пап и бекндов... Задача есть, надо её решать! В этой статье вы узнаете, с чем мне пришлось столкнуться и как я реализовал вырезание заднего фона у фотографий в приложении на iOS.
Зачем вообще вот это вот всё ?
В приложении, над которым я работаю, есть редактор создания коллажа. Коллаж создаётся из фотографий одежды, аксессуаров, обуви и прочего. Все фотографии нашего каталога, имеют белый фон, но для реализации качественного результата этот фон нужно вырезать, а иначе получится мозаика из фотографий. Исторически, обрезание фона происходило на бекенде на устройстве силами библиотеки GPUImage, но результат оставлял желать лучшего...
Как можно заметить, присутствует очень много "снега" и артефактов.
Решено, внедряем машинное обучение !
Ищем способ удаления заднего фона
На просторах интернета есть хорошие туториалы, описывающие удаление фона с использованием модели машинного обучения DeeplabV3, представленной на официальном сайте Аpple. Только есть одно но. Данная моделька распознаёт ограниченный набор натренированных объектов и если скормить ей картинку с кроссовками или сумкой, она просто не распознает объекта на ней.
Продолжительный гуглинг открыл для меня прекрасную модель под названием u2Net. Данная модель очень качественно сегментирует фотографию на объекты, расположенные на ней. Но на страничке гитхаба модельки лежат в неведомом формате .pth, а ведь наш родной CoreML принимает только формат .mlModel, что делать ?
Есть два пути:
Используем официальную python тулзу для конвертации моделей в .mlModel формат
Заходим на страничку гитхаба, где сконвертировали всё за вас
Выбор очевиден ?
Внедряем U2Net в своё приложение
И так, для того, что бы магия свершилась, нам нужно, чтобы моделька оказалась внутри нашего приложения. Самый простой путь - перетащить её в директорию. На выбор вам предоставляется две модели, u2Net и u2Netp, пока выбирайте любую, в чём их отличие помимо размера мы поговорим далее.
После того как модель окажется внутри вашего файлового дерева, будет сгенерирован класс для импортированной модели, название его вы увидите на страничке модельки.
Наконец-то настало время попрограммировать. Разминаем пальчики !
Давайте, для удобства сделаем функцию, которая будет присутствовать у каждого экземпляра UIImage класса, что бы, если нам вздумается, мы могли вырезать фон где угодно !! И забегая немного вперёд, уточню, что модель машинного обучения принимают изображения строго определённого размера, который описан на страничке модели во вкладке Predictions.
import UIKit
extension UIImage {
func removeBackgroudIfPosible() -> UIImage? {
let resizedImage: UIImage = resize(size: .init(width: 320, height: 320))
...
}
}Используемая функция resize
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.
Используемые материалы
Тулза по конвертации моделей в mlModel
Конвертированные популярные ML модели
Статья Apple по применению сегментированной маски