Замена объектов на изображении с использованием Stable Diffusion, SAM и Grounding DINO
Современные технологии компьютерного зрения и генерации изображений позволяют обычному пользователю (немного разбирающемуся как использовать Google colab) заменять части объектов на изображении. Здесь будет показано, как заменить автомобиль на медведя.
Используемые технологии:
- Stable Diffusion - модели и библиотека, создающая изображения по текстовым описаниям;
- Segment Anything Model (SAM) - модель, позволяющая очень точно сегментировать практически любой объект;
- Grounding DINO - модель, позволяющая по текстовому описанию детектировать объект (естественно вместо нее можно использовать например Yolo).
Установка и инициализация компонент
Будем использовать Google colab, поэтому предполагаем, что там на машине установлено уже стандартное для colab окружение. Загружаем описанные выше компоненты и инсталлируем:
Импортируем нужные модули:
Модели SAM и Grounding DINO нуждаются в ручной загрузке:
Загружаем изображение автомобиля (конечно оно может быть другим):
Указываем устройство и загружаем в память изображение:
Инициализируем SAM:
Далее скачиваем модели для Stable Diffusion
Инициализируем Grounding DINO и базовые настройки для детектирования объекта:
Детектирование объекта
Детектирование осуществляется с помощью функции predict:
Аннотированное изображение нужно для того, чтобы plt вместо нас нарисовал объекты на изображении.
Выводим изображение :
Находим маску с помощью SAM
Служебная функция:
Находим маску:
Аннотируем маску для plt:
Выводим маску:
Только маску можно посмотреть:
Замена объекта
Заменяем автомобиль на медведя:
И выводим изображения на экран (как вы поняли уже, я везде показывал второе изображение, т.к. всегда рисовались 2 - исходное и резульат):
Вы можете поэкспериментировать с colab notebook:
Только поправьте загрузку моделей, как указано выше.
