RobustVideoMatting herunterladen - RobustVideoMatting -Quellcode herunterladen

RobustVideoMatting

AI-Quellcode

1.0.0

Herunterladen

Robustes Video Matting (RVM)

Teaser

Englisch | 中文

Offizielles Repository für den Artikel Robust High-Resolution Video Matting with Temporal Guidance. RVM wurde speziell für robuste menschliche Videomatten entwickelt. Im Gegensatz zu bestehenden neuronalen Modellen, die Frames als unabhängige Bilder verarbeiten, verwendet RVM ein rekurrentes neuronales Netzwerk, um Videos mit zeitlichem Speicher zu verarbeiten. RVM kann ohne zusätzliche Eingaben die Mattierung aller Videos in Echtzeit durchführen. Es erreicht 4K 76FPS und HD 104FPS auf einer Nvidia GTX 1080 Ti GPU. Das Projekt wurde bei ByteDance Inc. entwickelt.

Nachricht

[03. November 2021] Ein Fehler in train.py wurde behoben.
[16. September 2021] Der Code wird unter der GPL-3.0-Lizenz erneut veröffentlicht.
[25. August 2021] Quellcode und vorab trainierte Modelle werden veröffentlicht.
[27. Juli 2021] Der Beitrag wird vom WACV 2022 angenommen.

Showreel

Sehen Sie sich das Showreel-Video (YouTube, Bilibili) an, um die Leistung des Modells zu sehen.

Das gesamte Filmmaterial im Video ist in Google Drive verfügbar.

Demo

Webcam-Demo: Führen Sie das Modell live in Ihrem Browser aus. Visualisieren Sie wiederkehrende Zustände.
Colab-Demo: Testen Sie unser Modell anhand Ihrer eigenen Videos mit kostenloser GPU.

Herunterladen

Für die meisten Anwendungsfälle empfehlen wir MobileNetv3-Modelle. ResNet50-Modelle sind die größere Variante mit kleinen Leistungsverbesserungen. Unser Modell ist auf verschiedenen Inferenz-Frameworks verfügbar. Weitere Anweisungen finden Sie in der Inferenzdokumentation.

Rahmen	Herunterladen	Notizen
PyTorch	rvm_mobilenetv3.pth rvm_resnet50.pth	Offizielle Gewichte für PyTorch. Dok
TorchHub	Nichts zum Herunterladen.	Der einfachste Weg, unser Modell in Ihrem PyTorch-Projekt zu verwenden. Dok
TorchScript	rvm_mobilenetv3_fp32.torchscript rvm_mobilenetv3_fp16.torchscript rvm_resnet50_fp32.torchscript rvm_resnet50_fp16.torchscript	Wenn die Inferenz auf Mobilgeräten erfolgt, sollten Sie in Erwägung ziehen, int8-quantisierte Modelle selbst zu exportieren. Dok
ONNX	rvm_mobilenetv3_fp32.onnx rvm_mobilenetv3_fp16.onnx rvm_resnet50_fp32.onnx rvm_resnet50_fp16.onnx	Getestet auf ONNX Runtime mit CPU und CUDA-Backends. Die bereitgestellten Modelle verwenden Opset 12. Doc, Exporter.
TensorFlow	rvm_mobilenetv3_tf.zip rvm_resnet50_tf.zip	TensorFlow 2 SavedModel. Dok
TensorFlow.js	rvm_mobilenetv3_tfjs_int8.zip	Führen Sie das Modell im Web aus. Demo, Startercode
CoreML	rvm_mobilenetv3_1280x720_s0.375_fp16.mlmodel rvm_mobilenetv3_1280x720_s0.375_int8.mlmodel rvm_mobilenetv3_1920x1080_s0.25_fp16.mlmodel rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel	CoreML unterstützt keine dynamische Auflösung. Andere Auflösungen können selbst exportiert werden. Für die Modelle ist iOS 13+ erforderlich. `s` bezeichnet `downsample_ratio` . Doc, Exporteur

Alle Modelle sind in Google Drive und Baidu Pan verfügbar (Code: gym7).

PyTorch-Beispiel

Abhängigkeiten installieren:

 pip install -r require_inference.txt

Laden Sie das Modell:

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

Zum Konvertieren von Videos stellen wir eine einfache Konvertierungs-API zur Verfügung:

 from inference import convert_videoconvert_video(model, # Das Modell kann auf jedem Gerät (CPU oder Cuda) sein.input_source='input.mp4', # Eine Videodatei oder ein Bildsequenzverzeichnis.output_type='video', # Wählen Sie „video " oder "png_sequence"output_composition='com.mp4', # Dateipfad bei Video; Verzeichnispfad bei png sequence.output_alpha="pha.mp4", # [Optional] Ausgabe die rohe Alpha-Vorhersage.output_foreground="fgr.mp4", # [Optional] Gibt die rohe Vordergrundvorhersage aus.output_video_mbps=4, # Gibt Video-MBit/s aus. Nicht erforderlich für PNG-Sequenz.downsample_ratio=None, # Ein Hyperparameter zum Anpassen oder Verwenden von None für auto.seq_chunk=12, # Verarbeiten Sie n Frames gleichzeitig für eine bessere Parallelität.)

Oder schreiben Sie Ihren eigenen Inferenzcode:

 aus Torch.utils.data Import DataLoader aus Torchvision.transforms Import ToTensor aus inference_utils Import 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() # Grüner Hintergrund.rec = [Keine] * 4 # Anfängliche wiederkehrende Zustände.downsample_ratio = 0,25 # Basierend auf Ihrem Video anpassen.mit Torch.no_grad():for src in DataLoader(reader): # RGB-Tensor normalisiert auf 0 ~ 1.fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio) # Durchlaufen Sie die wiederkehrenden Zustände.com = fgr * pha + bgr * (1 - pha) # Zusammengesetzt zum grünen Hintergrund. write.write(com) # Frame schreiben.

Die Modelle und die Konverter-API sind auch über TorchHub verfügbar.

 # Laden Sie das model.model = Torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # oder "resnet50"# Konverter API.convert_video = Torch.hub.load("PeterL1n/RobustVideoMatting", "Converter")

Einzelheiten zum downsample_ratio -Hyperparameter, weitere Konverterargumente und eine erweiterte Verwendung finden Sie in der Inferenzdokumentation.

Schulung und Bewertung

Informationen zum Trainieren und Bewerten Ihres eigenen Modells finden Sie in der Schulungsdokumentation.

Geschwindigkeit

Die Geschwindigkeit wird als Referenz mit inference_speed_test.py gemessen.

GPU	dType	HD (1920x1080)	4K (3840x2160)
RTX 3090	FP16	172 FPS	154 FPS
RTX 2060 Super	FP16	134 FPS	108 FPS
GTX 1080 Ti	FP32	104 FPS	74 FPS

Hinweis 1: HD verwendet downsample_ratio=0.25 , 4K verwendet downsample_ratio=0.125 . Alle Tests verwenden Batchgröße 1 und Frame-Chunk 1.
Hinweis 2: GPUs vor der Turing-Architektur unterstützen keine FP16-Inferenz, daher verwendet die GTX 1080 Ti FP32.
Hinweis 3: Wir messen nur den Tensordurchsatz. Das in diesem Repo bereitgestellte Videokonvertierungsskript wird voraussichtlich viel langsamer sein, da es keine Hardware-Videokodierung/-dekodierung nutzt und die Tensorübertragung nicht auf parallelen Threads erfolgt. Wenn Sie daran interessiert sind, Hardware-Videokodierung/-dekodierung in Python zu implementieren, lesen Sie bitte PyNvCodec.