Реализация AI-колоризации чёрно-белых фотографий в мобильном приложении

Реализация AI-колоризации чёрно-белых фотографий в мобильном приложении

Колоризация — задача с одним правильным входом и бесконечным множеством «правильных» выходов. Небо может быть голубым или серым. Пальто — чёрным или синим. Модель делает статистически вероятный выбор, а не восстанавливает реальность. Это важно объяснять пользователю UX-текстом, но не меняет техническую реализацию.

Какие модели используем

Классика — DeOldify (fastai + U-Net с self-attention). Хорошо обобщается, даёт насыщенные цвета. Минус: иногда «загрязняет» цвет в нежелательные пятна на лицах или одежде при сложном фоне.

DDColor (2023) — transformer-based архитектура, работает значительно лучше на портретах и архитектуре. Точнее передаёт цвет кожи и неба.

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 и блендингом на стыках. Тайловый подход создаёт проблему цветовой несогласованности между патчами: модель может покрасить один фрагмент неба в голубой, а соседний — в серый. Решение — global color histogram matching: нормализуем гистограмму каждого тайла к глобальной гистограмме.

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 недель.