AI Колоризація чорно-білих фотографій у мобільних додатках
Колоризація—задача з одним правильним входом та нескінченною множиною «правильних» виходів. Небо може бути голубим або сірим. Пальто—чорним або синім. Модель робить статистично вірогідний вибір, а не відновлює реальність. Важливо пояснити це через UX текст, але це не змінює технічну реалізацію.
Які моделі використовуємо
Класика—DeOldify (fastai + U-Net з self-attention). Добре узагальнюється, дає насичені кольори. Недолік: іноді «забруднює» нежиданими кольорами обличчя чи одяг при складному фоні.
DDColor (2023)—transformer архітектура, працює значно краще на портретах та архітектурі. Точніше передає колір шкіри та неба.
BigColor та ChromaGAN—для специфічних задач (історичні фотографії з сильною деградацією).
На мобіль йде одна з них—конвертована у Core ML або TFLite. Оригінальні ваги DeOldify—близько 250 МБ, після квантизації INT8—60–70 МБ. DDColor легше—110 МБ базова версія.
Конвертація DeOldify у Core ML
import coremltools as ct
import torch
from deoldify.visualize import get_image_colorizer
colorizer = get_image_colorizer(artistic=True) # або stable
model = colorizer.learn.model.eval()
# Експорт через torch.jit.trace
example = torch.zeros(1, 3, 256, 256)
traced = torch.jit.trace(model, example)
mlmodel = ct.convert(
traced,
inputs=[ct.TensorType(name="input", shape=(1, 3, 256, 256))],
compute_precision=ct.precision.FLOAT16,
minimum_deployment_target=ct.target.iOS16
)
mlmodel.save("DeOldify_artistic.mlpackage")
Важливо: DeOldify приймає RGB вхід (навіть для grayscale—він сам конвертує у Lab та працює з AB каналами). Перед передачею grayscale зображення його потрібно розширити у 3-канальний RGB повторенням: gray → [gray, gray, gray]. У CoreML це робиться в preprocessing, не в самій моделі.
iOS: запуск інференцу
let config = MLModelConfiguration()
config.computeUnits = .cpuAndNeuralEngine // ANE для FLOAT16
let model = try DeOldify_artistic(configuration: config)
// Підготовка: grayscale CVPixelBuffer → RGB
func prepareInput(from grayImage: UIImage) -> CVPixelBuffer? {
// Створюємо RGB піксельний буфер із grayscale, повторюючи канал
var pixelBuffer: CVPixelBuffer?
CVPixelBufferCreate(nil, width, height,
kCVPixelFormatType_32BGRA, nil, &pixelBuffer)
// ... копіюємо сірий канал у всі три
return pixelBuffer
}
let input = DeOldify_artisticInput(input: pixelBuffer)
let output = try model.prediction(input: input)
let colorizedImage = UIImage(cvPixelBuffer: output.output)
На iPhone 13 зображення 512×512 обробляється за 0.8–1.2 секунди. Для Full HD (1920×1080)—плиткове розбиття з патчами 512×512 та блендингом на стиках. Плиткований підхід створює проблему кольорової узгодженості між патчами: модель може пофарбувати один фрагмент неба в голубий, сусідній—у сірий. Рішення—глобальна гістограма кольорів: нормалізуємо гістограму кожної плитки до глобальної.
Android: TFLite з ONNX Runtime як альтернатива
// ONNX Runtime дає гнучкість: одна модель для iOS та Android
val env = OrtEnvironment.getEnvironment()
val session = env.createSession(
"deoldify_optimized.onnx",
OrtSession.SessionOptions().apply {
addNnapi() // або addCuda() якщо GPU доступна
}
)
val inputTensor = OnnxTensor.createTensor(env, inputArray, longArrayOf(1, 3, 512, 512))
val results = session.run(mapOf("input" to inputTensor))
val outputArray = (results[0].value as Array<*>)
ONNX Runtime Mobile—хороший вибір коли одна модель потрібна на обох платформах: не потрібно конвертувати двічі. NNAPI делегат працює на Android 8.1+.
UX: що важливо
Колоризація—повільна операція (1–5 секунд). Показуємо анімований прогрес. Хороший паттерн—«before/after» слайдер: користувач сдвигає розділювач, бачить оригінал та результат одночасно. Це і UX, і наглядна демонстрація роботи моделі.
Збереження результату—у .heic або .jpg максимальною якістю. Результат колоризації погано переносить повторне JPEG стиснення: артефакти на кольорових переходах стають помітні.
Процес
Вибір та конвертація моделі, оптимізація під цільові пристрої, реалізація плиткованого інференцу з кольоровою узгодженістю, UI з порівнянням before/after, тестування на історичних фотографіях різної якості.
Кошторис за часом
Одна платформа з базовою колоризацією займає 2–4 тижні. Обидві платформи з плиткуванням та нормалізацією кольорів вимагають 5–8 тижнів.