Marigold Computer Vision

Marigold был предложен в статье Repurposing Diffusion-Based Image Generators for Monocular Depth Estimation, представленной на CVPR 2024 в формате Oral, авторами которой являются Bingxin Ke, Anton Obukhov, Shengyu Huang, Nando Metzger, Rodrigo Caye Daudt и Konrad Schindler. Основная идея заключается в том, чтобы использовать генеративный априорный модель Text-to-Image Latent Diffusion Models (LDMs) для традиционных задач компьютерного зрения. Этот подход был исследован путем дообучения Stable Diffusion для задачи Monocular Depth Estimation, что продемонстрировано на иллюстрации выше.

Marigold был впоследствии расширен в продолжении, статье Marigold: Affordable Adaptation of Diffusion-Based Image Generators for Image Analysis, авторами которой являются Bingxin Ke, Kevin Qu, Tianfu Wang, Nando Metzger, Shengyu Huang, Bo Li, Anton Obukhov и Konrad Schindler. В этой работе Marigold был расширен для поддержки новых модальностей, таких как Surface Normals и Intrinsic Image Decomposition (IID), был предложен протокол обучения для Latent Consistency Models (LCM) и продемонстрирована возможность обработки изображений в High-Resolution (HR) режиме.

Ранние модели Marigold (v1-0 и более ранние) были оптимизированы для достижения наилучших результатов при использовании не менее 10 шагов инференса. Модели LCM были разработаны позже, чтобы обеспечить высококачественный инференс всего за 1–4 шага. Модели Marigold v1-1 и более поздние используют планировщик DDIM для достижения оптимальных результатов всего за 1–4 шага.

Доступные конвейеры

Каждый конвейер предназначен для конкретной задачи компьютерного зрения: он обрабатывает входное RGB-изображение и формирует соответствующее предсказание. В настоящее время реализованы следующие задачи компьютерного зрения:

Доступные контрольные точки

Все исходные контрольные точки доступны в организации PRS-ETH на Hugging Face. Они предназначены для использования с конвейерами diffusers и исходным кодом, который также можно использовать для обучения новых контрольных точек моделей. Ниже приведено краткое описание рекомендуемых контрольных точек, все из которых дают надежные результаты за 1–4 шага.

Обязательно ознакомьтесь с руководством по планировщикам, чтобы узнать, как исследовать компромисс между скоростью и качеством планировщика, а также с разделом повторное использование компонентов в разных конвейерах, чтобы узнать, как эффективно загружать одни и те же компоненты в несколько конвейеров. Кроме того, чтобы узнать больше о снижении потребления памяти этим конвейером, обратитесь к разделу [«Снижение потребления памяти»] здесь.

Пайплайны Marigold были разработаны и протестированы с планировщиком, встроенным в контрольную точку модели. Оптимальное количество шагов инференса зависит от планировщика, и универсального значения, которое было бы лучшим во всех случаях, не существует. Чтобы учесть это, параметр num_inference_steps в методе __call__ пайплайна по умолчанию принимает значение None (см. справочник API). Если значение не задано явно, оно наследуется из поля default_denoising_steps в файле конфигурации контрольной точки (model_index.json). Это обеспечивает высокое качество предсказаний при вызове пайплайна только с аргументом image.

Приведённые ниже примеры в основном относятся к предсказанию глубины, однако их можно универсально применять к другим поддерживаемым модальностям. Мы демонстрируем предсказания, используя одно и то же входное изображение Альберта Эйнштейна, сгенерированное Midjourney. Это упрощает сравнение визуализаций предсказаний для различных модальностей и контрольных точек.

Предсказание глубины

Чтобы получить предсказание глубины, загрузите чекпоинт prs-eth/marigold-depth-v1-1 в MarigoldDepthPipeline, пропустите изображение через пайплайн и сохраните предсказания:

import diffusers
import torch

pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
    "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

depth = pipe(image)

vis = pipe.image_processor.visualize_depth(depth.prediction)
vis[0].save("einstein_depth.png")

depth_16bit = pipe.image_processor.export_depth_to_16bit_png(depth.prediction)
depth_16bit[0].save("einstein_depth_16bit.png")

Функция visualize_depth() применяет одну из цветовых карт matplotlib (по умолчанию Spectral), чтобы отображать предсказанные значения пикселей из одноканального диапазона глубины [0, 1] в RGB-изображение. При использовании цветовой карты Spectral пиксели с малой глубиной окрашиваются в красный цвет, а пиксели с большой глубиной — в синий. Файл PNG с глубиной цвета 16 бит хранит значения одного канала, линейно отображённые из диапазона [0, 1] в диапазон [0, 65535]. Ниже приведены исходные и визуализированные предсказания. Тёмные и близкие области (усы) легче различить на визуализации.

Оценка нормалей поверхности

Загрузите чекпоинт prs-eth/marigold-normals-v1-1 в MarigoldNormalsPipeline, пропустите изображение через пайплайн и сохраните предсказания:

import diffusers
import torch

pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(
    "prs-eth/marigold-normals-v1-1", variant="fp16", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

normals = pipe(image)

vis = pipe.image_processor.visualize_normals(normals.prediction)
vis[0].save("einstein_normals.png")

Функция visualize_normals() отображает трёхмерное предсказание со значениями пикселей в диапазоне [-1, 1] в RGB-изображение. Функция визуализации поддерживает переворот осей нормалей поверхности, чтобы сделать визуализацию совместимой с другими выборами системы координат. Концептуально каждый пиксель окрашивается в соответствии с вектором нормали поверхности в системе координат, где ось X направлена вправо, ось Y — вверх, а ось Z — на зрителя. Ниже приведено визуализированное предсказание:

В этом примере на кончике носа почти наверняка есть точка на поверхности, нормаль к которой направлена прямо на наблюдателя, то есть её координаты равны [0, 0, 1]. Этот вектор отображается в RGB-значение [128, 128, 255], что соответствует фиолетово-синему цвету. Аналогично, нормаль к поверхности на щеке в правой части изображения имеет большую компоненту X, что усиливает красный оттенок. Точки на плечах, направленные вверх с большой компонентой Y, способствуют появлению зеленого цвета.

Разложение собственных изображений

Marigold предоставляет две модели для разложения изображений на внутренние компоненты (Intrinsic Image Decomposition, IID): «Appearance» и «Lighting». Каждая модель генерирует карты альбедо, полученные из аннотаций InteriorVerse и Hypersim соответственно.

• Модель «Appearance» также оценивает свойства материала: Roughness и Metallicity.
• Модель «Lighting» генерирует Diffuse Shading и Non-diffuse Residual.

Ниже приведён пример кода, сохраняющего предсказания, сделанные моделью «Appearance»:

import diffusers
import torch

pipe = diffusers.MarigoldIntrinsicsPipeline.from_pretrained(
    "prs-eth/marigold-iid-appearance-v1-1", variant="fp16", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

intrinsics = pipe(image)

vis = pipe.image_processor.visualize_intrinsics(intrinsics.prediction, pipe.target_properties)
vis[0]["albedo"].save("einstein_albedo.png")
vis[0]["roughness"].save("einstein_roughness.png")
vis[0]["metallicity"].save("einstein_metallicity.png")

Ещё один пример, демонстрирующий предсказания, сделанные моделью «Lighting»:

import diffusers
import torch

pipe = diffusers.MarigoldIntrinsicsPipeline.from_pretrained(
    "prs-eth/marigold-iid-lighting-v1-1", variant="fp16", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

intrinsics = pipe(image)

vis = pipe.image_processor.visualize_intrinsics(intrinsics.prediction, pipe.target_properties)
vis[0]["albedo"].save("einstein_albedo.png")
vis[0]["shading"].save("einstein_shading.png")
vis[0]["residual"].save("einstein_residual.png")

Обе модели используют один и тот же конвейер, поддерживая при этом различные типы декомпозиции. Точная параметризация декомпозиции (например, sRGB или линейное пространство) хранится в словаре pipe.target_properties, который передается в функцию visualize_intrinsics().

Ниже приведены примеры, демонстрирующие результаты предсказанного разложения. Все модальности можно просмотреть в пространстве Intrinsic Image Decomposition.

Ускорение вывода

Приведённые выше примеры быстрого старта уже оптимизированы по качеству и скорости: они загружают контрольную точку, используют вариант весов и вычислений fp16 и выполняют стандартное количество (4) шагов диффузии для удаления шума. Первый шаг к ускорению инференса, ценой снижения качества предсказаний, — это сократить количество шагов диффузии для удаления шума до минимума:

  import diffusers
  import torch

  pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
      "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
  ).to("cuda")  # or "mps", "xpu", "cpu"

  image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
  
- depth = pipe(image)
+ depth = pipe(image, num_inference_steps=1)

Благодаря этому изменению вызов pipe завершается за 280 мс на GPU RTX 3090. Внутри входное изображение сначала кодируется с помощью кодировщика VAE Stable Diffusion, после чего выполняется один шаг денойзинга, выполняемый U-Net. Наконец, предсказанное латентное представление декодируется декодером VAE в пиксельное пространство. В такой конфигурации два из трёх вызовов модулей предназначены для преобразования между пиксельным и латентным пространствами LDM. Поскольку латентное пространство Marigold совместимо со Stable Diffusion 2.0, инференс можно ускорить более чем в 3 раза, сократив время вызова до 85 мс на RTX 3090, используя лёгкую замену SD VAE. Обратите внимание, что использование лёгкого VAE может незначительно снизить визуальное качество предсказаний.

  import diffusers
  import torch

  pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
      "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
  ).to("cuda")  # or "mps", "xpu", "cpu"

+ pipe.vae = diffusers.AutoencoderTiny.from_pretrained(
+     "madebyollin/taesd", dtype=torch.float16
+ ).cuda()

  image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

  depth = pipe(image, num_inference_steps=1)

До сих пор мы оптимизировали количество шагов диффузии и компоненты модели. Операции self-attention занимают значительную долю вычислений. Их ускорение можно достичь с помощью более эффективного attention-процессора:

  import diffusers
  import torch
+ from diffusers.models.attention_processor import AttnProcessor2_0

  pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
      "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
  ).to("cuda")  # or "mps", "xpu", "cpu"

+ pipe.vae.set_attn_processor(AttnProcessor2_0()) 
+ pipe.unet.set_attn_processor(AttnProcessor2_0())

  image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

  depth = pipe(image, num_inference_steps=1)

Наконец, как предложено в разделе Оптимизации, включение torch.compile может дополнительно повысить производительность в зависимости от целевого оборудования. Однаком компиляция приводит к значительным накладным расходам при первом вызове конвейера, что делает её полезной только в том случае, если один и тот же экземпляр конвейера вызывается многократно, например, внутри цикла.

  import diffusers
  import torch
  from diffusers.models.attention_processor import AttnProcessor2_0

  pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
      "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
  ).to("cuda")  # or "mps", "xpu", "cpu"

  pipe.vae.set_attn_processor(AttnProcessor2_0()) 
  pipe.unet.set_attn_processor(AttnProcessor2_0())

+ pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
+ pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

  image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

  depth = pipe(image, num_inference_steps=1)

Максимизация точности и ансамблирование

Пайплайны Marigold имеют встроенный механизм ансамблирования, объединяющий несколько предсказаний из разных случайных латентов. Это грубый способ повышения точности предсказаний, использующий генеративную природу диффузии. Путь ансамблирования активируется автоматически, когда аргумент ensemble_size установлен в значение, большее или равное 3. При стремлении к максимальной точности имеет смысл одновременно настраивать num_inference_steps и ensemble_size. Рекомендуемые значения различаются в зависимости от контрольных точек, но в первую очередь зависят от типа планировщика. Эффект ансамблирования особенно хорошо виден на нормалях поверхности:

  import diffusers

  pipe = diffusers.MarigoldNormalsPipeline.from_pretrained("prs-eth/marigold-normals-v1-1").to("cuda")  # or "mps", "xpu", "cpu"

  image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

- depth = pipe(image)
+ depth = pipe(image, num_inference_steps=10, ensemble_size=5)

  vis = pipe.image_processor.visualize_normals(depth.prediction)
  vis[0].save("einstein_normals.png")

Как видно, все области с мелкомасштабной структурой, например волосы, получили более консервативные и в среднем более точные предсказания. Такой результат лучше подходит для задач, чувствительных к точности, таких как 3D-реконструкция.

Покдровая обработка видео с временной согласованностью

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

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

import imageio
import diffusers
import torch
from diffusers.models.attention_processor import AttnProcessor2_0
from PIL import Image
from tqdm import tqdm

device = "cuda"  # or "mps", "xpu", "cpu"
path_in = "https://huggingface.co/spaces/prs-eth/marigold-lcm/resolve/c7adb5427947d2680944f898cd91d386bf0d4924/files/video/obama.mp4"
path_out = "obama_depth.gif"

pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
    "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
).to(device)
pipe.vae = diffusers.AutoencoderTiny.from_pretrained(
    "madebyollin/taesd", dtype=torch.float16
).to(device)
pipe.unet.set_attn_processor(AttnProcessor2_0())
pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
pipe.set_progress_bar_config(disable=True)

with imageio.get_reader(path_in) as reader:
    size = reader.get_meta_data()['size']
    last_frame_latent = None
    latent_common = torch.randn(
        (1, 4, 768 * size[1] // (8 * max(size)), 768 * size[0] // (8 * max(size)))
    ).to(device=device, dtype=torch.float16)

    out = []
    for frame_id, frame in tqdm(enumerate(reader), desc="Processing Video"):
        frame = Image.fromarray(frame)
        latents = latent_common
        if last_frame_latent is not None:
            latents = 0.9 * latents + 0.1 * last_frame_latent

        depth = pipe(
            frame,
            num_inference_steps=1,
            match_input_resolution=False, 
            latents=latents, 
            output_latent=True,
        )
        last_frame_latent = depth.latent
        out.append(pipe.image_processor.visualize_depth(depth.prediction)[0])

    diffusers.utils.export_to_gif(out, path_out, fps=reader.get_meta_data()['fps'])

Здесь процесс диффузии начинается с заданного вычисленного латентного представления. Пайплайн устанавливает output_latent=True, чтобы получить доступ к out.latent и вычислить его вклад в инициализацию латентного представления следующего кадра. Теперь результат значительно более стабилен:

Marigold для ControlNet

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

import torch
import diffusers

device = "cuda"  # or "mps", "xpu", "cpu"
generator = torch.Generator(device=device).manual_seed(2024)
image = diffusers.utils.load_image(
    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/controlnet_depth_source.png"
)

pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
    "prs-eth/marigold-depth-v1-1", dtype=torch.float16, variant="fp16"
).to(device)

depth_image = pipe(image, generator=generator).prediction
depth_image = pipe.image_processor.visualize_depth(depth_image, color_map="binary")
depth_image[0].save("motorcycle_controlnet_depth.png")

controlnet = diffusers.ControlNetModel.from_pretrained(
    "diffusers/controlnet-depth-sdxl-1.0", dtype=torch.float16, variant="fp16"
).to(device)
pipe = diffusers.StableDiffusionXLControlNetPipeline.from_pretrained(
    "SG161222/RealVisXL_V4.0", dtype=torch.float16, variant="fp16", controlnet=controlnet
).to(device)
pipe.scheduler = diffusers.DPMSolverMultistepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)

controlnet_out = pipe(
    prompt="high quality photo of a sports bike, city",
    negative_prompt="",
    guidance_scale=6.5,
    num_inference_steps=25,
    image=depth_image,
    controlnet_conditioning_scale=0.7,
    control_guidance_end=0.7,
    generator=generator,
).images
controlnet_out[0].save("motorcycle_controlnet_out.png")

Количественная оценка

Для количественной оценки Marigold в стандартных рейтингах и бенчмарках (таких как NYU, KITTI и другие наборы данных) следуйте протоколу оценки, описанному в статье: загрузите модель с полной точностью fp32 и используйте соответствующие значения для num_inference_steps и ensemble_size. При необходимости задайте seed для обеспечения воспроизводимости. Максимальное значение batch_size обеспечит максимальную загрузку устройства.

import diffusers
import torch

device = "cuda"  # or "mps", "xpu", "cpu"
seed = 2024

generator = torch.Generator(device=device).manual_seed(seed)
pipe = diffusers.MarigoldDepthPipeline.from_pretrained("prs-eth/marigold-depth-v1-1").to(device)

image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

depth = pipe(
    image, 
    num_inference_steps=4,  # set according to the evaluation protocol from the paper
    ensemble_size=10,       # set according to the evaluation protocol from the paper
    generator=generator,
)

# evaluate metrics

Использование предиктивной неопределённости

Механизм ансамблирования, встроенный в конвейеры Marigold, объединяет несколько предсказаний, полученных из разных случайных латентов. Как побочный эффект, его можно использовать для количественной оценки эпистемической (модельной) неопределённости: достаточно указать ensemble_size больше или равно 3 и установить output_uncertainty=True. Полученная неопределённость будет доступна в поле uncertainty результата. Её можно визуализировать следующим образом:

import diffusers
import torch

pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
    "prs-eth/marigold-depth-v1-1", variant="fp16", dtype=torch.float16
).to("cuda")  # or "mps", "xpu", "cpu"

image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")

depth = pipe(
	image,
	ensemble_size=10,  # any number >= 3
	output_uncertainty=True,
)

uncertainty = pipe.image_processor.visualize_uncertainty(depth.uncertainty)
uncertainty[0].save("einstein_depth_uncertainty.png")

Интерпретация неопределённости проста: более высокие значения (белый цвет) соответствуют пикселям, на которых модели трудно делать согласованные предсказания.

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

API предсказания глубины Marigold

>>> import diffusers
>>> import torch

>>> pipe = diffusers.MarigoldDepthPipeline.from_pretrained(
...     "prs-eth/marigold-depth-v1-1", variant="fp16", torch_dtype=torch.float16
... ).to("cuda")

>>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
>>> depth = pipe(image)

>>> vis = pipe.image_processor.visualize_depth(depth.prediction)
>>> vis[0].save("einstein_depth.png")

>>> depth_16bit = pipe.image_processor.export_depth_to_16bit_png(depth.prediction)
>>> depth_16bit[0].save("einstein_depth_16bit.png")

API оценки нормалей Marigold

>>> import diffusers
>>> import torch

>>> pipe = diffusers.MarigoldNormalsPipeline.from_pretrained(
...     "prs-eth/marigold-normals-v1-1", variant="fp16", torch_dtype=torch.float16
... ).to("cuda")

>>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
>>> normals = pipe(image)

>>> vis = pipe.image_processor.visualize_normals(normals.prediction)
>>> vis[0].save("einstein_normals.png")

API Marigold Intrinsic Image Decomposition

>>> import diffusers
>>> import torch

>>> pipe = diffusers.MarigoldIntrinsicsPipeline.from_pretrained(
...     "prs-eth/marigold-iid-appearance-v1-1", variant="fp16", torch_dtype=torch.float16
... ).to("cuda")

>>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
>>> intrinsics = pipe(image)

>>> vis = pipe.image_processor.visualize_intrinsics(intrinsics.prediction, pipe.target_properties)
>>> vis[0]["albedo"].save("einstein_albedo.png")
>>> vis[0]["roughness"].save("einstein_roughness.png")
>>> vis[0]["metallicity"].save("einstein_metallicity.png")

>>> import diffusers
>>> import torch

>>> pipe = diffusers.MarigoldIntrinsicsPipeline.from_pretrained(
...     "prs-eth/marigold-iid-lighting-v1-1", variant="fp16", torch_dtype=torch.float16
... ).to("cuda")

>>> image = diffusers.utils.load_image("https://marigoldmonodepth.github.io/images/einstein.jpg")
>>> intrinsics = pipe(image)

>>> vis = pipe.image_processor.visualize_intrinsics(intrinsics.prediction, pipe.target_properties)
>>> vis[0]["albedo"].save("einstein_albedo.png")
>>> vis[0]["shading"].save("einstein_shading.png")
>>> vis[0]["residual"].save("einstein_residual.png")