Загрузка RobustVideoMatting - Загрузка исходного кода RobustVideoMatting

RobustVideoMatting

AI Исходный код

1.0.0

Скачать

Надежное видеоматирование (RVM)

Тизер

английский | 中文

Официальный репозиторий статьи «Надежное матирование видео высокого разрешения с временным руководством». RVM специально разработан для надежного матирования видеоизображения людей. В отличие от существующих нейронных моделей, обрабатывающих кадры как независимые изображения, RVM использует рекуррентную нейронную сеть для обработки видео с помощью временной памяти. RVM может выполнять матирование любых видео в режиме реального времени без дополнительных входных данных. Он обеспечивает разрешение 4K со скоростью 76 кадров в секунду и HD 104 кадра в секунду на графическом процессоре Nvidia GTX 1080 Ti. Проект разработан в ByteDance Inc.

Новости

[3 ноября 2021 г.] Исправлена ошибка в train.py.
[16 сентября 2021 г.] Код переиздан под лицензией GPL-3.0.
[25 августа 2021 г.] Опубликованы исходный код и предварительно обученные модели.
[27 июля 2021 г.] Статья принимается WACV 2022.

Шоурил

Посмотрите шоурил (YouTube, Bilibili), чтобы увидеть выступление модели.

Все кадры из видео доступны на Google Диске.

Демо

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

Скачать

Мы рекомендуем модели MobileNetv3 для большинства случаев использования. Модели ResNet50 представляют собой более крупный вариант с небольшими улучшениями производительности. Наша модель доступна в различных средах вывода. Дополнительные инструкции см. в документации вывода.

Рамки	Скачать	Примечания
PyTorch	rvm_mobilenetv3.pth rvm_resnet50.pth	Официальные веса для PyTorch. Док
TorchHub	Нечего скачивать.	Самый простой способ использовать нашу модель в вашем проекте PyTorch. Док
Торчскрипт	rvm_mobilenetv3_fp32.torchscript rvm_mobilenetv3_fp16.torchscript rvm_resnet50_fp32.torchscript rvm_resnet50_fp16.torchscript	Если вы делаете вывод на мобильных устройствах, рассмотрите возможность экспорта квантованных моделей int8 самостоятельно. Док
ОННКС	rvm_mobilenetv3_fp32.onnx rvm_mobilenetv3_fp16.onnx rvm_resnet50_fp32.onnx rvm_resnet50_fp16.onnx	Протестировано в среде выполнения ONNX с процессором и серверными модулями CUDA. В предоставленных моделях используется opset 12. Doc, Exporter.
ТензорФлоу	rvm_mobilenetv3_tf.zip rvm_resnet50_tf.zip	TensorFlow 2 Сохраненная модель. Док
TensorFlow.js	rvm_mobilenetv3_tfjs_int8.zip	Запустите модель в Интернете. Демо, Стартовый код
CoreML	rvm_mobilenetv3_1280x720_s0.375_fp16.mlмодель rvm_mobilenetv3_1280x720_s0.375_int8.mlмодель rvm_mobilenetv3_1920x1080_s0.25_fp16.mlмодель rvm_mobilenetv3_1920x1080_s0.25_int8.mlмодель	CoreML не поддерживает динамическое разрешение. Другие разрешения можно экспортировать самостоятельно. Для моделей требуется iOS 13+. `s` обозначает `downsample_ratio` . Док, Экспортер

Все модели доступны в Google Drive и Baidu Pan (код: Gym7).

Пример PyTorch

Установите зависимости:

 pip install -r требования_inference.txt

Загрузите модель:

 import torchfrom model import MattingNetworkmodel = MattingNetwork('mobilenetv3').eval().cuda() # или "resnet50"model.load_state_dict(torch.load('rvm_mobilenetv3.pth'))

Для конвертации видео мы предоставляем простой API конвертации:

 from inference import Convert_videoconvert_video(model, # Модель может находиться на любом устройстве (процессор или cuda).input_source='input.mp4', # Видеофайл или последовательность изображенийdirectory.output_type='video', # Выберите "video" " или "png_sequence"output_composition='com.mp4', # Путь к файлу, если видео; путь к каталогу, если png последовательность.output_alpha="pha.mp4", # [Необязательно] Вывод необработанного альфа-предсказания.output_foreground="fgr.mp4", # [Необязательно] Вывод необработанного переднего плана предсказания.output_video_mbps=4, # Вывод видео в Мбит/с. Не требуется для pngsequence.downsample_ratio=None, # Гиперпараметр для настройки или использования. Нет для auto.seq_chunk=12, # Обработка n кадров одновременно для лучшего параллелизма.)

Или напишите свой собственный код вывода:

 из torch.utils.data импортировать DataLoaderиз torchvision.transforms импортировать ToTensorиз inference_utils импортировать VideoReader, VideoWriterreader = VideoReader('input.mp4', Transform=ToTensor())writer = VideoWriter('output.mp4', Frame_rate=30)bgr = Torch .tensor([.47, 1, .6]).view(3, 1, 1).cuda() # Зеленый фон.rec = [Нет] * 4 # Начальные повторяющиеся состояния.downsample_ratio = 0,25 # Отрегулируйте на основе вашего видео.с torch.no_grad():for src в DataLoader(reader): # Тензор RGB нормализован до 0 ~ 1.fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio) # Цикл the recurrent States.com = fgr * pha + bgr * (1 - pha) # Комбинация с зеленым фоном. Writer.write(com) # Запись кадра.

Модели и API конвертера также доступны через TorchHub.

 # Загрузите model.model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # или "resnet50"# Converter API.convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")

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

Обучение и оценка

Пожалуйста, обратитесь к учебной документации, чтобы обучить и оценить свою собственную модель.

Скорость

Скорость измеряется с помощью inference_speed_test.py для справки.

графический процессор	dType	HD (1920x1080)	4К (3840x2160)
РТХ 3090	РП16	172 кадра в секунду	154 кадра в секунду
РТХ 2060 Супер	РП16	134 кадра в секунду	108 кадров в секунду
видеокарта GTX 1080 Ти	ФП32	104 кадра в секунду	74 кадра в секунду

Примечание 1: HD использует downsample_ratio=0.25 , 4K использует downsample_ratio=0.125 . Во всех тестах используется размер пакета 1 и фрагмент кадра 1.
Примечание 2. Графические процессоры до архитектуры Turing не поддерживают вывод FP16, поэтому GTX 1080 Ti использует FP32.
Примечание 3. Мы измеряем только тензорную пропускную способность. Ожидается, что предоставленный сценарий преобразования видео в этом репозитории будет намного медленнее, поскольку он не использует аппаратное кодирование/декодирование видео и не выполняет передачу тензора в параллельных потоках. Если вы заинтересованы в реализации аппаратного кодирования/декодирования видео в Python, обратитесь к PyNvCodec.