Генерация разнообразного контента с помощью ИИ продолжает быть на пике популярности. На смену картинкам по описанию пришли музыкальные композиции на основе текста и психоделические видео, на которых у людей меняется не только геометрия, но и вообще всё. Однако это лишь вершина айсберга. We need to go deeper. Хабру нужны не смешные нейро(де)генеративные мемы, а статьи от людей, которые работают с генеративным ИИ профессионально и на острие современных технологий пытаются сделать нечто крутое и полезное.

Привет, меня зовут Алексей Луговой, я занимаюсь Computer Vision в Самолете, и сегодня объявляю о старте автоген-челленджа. Этот челлендж — совместная инициатива Хабра и Самолета. Про призы лучшим авторам и другие детали расскажу подробнее в конце статьи, а начну с личного примера — расскажу, как мы научились подставлять другую мебель на фото интерьера.

Я уже рассказывал про наше решение для дизайна интерьеров на лету. Сегодня речь пойдёт о более узкой задаче: есть фото интерьера, и клиент спрашивает — а как бы смотрелось, если бы вместо вот этого дивана стоял другой? Возможность ответить на такой вопрос не расплывчатой фразой, а фотореалистичным изображением — это и удовлетворение клиента, и вау-эффект, и просто красиво. Сегодня я расскажу, какими средствами мы этого добились.

Этап 1. Как понять, что заменить

Какой бы ни была наша модель, ясно, что на вход она в той или иной форме будет получать три вещи:

• фото интерьера,

• указание, что нужно заменить,

• указание, на что нужно заменить.

С первым пунктом всё понятно (на самом деле нет, но об этом позже), так что начнём со второго. Как именно модель поймёт, что замене подлежит именно вот этот жёлтый диван?

Самый простой (с точки зрения модели) вариант — если ей на вход дадут маску со всеми пикселями, подлежащими замене. Это годится для лабораторной работы, но не для продукта. Пока сотрудник Самолета будет пипеткой выбирать пиксели, клиент заскучает и уйдёт. Значит, этот момент нужно автоматизировать. С точки зрения UX хорошо, если можно будет просто выбрать «Диван» из выпадающего списка.

Здесь мы встали на распутье. Можно обучить свою модель, чётко под свой кейс. Она будет в совершенстве понимать специфику мебели, причём именно нашей мебели, и, наверное, это поможет ей круто справляться с задачей. Но своя модель — это долго. Сбор и разметка данных, дообучение — это всё очень увеличивает тайм-ту-маркет.

Другой путь — взять какое-то решение из коробки и допилить под наши нужды. Возможно, такая нейросетка будет менее тонко чувствовать нюансы диванов. А возможно, и наоборот. Всё-таки SOTA-решения делались умными людьми, которые вложили в них много сил и машинного времени. Имеет смысл хотя бы попробовать, собрать MVP из чего бог послал, а дальше смотреть на метрики и решать, устраивает ли нас качество.

Поразмыслив, мы усмирили гордость и пошли вторым путём.

Этап 2. Детекция

Задача поиска пикселей, которые необходимо заменить, на практике декомпозируется в две — детекция и сегментация. Детекция — это когда объект локализуется на изображении и для него находится ограничивающий прямоугольник (bbox). Большинство пикселей ббокса не относится к объекту, поэтому нужен следующий этап — сегментация, на котором получается пиксельная маска.

Для детекции мы взяли Grounding DINO — SOTA-модель, подходящую для всего. Она понимает промпты на естественном языке и ищет по ним любые объекты в пределах своего понимания. Модель активно развивается, не так давно ребята из IDEA выкатили версию 1.5. Для начала мы взяли готовую имплементацию с HuggingFace, без какого-либо дообучения. Спойлер: этого оказалось достаточно.

Вот так Grounding DINO находит диван:

text_prompt = ‘couch’
box_threshold=0.25
text_threshold=0.25

import requests

import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection


def draw_boxes_on_image(img_pil, bbox_info):
    img_pil = img_pil.copy()
    draw = ImageDraw.Draw(img_pil)
    
    boxes = bbox_info[0]['boxes']
    labels = bbox_info[0]['labels']
    scores = bbox_info[0]['scores']

    font = ImageFont.load_default()
    for box, label, score in zip(boxes, labels, scores):
        x1, y1, x2, y2 = box
        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
        text = f'{label}: {score:.2f}'
        text_width, text_height = draw.textsize(text, font=font)
        draw.text((x1, y1 - text_height - 2), text, fill="red", font=font)

    return img_pil

model_id = "IDEA-Research/grounding-dino-base"
processor = AutoProcessor.from_pretrained(model_id, cache_dir=CACHE_DIR)
model = AutoModelForZeroShotObjectDetection.from_pretrained(model_id, cache_dir=CACHE_DIR).to(DEVICE)

# inference
text = 'couch'
inputs = processor(images=img_pil, text=text, return_tensors="pt").to(DEVICE)
with torch.no_grad():
    outputs = model(**inputs)

bbox_info = processor.post_process_grounded_object_detection(
    outputs,
    inputs.input_ids,
    box_threshold=0.25,
    text_threshold=0.25,
    target_sizes=[img_pil.size[::-1]]
)

# postprocess 
bbox_info[0]['boxes'] = bbox_info[0]['boxes'].cpu().numpy()
bbox_info[0]['scores'] = bbox_info[0]['scores'].cpu().numpy()

pprint(bbox_info)
img_pil_an = draw_boxes_on_image(img_pil, bbox_info)

А вот так, например, тумбу:

text_prompt = ‘nightstand’’
box_threshold=0.25
text_threshold=0.25

Есть нюанс — чтобы хорошо генерировались ббоксы, нужно правильно задавать трешхолды (пороговые значения уверенности модели). Слишком высокий трешхолд — модель засомневается в себе и вообще не выдаст ббокс. Слишком низкий — и она выдаст их десяток на всё хотя бы отдалённо похожее. Обычно всё хорошо работает из коробки, но если клиент приносит «проблемное» фото (низкое качество, неудачный ракурс, много мелких объектов), то высока вероятность, что возникнет одна из двух проблем, описанных выше.

Можно разруливать такие кейсы с помощью «естественных нейросетей»: сказать сотруднику, чтобы в таких случаях он просил у клиента другое фото. Но снять другое фото — это достаточно «дорогая» операция. И для клиента (время, усилия), и для нас (вау-эффект испорчен, клиент недоволен). 

Есть несколько вариантов, как решить проблему на стороне софта:

• Автоподбор трешхолдов в случае неудачной сегментации

• Поднять качество фото с помощью той или иной технологии апскейла (возможно, ещё одной генеративной моделью)

• Или всё-таки вернуться к варианту 1 и обучить/дообучить свою, более специализированную модель

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

Этап 3. Сегментация

Итак, у нас есть bbox, и нам нужна пиксельная маска. На этом шаге мы использовали Segment Anything Model (SAM) — ещё одну опенсорсную SOTA-модель. Точнее, на тот момент она была SOTA, сейчас вышла вторая версия, мы задумываемся о переходе, но пока что не переходили. SAM умеет принимать на вход разные вещи, в том числе bbox, а на выходе даёт пиксельную маску. 

Скрестить Grounding DINO и SAM — это не наше ноу-хау. Собственно, даже в репозитории Groundind DINO есть ссылка на связку Grounded-SAM. Мы не стали брать готовую связку, чтобы иметь больше контроля над компонентами, но концептуально делаем то же самое.

import torch
from PIL import Image
import requests
from transformers import SamModel, SamProcessor

model_id = "facebook/sam-vit-huge"
model = SamModel.from_pretrained(model_id, cache_dir=CACHE_DIR).to(DEVICE)
processor = SamProcessor.from_pretrained(model_id, cache_dir=CACHE_DIR)

boxes = bbox_info[0]['boxes'].tolist()
inputs = processor(img_pil, input_boxes=[boxes], return_tensors="pt").to(DEVICE)
with torch.no_grad():
    outputs = model(**inputs)

masks = processor.image_processor.post_process_masks(
    outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu(), mask_threshold=-10
)
scores = outputs.iou_scores

numpy_data = masks[0].numpy().astype('uint8') * 255
image_data = numpy_data[0, 0]
image = Image.fromarray(image_data)

В итоге получается что-то такое:

SAM чертовски хорош, но есть нюансы. В каких случаях всё пойдёт не так?

• Если у объекта сложные, нечёткие или размазанные границы, у модели могут возникнуть трудности с их точным выделением.

• SAM может запутаться в плотной сцене со множеством объектов (например, стол, на котором ложки-вилки-тарелки).

• Маленькие объекты — их выделять труднее.

• Плохое качество изображения..

• Не вполне удачно сгенерировался bbox на предыдущем этапе.

Решения этих проблем в основном очевидны. Плохое изображение улучшить, маленькие объекты увеличить, и т.п.

Самая интересная проблема — это сложная сцена. Возьмем стол с посудой, или, допустим, диван с подушками. SAM может решить, что подушки — это отдельные объекты, не связанные с диваном, и на их месте останутся дырки в маске. В каком-то смысле SAM даже прав — однако нам от этого не легче, нам нужен другой результат.

Вариантов решения может быть несколько:

• Задетектить всю сцену полностью — тогда будут ббоксы и для дивана, и для подушек, и для пледа, вообще для всего. Потом отдать SAM кучу ббоксов, которые географически связаны с основным, и пусть он сегментирует каждый. Однако детекция всего — грубое и ресурсоёмкое решение.

• Поиграть с настройками трешхолда для определения маски (параметр mask_threshold в нашем примере кода). Если придать SAM больше «уверенности», он может согласиться, что подушки — часть дивана.

• Сегментировать всё, Наташа, вообще всё — использовать опцию SAM «segment everything». Когда ему не нужно задавать никаких промптов, он просто разбивает на сегменты изображение целиком. В мире, где вычисления происходят бесплатно и мгновенно, это могло бы быть лучшей опцией, на практике это неоптимально.

• Можно оптимизировать предыдущий вариант, сегментируя только то, что рядом с ббоксом. Заодно так можно компенсировать недостаточно точный ббокс — если какая-то часть объекта не влезла на этапе детекции.

• Постобработка маски (об этом ниже).

Здесь мы опять же в фазе ресёрча, обдумываем эти варианты, а также с десяток других. Чтобы вы понимали, перечисленные проблемы — это краевые случаи, которые встречаются не так уж часто. Наш MVP уже достаточно хорош, чтобы работать в большей части случаев и иметь бизнес-ценность. Но пространство для его улучшения безгранично.

Этап 4. Постобработка маски

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

def refine_img_morphology(img_pil, kernel_size=15, erosion_iter=1, dilation_iter=3):
    """ 
    - убрать черные засечки вне маски
    - белые засечки внутри маски
    - увеличить маскируемую область (по форме)
    
    """

    img_np = np.array(img_pil)
    
    # Создание ядра для морфологических операций
    kernel = np.ones((kernel_size, kernel_size), np.uint8)

    # Удаление белых засечек вне объекта
    erosion = cv2.erode(img_np, kernel, iterations=erosion_iter)

    # Закрашивание черных областей внутри объекта
    dilation = cv2.dilate(erosion, kernel, iterations=dilation_iter)
    closing = cv2.erode(dilation, kernel, iterations=erosion_iter)

    # Конвертация обработанного изображения обратно в формат PIL
    if closing.ndim == 3:
        closing = cv2.cvtColor(closing, cv2.COLOR_BGR2RGB)
    processed_pil_image = Image.fromarray(closing)

    return processed_pil_image

mask_l = refine_img_morphology(mask_l)

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

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

Кстати, на этом этапе мы можем не только заменить предмет на фото, но и удалить его совсем с помощью LaMa. Так можно легко убрать с картинки мусор или мелкие неугодные вещи. Чем больше предмет, тем сложнее его «затереть» — нейронка не может представить, что находится за этим объектом.

Например, тут мы убрали телевизор:

Этап 5. Генерация конечного результата

Итак, у нас есть оригинальное фото и маска объекта на нём. Теперь мы готовы к генерации. Вроде бы. Берём модель для инпейнтинга (например, эту), даём ей фото, маску и промпт — и вуаля, мы великолепны.

from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image
import torch

pipe = AutoPipelineForInpainting.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype=torch.float16, variant="fp16").to("cuda")

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

image = load_image(img_url).resize((1024, 1024))
mask_image = load_image(mask_url).resize((1024, 1024))

prompt = "a tiger sitting on a park bench"
generator = torch.Generator(device="cuda").manual_seed(0)

image = pipe(
  prompt=prompt,
  image=image,
  mask_image=mask_image,
  guidance_scale=8.0,
  num_inference_steps=20,  # steps between 15 and 30 work well for us
  strength=0.99,  # make sure to use `strength` below 1.0
  generator=generator,
).images[0]

Было:

Стало:

Однако дьявол кроется в деталях, и если приглядеться, начинаешь эти детали замечать. Дело в том, что при инпейнтинге по маске изображение перегенерируется целиком (как минимум, этап кодера/декодера), и при этом неизбежно происходят хоть и небольшие, но изменения по сравнению с исходником. Появляются шумы, замыливание (например, обратите внимание на вотермарку). Вроде бы и не критично, но всё равно неприятно.

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

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

Здесь главный нюанс — выбор фрагмента. Он должен иметь размер, оптимальный для генеративной модели, содержать в себе нужный объект и не выходить за границы оригинала. Тут уже задача не машинлёрнинговая, а вполне себе классическая алгоритмическая.

Что касается оптимального размера — мы в итоге выбрали stable-diffusion-2-inpainting, которая обучалась на изображениях 512х512. Модель умеет и в другие размеры, но согласно совету разработчиков для наилучшего качества лучше придерживаться этих.

def expand_crop_side(current_coords, original_size, target_length):
    start, end = current_coords
    current_length = end - start

    # Вычисление, сколько не хватает до желаемого размера
    delta = target_length - current_length
    if delta <= 0:
        return current_coords  # Текущий размер уже соответствует или превышает целевой

    # Рассчитываем расширение с каждой стороны
    expand_one_side = min(delta // 2, start)
    expand_other_side = min(delta - expand_one_side, original_size - end)

    # Расширяем обрезку
    new_start = start - expand_one_side
    new_end = end + expand_other_side

    # Компенсация, если одна сторона достигла границы
    remaining = delta - (expand_one_side + expand_other_side)
    new_start = max(new_start - remaining, 0) if new_end == original_size else new_start
    new_end = min(new_end + remaining, original_size) if new_start == 0 else new_end

    return (new_start, new_end)


def to_multiplicity_size(start_coord, end_coord, original_size, target_multiple=8):
    """
    делает кратность к target_multiple
    если доступно - расширяет в большую сторону
    если нет - в меньшую
    
    """
    
    size = end_coord - start_coord
    new_size = ((size + target_multiple - 1) // target_multiple) * target_multiple
    diff = new_size - size

    # Попытка увеличить размер
    if end_coord + diff <= original_size:
        end_coord += diff
    elif start_coord - diff >= 0:
        start_coord -= diff
    else:
        # Если увеличить не получается, уменьшаем
        new_size = (size // target_multiple) * target_multiple
        diff = size - new_size
        end_coord = start_coord + new_size

    return start_coord, end_coord
    
    
target_width, target_height = 512, 512
original_width, original_height = maskl.size

maskl_np = np.array(maskl)
img_np = np.array(img_pil)

# описываем границы маски прямоугольником
rows, cols = np.where(maskl_np == 255)
top, bottom = np.min(rows), np.max(rows)
left, right = np.min(cols), np.max(cols)

# пытаемся обрезать изображение до желаемого размера
new_left, new_right = expand_crop_side((left, right), original_width, target_width)
new_top, new_bottom = expand_crop_side((top, bottom), original_height, target_height)

# корректируем обрезку до размеров, кратным 8 (условие модели stable diffusion)
new_left, new_right = to_multiplicity_size(new_left, new_right, original_width, 8)
new_top, new_bottom = to_multiplicity_size(new_top, new_bottom, original_height, 8)

crop_coords = (new_left, new_top, new_right, new_bottom)
img_pil_cropped = img_pil.crop(crop_coords).copy()
maskl_cropped = maskl.crop(crop_coords).copy()

width, height = img_orig_pil_cropped.size

display_images([img_pil_cropped, maskl_cropped])

from diffusers import StableDiffusionInpaintPipeline, AutoPipelineForInpainting

pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-inpainting",
    torch_dtype=torch.float16,
    safety_checker=None
).to(DEVICE)

# inference
prompt = "red couch"
generator = torch.Generator(device=DEVICE).manual_seed(4)

img_generated_pil = pipe(
    prompt=prompt,
    image=img_pil_cropped,
    mask_image=maskl_cropped,
    guidance_scale=8.0,
    num_inference_steps=120, 
    strength=0.9,
    height=height,
    width=width,
    generator=generator,
)

img_generated_pil = img_generated_pil.images[0]
display_images([img_generated_pil])

repainted_image = img_generated_pil.copy()
init_image = img_pil_cropped.copy()
mask_image_arr = np.array(maskl_cropped.convert("L")).copy()

# Add a channel dimension to the end of the grayscale mask
mask_image_arr = mask_image_arr[:, :, None]

# Binarize the mask: 1s correspond to the pixels which are repainted
mask_image_arr = mask_image_arr.astype(np.float32) / 255.0
mask_image_arr[mask_image_arr < 0.5] = 0
mask_image_arr[mask_image_arr >= 0.5] = 1

# Take the masked pixels from the repainted image and the unmasked pixels from the initial image
unmasked_unchanged_image_arr = (1 - mask_image_arr) * init_image + mask_image_arr * repainted_image
unmasked_unchanged_image = Image.fromarray(unmasked_unchanged_image_arr.round().astype("uint8"))

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

Этап 6. Дообучение на своих данных

Итак, мы научились вместо какого-то жёлтого дивана вставлять какой-то красный диван. Это круто? Как шоукейс генеративных технологий — безусловно да. Как продукт — не совсем. Нам нужно показать не просто красный диван, а конкретный красный диван из нашего каталога — и для вау-эффекта, и чтобы можно было этот самый диван немедленно продать.

К сожалению, наша мебель недостаточно знаменита, чтобы Stable Diffusion всё знала про неё «из коробки». Можно попытаться хорошо описать его промптами, но это ненадёжно. А значит, тут нас коробочное решение уже не устраивает, и придётся не играть в конструктор, а реально заморочиться.

Вообще diffusion-based модели огромные и неповоротливые. Чтобы дообучать такую «в лоб», надо быть очень богатым и очень терпеливым человеком, — времени и ресурсов уйдёт прорва. Однако мы не первые, у кого возникло желание дообучить большую генеративную модель. Комьюнити уже придумало несколько легковесных методов. Например, можно использовать адаптеры типа LoRa — добавить к модели сверху ещё слой и обучать только его. Альтернативно можно «хакнуть» обучение, и за малое количество проходов на маленькой выборке добиться оверфита по этой выборке (обычно оверфит — это плохо, но в нашем случае это как раз то, что нужно). Это метод dreambooth, достаточно ресурсоёмкий по сравнению с LoRa, но при этом он обещает лучшее качество. Мы решили не экономить на вау-эффекте и использовать dreambooth.

Генеративки на хайпе, комьюнити большое, опенсорс-модели творят чудеса, поэтому заходим на гитхаб диффузерс и выбираем любой скрипт обучения под нашу задачу и версию модели. Некоторые скрипты могут быть нерабочие (добро пожаловать в Issues проекта), иногда даже можно похардкодить и скорректировать скрипты вручную под свои хотелки. Разумеется, нужно понимать что делаешь, что такое prior loss и зачем нужны class images — иначе чудо не произойдёт.

В итоге мы взяли вот такое решение, оно показалось наиболее удобным и масштабируемым (было меньше ошибок при попытке его завести).

Это картинки, на которых дообучали модель:

А вот что имеем на выходе:

Собирать ИИ-многоножку из нейросетей — это весело и увлекательно, но не стоит забывать про такую скучную вещь, как подготовка данных. На этом этапе можно выстрелить себе в ногу, например, не добавив в выборку нужный ракурс. У Stable Diffusion может не хватить воображения, чтобы представить, как диван будет выглядеть сзади.

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

Демку проекта целиком можете посмотреть у меня на гитхабе.

Эта статья — пример того, что мы хотели бы видеть на автоген-челлендже. Нам нужны не обзоры и не научпоп, а технически насыщенные материалы от тех, кто имеет непосредственное отношение к разработке в области генеративного ИИ. Лишь благодаря технохардкору Хабр — торт.

А теперь подробнее про условия.

• Челлендж будет идти 6 недель с момента публикации этой статьи

• К участию допускаются статьи и посты в личных и корпоративных блогах, у которых выставлен специальный тег — автоген-челлендж.

• Обратите внимание: в челлендже могут участвовать не только статьи, но и более короткие посты. Если вам есть что сказать коротко и по делу, этого может быть достаточно

• В центре внимания — генеративный ИИ, большие языковые модели, ИИ-агенты. Применение в реальных задачах, настройка инфраструктуры, обучение и тестирование.

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

«Но зачем нам соблюдать какие-то условия?» — спросите вы. А затем, чтобы получить призы. Из пятерки статей и постов-участников с самым высоким рейтингом эксперты Самолета выберут от одного до трёх победителей. Этим счастливчикам достанутся информационные гранты от Хабра, которые они смогут использовать для продвижения своего проекта.

• 1 место: блог по тарифу «Бизнес» на полгода, истории на Хабре и пост в социальных сетях о вашем кейсе.

• 2 место: блог по тарифу «Бизнес» на полгода, пост в социальных сетях о вашем кейсе.

• 3 место: блог по тарифу «Бизнес» на полгода.

Кроме того, все победители получат мерч от Самолета — вот такой:

На этом всё. Авторы, жгите автогеном, буду ждать ваших крутых статей.