Сегментация чего угодно с высоким качеством
Segment Anything in High Quality
June 2, 2023
Авторы: Lei Ke, Mingqiao Ye, Martin Danelljan, Yifan Liu, Yu-Wing Tai, Chi-Keung Tang, Fisher Yu
cs.AI
Аннотация
Недавно представленная модель Segment Anything Model (SAM) представляет собой значительный шаг вперед в масштабировании моделей сегментации, обеспечивая мощные возможности zero-shot и гибкость в использовании подсказок. Несмотря на обучение с использованием 1,1 миллиарда масок, качество предсказания масок SAM во многих случаях оказывается недостаточным, особенно при работе с объектами, имеющими сложную структуру. Мы предлагаем HQ-SAM, наделяя SAM способностью точно сегментировать любой объект, сохраняя при этом оригинальную дизайн-концепцию SAM, эффективность и обобщаемость zero-shot. Наш тщательно продуманный подход повторно использует и сохраняет предварительно обученные веса модели SAM, вводя лишь минимальное количество дополнительных параметров и вычислений. Мы разработали обучаемый токен High-Quality Output Token, который внедряется в декодер масок SAM и отвечает за предсказание высококачественной маски. Вместо того чтобы применять его только к признакам декодера масок, мы сначала объединяем их с ранними и финальными признаками ViT для улучшения деталей маски. Для обучения введенных нами обучаемых параметров мы составили набор данных из 44 тысяч детализированных масок из нескольких источников. HQ-SAM обучается только на введенном наборе данных из 44 тысяч масок, что занимает всего 4 часа на 8 GPU. Мы демонстрируем эффективность HQ-SAM на наборе из 9 разнообразных наборов данных для сегментации в различных задачах, где 7 из них оцениваются по протоколу zero-shot переноса. Наш код и модели будут доступны по адресу https://github.com/SysCV/SAM-HQ.
English
The recent Segment Anything Model (SAM) represents a big leap in scaling up
segmentation models, allowing for powerful zero-shot capabilities and flexible
prompting. Despite being trained with 1.1 billion masks, SAM's mask prediction
quality falls short in many cases, particularly when dealing with objects that
have intricate structures. We propose HQ-SAM, equipping SAM with the ability to
accurately segment any object, while maintaining SAM's original promptable
design, efficiency, and zero-shot generalizability. Our careful design reuses
and preserves the pre-trained model weights of SAM, while only introducing
minimal additional parameters and computation. We design a learnable
High-Quality Output Token, which is injected into SAM's mask decoder and is
responsible for predicting the high-quality mask. Instead of only applying it
on mask-decoder features, we first fuse them with early and final ViT features
for improved mask details. To train our introduced learnable parameters, we
compose a dataset of 44K fine-grained masks from several sources. HQ-SAM is
only trained on the introduced detaset of 44k masks, which takes only 4 hours
on 8 GPUs. We show the efficacy of HQ-SAM in a suite of 9 diverse segmentation
datasets across different downstream tasks, where 7 out of them are evaluated
in a zero-shot transfer protocol. Our code and models will be released at
https://github.com/SysCV/SAM-HQ.