Как я сделал Smart Select для Krita: локальное AI-выделение объектов по лассо

от автора

Я недавно начал пользоваться Krita, и после Фотошопа основной болью для меня было отсутствие удобного инструмента для умного выделения объектов.

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

Идея простая:

1. Пользователь обводит объект лассо.
2. Нажимает кнопку Select object with AI.
3. Плагин локально строит мягкую alpha-маску.
4. Krita получает обычное выделение, с которым дальше можно работать штатными
   инструментами.

Назвать плагин я решил Krita Smart Select.

Репозиторий:
https://github.com/BMFreed/krita-smart-select

Релизы:
https://github.com/BMFreed/krita-smart-select/releases

Зачем это нужно

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

Мне хотелось сделать инструмент, который на основе выделенной пользователем области сам понимает, как «вырезать» объект из окружения:

До

До
После (внимание на очки)

После (внимание на очки)
Пример с волосами - до

Пример с волосами — до
После

После

Как выглядит пользовательский сценарий

В текущей версии сценарий такой:

1. Выбираем инструмент Smart Select.
2. Обводим объект обычным лассо.
3. При необходимости редактируем область выделения стандартными действиями
   Krita: replace, add, subtract, intersect.
4. Нажимаем Select object with AI.
5. Получаем soft selection внутри выбранной области.

Результат жёстко ограничивается исходным пользовательским выделением и
alpha-каналом документа. То есть модель не может внезапно выделить объект за
пределами области, которую пользователь сам разрешил.

Почему не SAM

Изначально я экспериментировал с SAM2 и MAM refinement-пайплайном. На
бумаге это выглядело логично: SAM строит маску, matting-модель уточняет края.

На практике оказалось, что для пользовательского инструмента это слишком
капризно:

— SAM может недобрать или перебрать объект;
— refinement может улучшить края, а может и ухудшить. И не всегда исправляет плохой prior;
— нужно подбирать prompt-стратегию;
— pipeline становится сложным и плохо предсказуемым.

В итоге я стал копать в другую сторону — начал искать модель, которая выполнит поставленную задачу с минимумом настроек (вставил и завёл), но при этом куда эффективнее, чем SAM+MAM. Звучит как довольно амбициозный запрос, но модель такую я в итоге всё же нашёл, и это ZhengPeng7/BiRefNet_HR-matting.

Модель можно посмотреть здесь.

Плагин использует BiRefNet как direct alpha extraction model: на входе crop
вокруг пользовательского выделения, на выходе soft alpha mask.

Архитектура

Плагин состоит из двух частей:

— C++ native Krita tool;
— Python worker process.

C++ отвечает за интеграцию с canvas/tool API Krita:

— рисование лассо;
— работу с текущим выделением;
— экспорт flattened snapshot;
— экспорт document alpha;
— экспорт region mask;
— применение итоговой selection mask обратно в Krita.

Python worker отвечает за ML-часть:

— managed runtime;
— загрузку модели;
— inference;
— построение итоговой alpha mask;
— debug exports;
— structured logs.

Связь между C++ и Python сделана через QProcess и stdio NDJSON. Большие
данные через JSON не передаются: изображения и маски пишутся во временные
файлы.

Пример запроса worker-а:
{
  "id": 1,
  "method": "prepare_selection",
  "params": {
    "image_path": "/tmp/.../snapshot.png",
    "alpha_mask_path": "/tmp/.../alpha_mask.npy",
    "region_mask_path": "/tmp/.../region_mask.npy",
    "region_bbox": [100, 120, 420, 360],
    "model_id": "birefnet_hr_matting",
    "algorithm_version": "birefnet-hr-matting-lasso-v1",
    "canvas_width": 1024,
    "canvas_height": 768,
    "model_cache_dir": ".../smart_select/models",
    "runtime_cache_dir": ".../smart_select/runtime"
  }
}

На выходе worker возвращает .npy-маску:
dtype: uint8
shape: [height, width]
values: 0..255

Важно, что C++ не бинаризует маску. Значения 1..254 сохраняются как soft
alpha.

Runtime и зависимости

Плагин не кладет PyTorch и модель внутрь самого ZIP. Вместо этого при первом
использовании он создаёт managed runtime в пользовательской директории Krita.

Упрощенно структура такая:
~/.local/share/krita/smart_select/
  runtime/
    venv/
    pip-cache/
    runtime.json
  models/
    birefnet_hr_matting/
      snapshot/
      model.json

Runtime создается автоматически. Пользователю не нужно руками ставить
Python-зависимости в системный Python.

При этом модель и зависимости работают локально. Никакие изображения никуда не
отправляются.

Что под капотом

Основной flow:

1. C++ получает текущую Krita selection region.
2. Экспортирует flattened RGB snapshot.
3. Экспортирует alpha mask документа.
4. Экспортирует region mask текущего выделения.
5. Worker вырезает crop по bbox выделения.
6. BiRefNet строит alpha mask для crop.
7. Worker клипует результат:
final_alpha = birefnet_alpha & user_region & document_alpha

8. C++ применяет итоговую soft alpha mask как Krita selection.

Krita-параметры вроде Grow, Feather и Anti-aliasing остаются штатными
пост-операциями Krita. Плагин не добавляет скрытый grow и не пытается «сам за
пользователя» расширять выделение.

Observability

Я отдельно заморочился с логами, потому что debugging C++/Python/ML-плагина
внутри Krita без нормальных событий быстро превращается в археологию.

Логи сделаны в OpenTelemetry-friendly формате, но без OpenTelemetry SDK и без
отправки telemetry наружу.

Пример:
[SmartSelectWorker] event.name=alpha_model.predict.result event.domain=smart_select severity=info component=worker model.id=birefnet_hr_matting device.type=cuda duration_ms=843 selection.area=12345

Сейчас логируются:

— lifecycle worker-а;
— runtime setup;
— model cache;
— inference;
— cancellation;
— cleanup временных файлов;
— ошибки с user-safe error types.

Установка

Пока плагин заточен под:

— Linux;
— Krita 5.3.1;
— NVIDIA GPU желательно, но CPU fallback возможен;

Первый запуск может занять время из-за установки runtime и загрузки модели.

Скачать ZIP можно на странице релизов:

https://github.com/BMFreed/krita-smart-select/releases

В Krita:
Tools -> Scripts -> Import Python Plugin

Выбираем скачанный ZIP, перезапускаем Krita.

Подводя итоги

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

Сначала мне казалось, что SAM вроде как должен был эту задачу решить. Потом я стал строить сложный pipeline вокруг SAM/MAM. И только после огромной череды неудач я смог задать себе судьбоносный вопрос: «А не фигню ли я делаю?», откатиться обратно до статии ресёрча и смотреть другие классы решений, что меня в итоге и привело к BiRefNet.

Буду рад фидбеку, issue и тестовым кейсам.

Репозиторий:
https://github.com/BMFreed/krita-smart-select

ссылка на оригинал статьи https://habr.com/ru/articles/1045846/