English | 中文
+ +Official repository for the paper [Robust High-Resolution Video Matting with Temporal Guidance](https://peterl1n.github.io/RobustVideoMatting/). RVM is specifically designed for robust human video matting. Unlike existing neural models that process frames as independent images, RVM uses a recurrent neural network to process videos with temporal memory. RVM can perform matting in real-time on any videos without additional inputs. It achieves **4K 76FPS** and **HD 104FPS** on an Nvidia GTX 1080 Ti GPU. The project was developed at [ByteDance Inc.](https://www.bytedance.com/) + ++ +## News + +* [Aug 25 2021] Source code and pretrained models are published. +* [Jul 27 2021] Paper is accepted by WACV 2022. + +
+ +## Showreel +Watch the showreel video ([YouTube](https://youtu.be/Jvzltozpbpk), [Bilibili](https://www.bilibili.com/video/BV1Z3411B7g7/)) to see the model's performance. + +
+
+ 
+
+
+ + +## Demo +* [Webcam Demo](https://peterl1n.github.io/RobustVideoMatting/#/demo): Run the model live in your browser. Visualize recurrent states. +* [Colab Demo](https://colab.research.google.com/drive/10z-pNKRnVNsp0Lq9tH1J_XPZ7CBC_uHm?usp=sharing): Test our model on your own videos with free GPU. + +
+ +## Download + +We recommend MobileNetv3 models for most use cases. ResNet50 models are the larger variant with small performance improvements. Our model is available on various inference frameworks. See [inference documentation](documentation/inference.md) for more instructions. + +
| Framework | +Download | +Notes | +
| PyTorch | +
+ rvm_mobilenetv3.pth + rvm_resnet50.pth + |
+ + Official weights for PyTorch. Doc + | +
| TorchHub | ++ Nothing to Download. + | ++ Easiest way to use our model in your PyTorch project. Doc + | +
| TorchScript | +
+ rvm_mobilenetv3_fp32.torchscript + rvm_mobilenetv3_fp16.torchscript + rvm_resnet50_fp32.torchscript + rvm_resnet50_fp16.torchscript + |
+ + If inference on mobile, consider export int8 quantized models yourself. Doc + | +
| ONNX | +
+ rvm_mobilenetv3_fp32.onnx + rvm_mobilenetv3_fp16.onnx + rvm_resnet50_fp32.onnx + rvm_resnet50_fp16.onnx + |
+ + Tested on ONNX Runtime with CPU and CUDA backends. Provided models use opset 12. Doc, Exporter. + | +
| TensorFlow | +
+ rvm_mobilenetv3_tf.zip + rvm_resnet50_tf.zip + |
+ + TensorFlow 2 SavedModel. Doc + | +
| TensorFlow.js | +
+ rvm_mobilenetv3_tfjs_int8.zip + |
+ + Run the model on the web. Demo, Starter Code + | +
| 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 does not support dynamic resolution. Other resolutions can be exported yourself. Models require iOS 13+. s denotes downsample_ratio. Doc, Exporter
+ |
+
+ +## PyTorch Example + +1. Install dependencies: +```sh +pip install -r requirements_inference.txt +``` + +2. Load the model: + +```python +import torch +from model import MattingNetwork + +model = MattingNetwork('mobilenetv3').eval().cuda() # or "resnet50" +model.load_state_dict(torch.load('rvm_mobilenetv3.pth')) +``` + +3. To convert videos, we provide a simple conversion API: + +```python +from inference import convert_video + +convert_video( + model, # The model, can be on any device (cpu or cuda). + input_source='input.mp4', # A video file or an image sequence directory. + output_type='video', # Choose "video" or "png_sequence" + output_composition='output.mp4', # File path if video; directory path if png sequence. + output_video_mbps=4, # Output video mbps. Not needed for png sequence. + downsample_ratio=None, # A hyperparameter to adjust or use None for auto. + seq_chunk=12, # Process n frames at once for better parallelism. +) +``` + +4. Or write your own inference code: +```python +from torch.utils.data import DataLoader +from torchvision.transforms import ToTensor +from inference_utils import VideoReader, VideoWriter + +reader = VideoReader('input.mp4', transform=ToTensor()) +writer = VideoWriter('output.mp4', frame_rate=30) + +bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1).cuda() # Green background. +rec = [None] * 4 # Initial recurrent states. +downsample_ratio = 0.25 # Adjust based on your video. + +with torch.no_grad(): + for src in DataLoader(reader): # RGB tensor normalized to 0 ~ 1. + fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio) # Cycle the recurrent states. + com = fgr * pha + bgr * (1 - pha) # Composite to green background. + writer.write(com) # Write frame. +``` + +5. The models and converter API are also available through TorchHub. + +```python +# Load the model. +model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # or "resnet50" + +# Converter API. +convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter") +``` + +Please see [inference documentation](documentation/inference.md) for details on `downsample_ratio` hyperparameter, more converter arguments, and more advanced usage. + +
+ +## Training and Evaluation + +Please refer to the [training documentation](documentation/training.md) to train and evaluate your own model. + +
+ +## Speed + +Speed is measured with `inference_speed_test.py` for reference. + +| 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 | + +* Note 1: HD uses `downsample_ratio=0.25`, 4K uses `downsample_ratio=0.125`. All tests use batch size 1 and frame chunk 1. +* Note 2: GPUs before Turing architecture does not support FP16 inference, so GTX 1080 Ti uses FP32. +* Note 3: We only measure tensor throughput. The provided video conversion script in this repo is expected to be much slower, because it does not utilize hardware video encoding/decoding and does not have the tensor transfer done on parallel threads. If you are interested in implementing hardware video encoding/decoding in Python, please refer to [PyNvCodec](https://github.com/NVIDIA/VideoProcessingFramework). + +
+ +## Project Members +* [Shanchuan Lin](https://www.linkedin.com/in/shanchuanlin/) +* [Linjie Yang](https://sites.google.com/site/linjieyang89/) +* [Imran Saleemi](https://www.linkedin.com/in/imran-saleemi/) +* [Soumyadip Sengupta](https://homes.cs.washington.edu/~soumya91/) diff --git a/README_zh_Hans.md b/README_zh_Hans.md new file mode 100644 index 0000000..5cdfd72 --- /dev/null +++ b/README_zh_Hans.md @@ -0,0 +1,224 @@ +# 稳定视频抠像 (RVM) + + + +
English | 中文
+ +论文 [Robust High-Resolution Video Matting with Temporal Guidance](https://peterl1n.github.io/RobustVideoMatting/) 的官方 GitHub 库。RVM 专为稳定人物视频抠像设计。不同于现有神经网络将每一帧作为单独图片处理，RVM 使用循环神经网络，在处理视频流时有时间记忆。RVM 可在任意视频上做实时高清抠像。在 Nvidia GTX 1080Ti 上实现 **4K 76FPS** 和 **HD 104FPS**。此研究项目来自[字节跳动](https://www.bytedance.com/)。 + ++ +## 更新 + +* [2021年8月25日] 公开代码和模型。 +* [2021年7月27日] 论文被 WACV 2022 收录。 + +
+ +## 展示视频 +观看展示视频 ([YouTube](https://youtu.be/Jvzltozpbpk), [Bilibili](https://www.bilibili.com/video/BV1Z3411B7g7/))，了解模型能力。 +
+
+
+
+
+ + +## Demo +* [网页](https://peterl1n.github.io/RobustVideoMatting/#/demo): 在浏览器里看摄像头抠像效果，展示模型内部循环记忆值。 +* [Colab](https://colab.research.google.com/drive/10z-pNKRnVNsp0Lq9tH1J_XPZ7CBC_uHm?usp=sharing): 用我们的模型转换你的视频。 + +
+ +## 下载 + +推荐在通常情况下使用 MobileNetV3 的模型。ResNet50 的模型大很多，效果稍有提高。我们的模型支持很多框架。详情请阅读[推断文档](documentation/inference_zh_Hans.md)。 + +
+ +## PyTorch 范例 + +1. 安装 Python 库: +```sh +pip install -r requirements_inference.txt +``` + +2. 加载模型: + +```python +import torch +from model import MattingNetwork + +model = MattingNetwork('mobilenetv3').eval().cuda() # 或 "resnet50" +model.load_state_dict(torch.load('rvm_mobilenetv3.pth')) +``` + +3. 若只需要做视频抠像处理，我们提供简单的 API: + +```python +from inference import convert_video + +convert_video( + model, # 模型，可以加载到任何设备（cpu 或 cuda） + input_source='input.mp4', # 视频文件，或图片序列文件夹 + output_type='video', # 可选 "video"（视频）或 "png_sequence"（PNG 序列） + output_composition='output.mp4', # 若导出视频，提供文件路径。若导出 PNG 序列，提供文件夹路径 + output_video_mbps=4, # 若导出视频，提供视频码率 + downsample_ratio=None, # 下采样比，可根据具体视频调节，或 None 选择自动 + seq_chunk=12, # 设置多帧并行计算 +) +``` + +4. 或自己写推断逻辑: +```python +from torch.utils.data import DataLoader +from torchvision.transforms import ToTensor +from inference_utils import VideoReader, VideoWriter + +reader = VideoReader('input.mp4', transform=ToTensor()) +writer = VideoWriter('output.mp4', frame_rate=30) + +bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1).cuda() # 绿背景 +rec = [None] * 4 # 初始循环记忆（Recurrent States） +downsample_ratio = 0.25 # 下采样比，根据视频调节 + +with torch.no_grad(): + for src in DataLoader(reader): # 输入张量，RGB通道，范围为 0～1 + fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio) # 将上一帧的记忆给下一帧 + com = fgr * pha + bgr * (1 - pha) # 将前景合成到绿色背景 + writer.write(com) # 输出帧 +``` + +5. 模型和 API 也可通过 TorchHub 快速载入。 + +```python +# 加载模型 +model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # 或 "resnet50" + +# 转换 API +convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter") +``` + +[推断文档](documentation/inference_zh_Hans.md)里有对 `downsample_ratio` 参数，API 使用，和高阶使用的讲解。 + +
+ +## 训练和评估 + +请参照[训练文档（英文）](documentation/training.md)。 + +
+ +## 速度 + +速度用 `inference_speed_test.py` 测量以供参考。 + +| 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 | + +* 注释1：HD 使用 `downsample_ratio=0.25`，4K 使用 `downsample_ratio=0.125`。 所有测试都使用 batch size 1 和 frame chunk 1。 +* 注释2：图灵架构之前的 GPU 不支持 FP16 推理，所以 GTX 1080 Ti 使用 FP32。 +* 注释3：我们只测量张量吞吐量（tensor throughput）。 提供的视频转换脚本会慢得多，因为它不使用硬件视频编码/解码，也没有在并行线程上完成张量传输。如果您有兴趣在 Python 中实现硬件视频编码/解码，请参考 [PyNvCodec](https://github.com/NVIDIA/VideoProcessingFramework)。 + +
+ +## 项目成员 +* [Shanchuan Lin](https://www.linkedin.com/in/shanchuanlin/) +* [Linjie Yang](https://sites.google.com/site/linjieyang89/) +* [Imran Saleemi](https://www.linkedin.com/in/imran-saleemi/) +* [Soumyadip Sengupta](https://homes.cs.washington.edu/~soumya91/) diff --git a/dataset/augmentation.py b/dataset/augmentation.py new file mode 100644 index 0000000..4a12869 --- /dev/null +++ b/dataset/augmentation.py @@ -0,0 +1,260 @@ +import easing_functions as ef +import random +import torch +from torchvision import transforms +from torchvision.transforms import functional as F + + +class MotionAugmentation: + def __init__(self, + size, + prob_fgr_affine, + prob_bgr_affine, + prob_noise, + prob_color_jitter, + prob_grayscale, + prob_sharpness, + prob_blur, + prob_hflip, + prob_pause, + static_affine=True, + aspect_ratio_range=(0.9, 1.1)): + self.size = size + self.prob_fgr_affine = prob_fgr_affine + self.prob_bgr_affine = prob_bgr_affine + self.prob_noise = prob_noise + self.prob_color_jitter = prob_color_jitter + self.prob_grayscale = prob_grayscale + self.prob_sharpness = prob_sharpness + self.prob_blur = prob_blur + self.prob_hflip = prob_hflip + self.prob_pause = prob_pause + self.static_affine = static_affine + self.aspect_ratio_range = aspect_ratio_range + + def __call__(self, fgrs, phas, bgrs): + # Foreground affine + if random.random() < self.prob_fgr_affine: + fgrs, phas = self._motion_affine(fgrs, phas) + + # Background affine + if random.random() < self.prob_bgr_affine / 2: + bgrs = self._motion_affine(bgrs) + if random.random() < self.prob_bgr_affine / 2: + fgrs, phas, bgrs = self._motion_affine(fgrs, phas, bgrs) + + # Still Affine + if self.static_affine: + fgrs, phas = self._static_affine(fgrs, phas, scale_ranges=(0.5, 1)) + bgrs = self._static_affine(bgrs, scale_ranges=(1, 1.5)) + + # To tensor + fgrs = torch.stack([F.to_tensor(fgr) for fgr in fgrs]) + phas = torch.stack([F.to_tensor(pha) for pha in phas]) + bgrs = torch.stack([F.to_tensor(bgr) for bgr in bgrs]) + + # Resize + params = transforms.RandomResizedCrop.get_params(fgrs, scale=(1, 1), ratio=self.aspect_ratio_range) + fgrs = F.resized_crop(fgrs, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + phas = F.resized_crop(phas, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + params = transforms.RandomResizedCrop.get_params(bgrs, scale=(1, 1), ratio=self.aspect_ratio_range) + bgrs = F.resized_crop(bgrs, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + + # Horizontal flip + if random.random() < self.prob_hflip: + fgrs = F.hflip(fgrs) + phas = F.hflip(phas) + if random.random() < self.prob_hflip: + bgrs = F.hflip(bgrs) + + # Noise + if random.random() < self.prob_noise: + fgrs, bgrs = self._motion_noise(fgrs, bgrs) + + # Color jitter + if random.random() < self.prob_color_jitter: + fgrs = self._motion_color_jitter(fgrs) + if random.random() < self.prob_color_jitter: + bgrs = self._motion_color_jitter(bgrs) + + # Grayscale + if random.random() < self.prob_grayscale: + fgrs = F.rgb_to_grayscale(fgrs, num_output_channels=3).contiguous() + bgrs = F.rgb_to_grayscale(bgrs, num_output_channels=3).contiguous() + + # Sharpen + if random.random() < self.prob_sharpness: + sharpness = random.random() * 8 + fgrs = F.adjust_sharpness(fgrs, sharpness) + phas = F.adjust_sharpness(phas, sharpness) + bgrs = F.adjust_sharpness(bgrs, sharpness) + + # Blur + if random.random() < self.prob_blur / 3: + fgrs, phas = self._motion_blur(fgrs, phas) + if random.random() < self.prob_blur / 3: + bgrs = self._motion_blur(bgrs) + if random.random() < self.prob_blur / 3: + fgrs, phas, bgrs = self._motion_blur(fgrs, phas, bgrs) + + # Pause + if random.random() < self.prob_pause: + fgrs, phas, bgrs = self._motion_pause(fgrs, phas, bgrs) + + return fgrs, phas, bgrs + + def _static_affine(self, *imgs, scale_ranges): + params = transforms.RandomAffine.get_params( + degrees=(-10, 10), translate=(0.1, 0.1), scale_ranges=scale_ranges, + shears=(-5, 5), img_size=imgs[0][0].size) + imgs = [[F.affine(t, *params, F.InterpolationMode.BILINEAR) for t in img] for img in imgs] + return imgs if len(imgs) > 1 else imgs[0] + + def _motion_affine(self, *imgs): + config = dict(degrees=(-10, 10), translate=(0.1, 0.1), + scale_ranges=(0.9, 1.1), shears=(-5, 5), img_size=imgs[0][0].size) + angleA, (transXA, transYA), scaleA, (shearXA, shearYA) = transforms.RandomAffine.get_params(**config) + angleB, (transXB, transYB), scaleB, (shearXB, shearYB) = transforms.RandomAffine.get_params(**config) + + T = len(imgs[0]) + easing = random_easing_fn() + for t in range(T): + percentage = easing(t / (T - 1)) + angle = lerp(angleA, angleB, percentage) + transX = lerp(transXA, transXB, percentage) + transY = lerp(transYA, transYB, percentage) + scale = lerp(scaleA, scaleB, percentage) + shearX = lerp(shearXA, shearXB, percentage) + shearY = lerp(shearYA, shearYB, percentage) + for img in imgs: + img[t] = F.affine(img[t], angle, (transX, transY), scale, (shearX, shearY), F.InterpolationMode.BILINEAR) + return imgs if len(imgs) > 1 else imgs[0] + + def _motion_noise(self, *imgs): + grain_size = random.random() * 3 + 1 # range 1 ~ 4 + monochrome = random.random() < 0.5 + for img in imgs: + T, C, H, W = img.shape + noise = torch.randn((T, 1 if monochrome else C, round(H / grain_size), round(W / grain_size))) + noise.mul_(random.random() * 0.2 / grain_size) + if grain_size != 1: + noise = F.resize(noise, (H, W)) + img.add_(noise).clamp_(0, 1) + return imgs if len(imgs) > 1 else imgs[0] + + def _motion_color_jitter(self, *imgs): + brightnessA, brightnessB, contrastA, contrastB, saturationA, saturationB, hueA, hueB \ + = torch.randn(8).mul(0.1).tolist() + strength = random.random() * 0.2 + easing = random_easing_fn() + T = len(imgs[0]) + for t in range(T): + percentage = easing(t / (T - 1)) * strength + for img in imgs: + img[t] = F.adjust_brightness(img[t], max(1 + lerp(brightnessA, brightnessB, percentage), 0.1)) + img[t] = F.adjust_contrast(img[t], max(1 + lerp(contrastA, contrastB, percentage), 0.1)) + img[t] = F.adjust_saturation(img[t], max(1 + lerp(brightnessA, brightnessB, percentage), 0.1)) + img[t] = F.adjust_hue(img[t], min(0.5, max(-0.5, lerp(hueA, hueB, percentage) * 0.1))) + return imgs if len(imgs) > 1 else imgs[0] + + def _motion_blur(self, *imgs): + blurA = random.random() * 10 + blurB = random.random() * 10 + + T = len(imgs[0]) + easing = random_easing_fn() + for t in range(T): + percentage = easing(t / (T - 1)) + blur = max(lerp(blurA, blurB, percentage), 0) + if blur != 0: + kernel_size = int(blur * 2) + if kernel_size % 2 == 0: + kernel_size += 1 # Make kernel_size odd + for img in imgs: + img[t] = F.gaussian_blur(img[t], kernel_size, sigma=blur) + + return imgs if len(imgs) > 1 else imgs[0] + + def _motion_pause(self, *imgs): + T = len(imgs[0]) + pause_frame = random.choice(range(T - 1)) + pause_length = random.choice(range(T - pause_frame)) + for img in imgs: + img[pause_frame + 1 : pause_frame + pause_length] = img[pause_frame] + return imgs if len(imgs) > 1 else imgs[0] + + +def lerp(a, b, percentage): + return a * (1 - percentage) + b * percentage + + +def random_easing_fn(): + if random.random() < 0.2: + return ef.LinearInOut() + else: + return random.choice([ + ef.BackEaseIn, + ef.BackEaseOut, + ef.BackEaseInOut, + ef.BounceEaseIn, + ef.BounceEaseOut, + ef.BounceEaseInOut, + ef.CircularEaseIn, + ef.CircularEaseOut, + ef.CircularEaseInOut, + ef.CubicEaseIn, + ef.CubicEaseOut, + ef.CubicEaseInOut, + ef.ExponentialEaseIn, + ef.ExponentialEaseOut, + ef.ExponentialEaseInOut, + ef.ElasticEaseIn, + ef.ElasticEaseOut, + ef.ElasticEaseInOut, + ef.QuadEaseIn, + ef.QuadEaseOut, + ef.QuadEaseInOut, + ef.QuarticEaseIn, + ef.QuarticEaseOut, + ef.QuarticEaseInOut, + ef.QuinticEaseIn, + ef.QuinticEaseOut, + ef.QuinticEaseInOut, + ef.SineEaseIn, + ef.SineEaseOut, + ef.SineEaseInOut, + Step, + ])() + +class Step: # Custom easing function for sudden change. + def __call__(self, value): + return 0 if value < 0.5 else 1 + + +# ---------------------------- Frame Sampler ---------------------------- + + +class TrainFrameSampler: + def __init__(self, speed=[0.5, 1, 2, 3, 4, 5]): + self.speed = speed + + def __call__(self, seq_length): + frames = list(range(seq_length)) + + # Speed up + speed = random.choice(self.speed) + frames = [int(f * speed) for f in frames] + + # Shift + shift = random.choice(range(seq_length)) + frames = [f + shift for f in frames] + + # Reverse + if random.random() < 0.5: + frames = frames[::-1] + + return frames + +class ValidFrameSampler: + def __call__(self, seq_length): + return range(seq_length) diff --git a/dataset/coco.py b/dataset/coco.py new file mode 100644 index 0000000..3c183cb --- /dev/null +++ b/dataset/coco.py @@ -0,0 +1,103 @@ +import os +import numpy as np +import random +import json +import os +from torch.utils.data import Dataset +from torchvision import transforms +from torchvision.transforms import functional as F +from PIL import Image + + +class CocoPanopticDataset(Dataset): + def __init__(self, + imgdir: str, + anndir: str, + annfile: str, + transform=None): + with open(annfile) as f: + self.data = json.load(f)['annotations'] + self.data = list(filter(lambda data: any(info['category_id'] == 1 for info in data['segments_info']), self.data)) + self.imgdir = imgdir + self.anndir = anndir + self.transform = transform + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + data = self.data[idx] + img = self._load_img(data) + seg = self._load_seg(data) + + if self.transform is not None: + img, seg = self.transform(img, seg) + + return img, seg + + def _load_img(self, data): + with Image.open(os.path.join(self.imgdir, data['file_name'].replace('.png', '.jpg'))) as img: + return img.convert('RGB') + + def _load_seg(self, data): + with Image.open(os.path.join(self.anndir, data['file_name'])) as ann: + ann.load() + + ann = np.array(ann, copy=False).astype(np.int32) + ann = ann[:, :, 0] + 256 * ann[:, :, 1] + 256 * 256 * ann[:, :, 2] + seg = np.zeros(ann.shape, np.uint8) + + for segments_info in data['segments_info']: + if segments_info['category_id'] in [1, 27, 32]: # person, backpack, tie + seg[ann == segments_info['id']] = 255 + + return Image.fromarray(seg) + + +class CocoPanopticTrainAugmentation: + def __init__(self, size): + self.size = size + self.jitter = transforms.ColorJitter(0.1, 0.1, 0.1, 0.1) + + def __call__(self, img, seg): + # Affine + params = transforms.RandomAffine.get_params(degrees=(-20, 20), translate=(0.1, 0.1), + scale_ranges=(1, 1), shears=(-10, 10), img_size=img.size) + img = F.affine(img, *params, interpolation=F.InterpolationMode.BILINEAR) + seg = F.affine(seg, *params, interpolation=F.InterpolationMode.NEAREST) + + # Resize + params = transforms.RandomResizedCrop.get_params(img, scale=(0.5, 1), ratio=(0.7, 1.3)) + img = F.resized_crop(img, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + seg = F.resized_crop(seg, *params, self.size, interpolation=F.InterpolationMode.NEAREST) + + # Horizontal flip + if random.random() < 0.5: + img = F.hflip(img) + seg = F.hflip(seg) + + # Color jitter + img = self.jitter(img) + + # To tensor + img = F.to_tensor(img) + seg = F.to_tensor(seg) + + return img, seg + + +class CocoPanopticValidAugmentation: + def __init__(self, size): + self.size = size + + def __call__(self, img, seg): + # Resize + params = transforms.RandomResizedCrop.get_params(img, scale=(1, 1), ratio=(1., 1.)) + img = F.resized_crop(img, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + seg = F.resized_crop(seg, *params, self.size, interpolation=F.InterpolationMode.NEAREST) + + # To tensor + img = F.to_tensor(img) + seg = F.to_tensor(seg) + + return img, seg \ No newline at end of file diff --git a/dataset/imagematte.py b/dataset/imagematte.py new file mode 100644 index 0000000..124faca --- /dev/null +++ b/dataset/imagematte.py @@ -0,0 +1,98 @@ +import os +import random +from torch.utils.data import Dataset +from PIL import Image + +from .augmentation import MotionAugmentation + + +class ImageMatteDataset(Dataset): + def __init__(self, + imagematte_dir, + background_image_dir, + background_video_dir, + size, + seq_length, + seq_sampler, + transform): + self.imagematte_dir = imagematte_dir + self.imagematte_files = os.listdir(os.path.join(imagematte_dir, 'fgr')) + self.background_image_dir = background_image_dir + self.background_image_files = os.listdir(background_image_dir) + self.background_video_dir = background_video_dir + self.background_video_clips = os.listdir(background_video_dir) + self.background_video_frames = [sorted(os.listdir(os.path.join(background_video_dir, clip))) + for clip in self.background_video_clips] + self.seq_length = seq_length + self.seq_sampler = seq_sampler + self.size = size + self.transform = transform + + def __len__(self): + return max(len(self.imagematte_files), len(self.background_image_files) + len(self.background_video_clips)) + + def __getitem__(self, idx): + if random.random() < 0.5: + bgrs = self._get_random_image_background() + else: + bgrs = self._get_random_video_background() + + fgrs, phas = self._get_imagematte(idx) + + if self.transform is not None: + return self.transform(fgrs, phas, bgrs) + + return fgrs, phas, bgrs + + def _get_imagematte(self, idx): + with Image.open(os.path.join(self.imagematte_dir, 'fgr', self.imagematte_files[idx % len(self.imagematte_files)])) as fgr, \ + Image.open(os.path.join(self.imagematte_dir, 'pha', self.imagematte_files[idx % len(self.imagematte_files)])) as pha: + fgr = self._downsample_if_needed(fgr.convert('RGB')) + pha = self._downsample_if_needed(pha.convert('L')) + fgrs = [fgr] * self.seq_length + phas = [pha] * self.seq_length + return fgrs, phas + + def _get_random_image_background(self): + with Image.open(os.path.join(self.background_image_dir, self.background_image_files[random.choice(range(len(self.background_image_files)))])) as bgr: + bgr = self._downsample_if_needed(bgr.convert('RGB')) + bgrs = [bgr] * self.seq_length + return bgrs + + def _get_random_video_background(self): + clip_idx = random.choice(range(len(self.background_video_clips))) + frame_count = len(self.background_video_frames[clip_idx]) + frame_idx = random.choice(range(max(1, frame_count - self.seq_length))) + clip = self.background_video_clips[clip_idx] + bgrs = [] + for i in self.seq_sampler(self.seq_length): + frame_idx_t = frame_idx + i + frame = self.background_video_frames[clip_idx][frame_idx_t % frame_count] + with Image.open(os.path.join(self.background_video_dir, clip, frame)) as bgr: + bgr = self._downsample_if_needed(bgr.convert('RGB')) + bgrs.append(bgr) + return bgrs + + def _downsample_if_needed(self, img): + w, h = img.size + if min(w, h) > self.size: + scale = self.size / min(w, h) + w = int(scale * w) + h = int(scale * h) + img = img.resize((w, h)) + return img + +class ImageMatteAugmentation(MotionAugmentation): + def __init__(self, size): + super().__init__( + size=size, + prob_fgr_affine=0.95, + prob_bgr_affine=0.3, + prob_noise=0.05, + prob_color_jitter=0.3, + prob_grayscale=0.03, + prob_sharpness=0.05, + prob_blur=0.02, + prob_hflip=0.5, + prob_pause=0.03, + ) diff --git a/dataset/spd.py b/dataset/spd.py new file mode 100644 index 0000000..7f76b82 --- /dev/null +++ b/dataset/spd.py @@ -0,0 +1,27 @@ +import os +from torch.utils.data import Dataset +from PIL import Image + + +class SuperviselyPersonDataset(Dataset): + def __init__(self, imgdir, segdir, transform=None): + self.img_dir = imgdir + self.img_files = sorted(os.listdir(imgdir)) + self.seg_dir = segdir + self.seg_files = sorted(os.listdir(segdir)) + assert len(self.img_files) == len(self.seg_files) + self.transform = transform + + def __len__(self): + return len(self.img_files) + + def __getitem__(self, idx): + with Image.open(os.path.join(self.img_dir, self.img_files[idx])) as img, \ + Image.open(os.path.join(self.seg_dir, self.seg_files[idx])) as seg: + img = img.convert('RGB') + seg = seg.convert('L') + + if self.transform is not None: + img, seg = self.transform(img, seg) + + return img, seg diff --git a/dataset/videomatte.py b/dataset/videomatte.py new file mode 100644 index 0000000..555911b --- /dev/null +++ b/dataset/videomatte.py @@ -0,0 +1,125 @@ +import os +import random +from torch.utils.data import Dataset +from PIL import Image + +from .augmentation import MotionAugmentation + + +class VideoMatteDataset(Dataset): + def __init__(self, + videomatte_dir, + background_image_dir, + background_video_dir, + size, + seq_length, + seq_sampler, + transform=None): + self.background_image_dir = background_image_dir + self.background_image_files = os.listdir(background_image_dir) + self.background_video_dir = background_video_dir + self.background_video_clips = sorted(os.listdir(background_video_dir)) + self.background_video_frames = [sorted(os.listdir(os.path.join(background_video_dir, clip))) + for clip in self.background_video_clips] + + self.videomatte_dir = videomatte_dir + self.videomatte_clips = sorted(os.listdir(os.path.join(videomatte_dir, 'fgr'))) + self.videomatte_frames = [sorted(os.listdir(os.path.join(videomatte_dir, 'fgr', clip))) + for clip in self.videomatte_clips] + self.videomatte_idx = [(clip_idx, frame_idx) + for clip_idx in range(len(self.videomatte_clips)) + for frame_idx in range(0, len(self.videomatte_frames[clip_idx]), seq_length)] + self.size = size + self.seq_length = seq_length + self.seq_sampler = seq_sampler + self.transform = transform + + def __len__(self): + return len(self.videomatte_idx) + + def __getitem__(self, idx): + if random.random() < 0.5: + bgrs = self._get_random_image_background() + else: + bgrs = self._get_random_video_background() + + fgrs, phas = self._get_videomatte(idx) + + if self.transform is not None: + return self.transform(fgrs, phas, bgrs) + + return fgrs, phas, bgrs + + def _get_random_image_background(self): + with Image.open(os.path.join(self.background_image_dir, random.choice(self.background_image_files))) as bgr: + bgr = self._downsample_if_needed(bgr.convert('RGB')) + bgrs = [bgr] * self.seq_length + return bgrs + + def _get_random_video_background(self): + clip_idx = random.choice(range(len(self.background_video_clips))) + frame_count = len(self.background_video_frames[clip_idx]) + frame_idx = random.choice(range(max(1, frame_count - self.seq_length))) + clip = self.background_video_clips[clip_idx] + bgrs = [] + for i in self.seq_sampler(self.seq_length): + frame_idx_t = frame_idx + i + frame = self.background_video_frames[clip_idx][frame_idx_t % frame_count] + with Image.open(os.path.join(self.background_video_dir, clip, frame)) as bgr: + bgr = self._downsample_if_needed(bgr.convert('RGB')) + bgrs.append(bgr) + return bgrs + + def _get_videomatte(self, idx): + clip_idx, frame_idx = self.videomatte_idx[idx] + clip = self.videomatte_clips[clip_idx] + frame_count = len(self.videomatte_frames[clip_idx]) + fgrs, phas = [], [] + for i in self.seq_sampler(self.seq_length): + frame = self.videomatte_frames[clip_idx][(frame_idx + i) % frame_count] + with Image.open(os.path.join(self.videomatte_dir, 'fgr', clip, frame)) as fgr, \ + Image.open(os.path.join(self.videomatte_dir, 'pha', clip, frame)) as pha: + fgr = self._downsample_if_needed(fgr.convert('RGB')) + pha = self._downsample_if_needed(pha.convert('L')) + fgrs.append(fgr) + phas.append(pha) + return fgrs, phas + + def _downsample_if_needed(self, img): + w, h = img.size + if min(w, h) > self.size: + scale = self.size / min(w, h) + w = int(scale * w) + h = int(scale * h) + img = img.resize((w, h)) + return img + +class VideoMatteTrainAugmentation(MotionAugmentation): + def __init__(self, size): + super().__init__( + size=size, + prob_fgr_affine=0.3, + prob_bgr_affine=0.3, + prob_noise=0.1, + prob_color_jitter=0.3, + prob_grayscale=0.02, + prob_sharpness=0.1, + prob_blur=0.02, + prob_hflip=0.5, + prob_pause=0.03, + ) + +class VideoMatteValidAugmentation(MotionAugmentation): + def __init__(self, size): + super().__init__( + size=size, + prob_fgr_affine=0, + prob_bgr_affine=0, + prob_noise=0, + prob_color_jitter=0, + prob_grayscale=0, + prob_sharpness=0, + prob_blur=0, + prob_hflip=0, + prob_pause=0, + ) diff --git a/dataset/youtubevis.py b/dataset/youtubevis.py new file mode 100644 index 0000000..babbf0e --- /dev/null +++ b/dataset/youtubevis.py @@ -0,0 +1,123 @@ +import torch +import os +import json +import numpy as np +import random +from torch.utils.data import Dataset +from PIL import Image +from torchvision import transforms +from torchvision.transforms import functional as F + + +class YouTubeVISDataset(Dataset): + def __init__(self, videodir, annfile, size, seq_length, seq_sampler, transform=None): + self.videodir = videodir + self.size = size + self.seq_length = seq_length + self.seq_sampler = seq_sampler + self.transform = transform + + with open(annfile) as f: + data = json.load(f) + + self.masks = {} + for ann in data['annotations']: + if ann['category_id'] == 26: # person + video_id = ann['video_id'] + if video_id not in self.masks: + self.masks[video_id] = [[] for _ in range(len(ann['segmentations']))] + for frame, mask in zip(self.masks[video_id], ann['segmentations']): + if mask is not None: + frame.append(mask) + + self.videos = {} + for video in data['videos']: + video_id = video['id'] + if video_id in self.masks: + self.videos[video_id] = video + + self.index = [] + for video_id in self.videos.keys(): + for frame in range(len(self.videos[video_id])): + self.index.append((video_id, frame)) + + def __len__(self): + return len(self.index) + + def __getitem__(self, idx): + video_id, frame_id = self.index[idx] + video = self.videos[video_id] + frame_count = len(self.videos[video_id]['file_names']) + H, W = video['height'], video['width'] + + imgs, segs = [], [] + for t in self.seq_sampler(self.seq_length): + frame = (frame_id + t) % frame_count + + filename = video['file_names'][frame] + masks = self.masks[video_id][frame] + + with Image.open(os.path.join(self.videodir, filename)) as img: + imgs.append(self._downsample_if_needed(img.convert('RGB'), Image.BILINEAR)) + + seg = np.zeros((H, W), dtype=np.uint8) + for mask in masks: + seg |= self._decode_rle(mask) + segs.append(self._downsample_if_needed(Image.fromarray(seg), Image.NEAREST)) + + if self.transform is not None: + imgs, segs = self.transform(imgs, segs) + + return imgs, segs + + def _decode_rle(self, rle): + H, W = rle['size'] + msk = np.zeros(H * W, dtype=np.uint8) + encoding = rle['counts'] + skip = 0 + for i in range(0, len(encoding) - 1, 2): + skip += encoding[i] + draw = encoding[i + 1] + msk[skip : skip + draw] = 255 + skip += draw + return msk.reshape(W, H).transpose() + + def _downsample_if_needed(self, img, resample): + w, h = img.size + if min(w, h) > self.size: + scale = self.size / min(w, h) + w = int(scale * w) + h = int(scale * h) + img = img.resize((w, h), resample) + return img + + +class YouTubeVISAugmentation: + def __init__(self, size): + self.size = size + self.jitter = transforms.ColorJitter(0.3, 0.3, 0.3, 0.15) + + def __call__(self, imgs, segs): + + # To tensor + imgs = torch.stack([F.to_tensor(img) for img in imgs]) + segs = torch.stack([F.to_tensor(seg) for seg in segs]) + + # Resize + params = transforms.RandomResizedCrop.get_params(imgs, scale=(0.8, 1), ratio=(0.9, 1.1)) + imgs = F.resized_crop(imgs, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + segs = F.resized_crop(segs, *params, self.size, interpolation=F.InterpolationMode.BILINEAR) + + # Color jitter + imgs = self.jitter(imgs) + + # Grayscale + if random.random() < 0.05: + imgs = F.rgb_to_grayscale(imgs, num_output_channels=3) + + # Horizontal flip + if random.random() < 0.5: + imgs = F.hflip(imgs) + segs = F.hflip(segs) + + return imgs, segs \ No newline at end of file diff --git a/documentation/image/showreel.gif b/documentation/image/showreel.gif new file mode 100644 index 0000000..5af4c45 Binary files /dev/null and b/documentation/image/showreel.gif differ diff --git a/documentation/image/teaser.gif b/documentation/image/teaser.gif new file mode 100644 index 0000000..26c9ec2 Binary files /dev/null and b/documentation/image/teaser.gif differ diff --git a/documentation/inference.md b/documentation/inference.md new file mode 100644 index 0000000..413d720 --- /dev/null +++ b/documentation/inference.md @@ -0,0 +1,352 @@ +# Inference + +
English | 中文
+ +## Content + +* [Concepts](#concepts) + * [Downsample Ratio](#downsample-ratio) + * [Recurrent States](#recurrent-states) +* [PyTorch](#pytorch) +* [TorchHub](#torchhub) +* [TorchScript](#torchscript) +* [ONNX](#onnx) +* [TensorFlow](#tensorflow) +* [TensorFlow.js](#tensorflowjs) +* [CoreML](#coreml) + ++ + +## Concepts + +### Downsample Ratio + +The table provides a general guideline. Please adjust based on your video content. + +| Resolution | Portrait | Full-Body | +| ------------- | ------------- | -------------- | +| <= 512x512 | 1 | 1 | +| 1280x720 | 0.375 | 0.6 | +| 1920x1080 | 0.25 | 0.4 | +| 3840x2160 | 0.125 | 0.2 | + +Internally, the model resizes down the input for stage 1. Then, it refines at high-resolution for stage 2. + +Set `downsample_ratio` so that the downsampled resolution is between 256 and 512. For example, for `1920x1080` input with `downsample_ratio=0.25`, the resized resolution `480x270` is between 256 and 512. + +Adjust `downsample_ratio` base on the video content. If the shot is portrait, a lower `downsample_ratio` is sufficient. If the shot contains the full human body, use high `downsample_ratio`. Note that higher `downsample_ratio` is not always better. + + +
+ +### Recurrent States +The model is a recurrent neural network. You must process frames sequentially and recycle its recurrent states. + +**Correct Way** + +The recurrent outputs are recycled back as input when processing the next frame. The states are essentially the model's memory. + +```python +rec = [None] * 4 # Initial recurrent states are None + +for frame in YOUR_VIDEO: + fgr, pha, *rec = model(frame, *rec, downsample_ratio) +``` + +**Wrong Way** + +The model does not utilize the recurrent states. Only use it to process independent images. + +```python +for frame in YOUR_VIDEO: + fgr, pha = model(frame, downsample_ratio)[:2] +``` + +More technical details are in the [paper](https://peterl1n.github.io/RobustVideoMatting/). + +
+ + +## PyTorch + +Model loading: + +```python +import torch +from model import MattingNetwork + +model = MattingNetwork(variant='mobilenetv3').eval().cuda() # Or variant="resnet50" +model.load_state_dict(torch.load('rvm_mobilenetv3.pth')) +``` + +Example inference loop: +```python +rec = [None] * 4 # Set initial recurrent states to None + +for src in YOUR_VIDEO: # src can be [B, C, H, W] or [B, T, C, H, W] + fgr, pha, *rec = model(src, *rec, downsample_ratio=0.25) +``` + +* `src`: Input frame. + * Can be of shape `[B, C, H, W]` or `[B, T, C, H, W]`. + * If `[B, T, C, H, W]`, a chunk of `T` frames can be given at once for better parallelism. + * RGB input is normalized to `0~1` range. + +* `fgr, pha`: Foreground and alpha predictions. + * Can be of shape `[B, C, H, W]` or `[B, T, C, H, W]` depends on `src`. + * `fgr` has `C=3` for RGB, `pha` has `C=1`. + * Outputs normalized to `0~1` range. +* `rec`: Recurrent states. + * Type of `List[Tensor, Tensor, Tensor, Tensor]`. + * Initial `rec` can be `List[None, None, None, None]`. + * It has 4 recurrent states because the model has 4 ConvGRU layers. + * All tensors are rank 4 regardless of `src` rank. + * If a chunk of `T` frames is given, only the last frame's recurrent states will be returned. + +To inference on video, here is a complete example: + +```python +from torch.utils.data import DataLoader +from torchvision.transforms import ToTensor +from inference_utils import VideoReader, VideoWriter + +reader = VideoReader('input.mp4', transform=ToTensor()) +writer = VideoWriter('output.mp4', frame_rate=30) + +bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1).cuda() # Green background. +rec = [None] * 4 # Initial recurrent states. + +with torch.no_grad(): + for src in DataLoader(reader): + fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio=0.25) # Cycle the recurrent states. + writer.write(fgr * pha + bgr * (1 - pha)) +``` + +Or you can use the provided video converter: + +```python +from inference import convert_video + +convert_video( + model, # The loaded model, can be on any device (cpu or cuda). + input_source='input.mp4', # A video file or an image sequence directory. + input_resize=(1920, 1080), # [Optional] Resize the input (also the output). + downsample_ratio=0.25, # [Optional] If None, make downsampled max size be 512px. + output_type='video', # Choose "video" or "png_sequence" + output_composition='com.mp4', # File path if video; directory path if png sequence. + output_alpha="pha.mp4", # [Optional] Output the raw alpha prediction. + output_foreground="fgr.mp4", # [Optional] Output the raw foreground prediction. + output_video_mbps=4, # Output video mbps. Not needed for png sequence. + seq_chunk=12, # Process n frames at once for better parallelism. + num_workers=1, # Only for image sequence input. Reader threads. + progress=True # Print conversion progress. +) +``` + +The converter can also be invoked in command line: + +```sh +python inference.py \ + --variant mobilenetv3 \ + --checkpoint "CHECKPOINT" \ + --device cuda \ + --input-source "input.mp4" \ + --downsample-ratio 0.25 \ + --output-type video \ + --output-composition "composition.mp4" \ + --output-alpha "alpha.mp4" \ + --output-foreground "foreground.mp4" \ + --output-video-mbps 4 \ + --seq-chunk 12 +``` + +
+ +## TorchHub + +Model loading: + +```python +model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # or "resnet50" +``` + +Use the conversion function. Refer to the documentation for `convert_video` function above. + +```python +convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter") + +convert_video(model, ...args...) +``` + +
+ +## TorchScript + +Model loading: + +```python +import torch +model = torch.jit.load('rvm_mobilenetv3.torchscript') +``` + +Optionally, freeze the model. This will trigger graph optimization, such as BatchNorm fusion etc. Frozen models are faster. + +```python +model = torch.jit.freeze(model) +``` + +Then, you can use the `model` exactly the same as a PyTorch model, with the exception that you must manually provide `device` and `dtype` to the converter API for frozen model. For example: + +```python +convert_video(frozen_model, ...args..., device='cuda', dtype=torch.float32) +``` + +
+ +## ONNX + +Model spec: +* Inputs: [`src`, `r1i`, `r2i`, `r3i`, `r4i`, `downsample_ratio`]. + * `src` is the RGB input frame of shape `[B, C, H, W]` normalized to `0~1` range. + * `rXi` are the recurrent state inputs. Initial recurrent states are zero value tensors of shape `[1, 1, 1, 1]`. + * `downsample_ratio` is a tensor of shape `[1]`. + * Only `downsample_ratio` must have `dtype=FP32`. Other inputs must have `dtype` matching the loaded model's precision. +* Outputs: [`fgr`, `pha`, `r1o`, `r2o`, `r3o`, `r4o`] + * `fgr, pha` are the foreground and alpha prediction. Normalized to `0~1` range. + * `rXo` are the recurrent state outputs. + +We only show examples of using onnxruntime CUDA backend in Python. + +Model loading + +```python +import onnxruntime as ort + +sess = ort.InferenceSession('rvm_mobilenetv3_fp16.onnx') +``` + +Naive inference loop + +```python +import numpy as np + +rec = [ np.zeros([1, 1, 1, 1], dtype=np.float16) ] * 4 # Must match dtype of the model. +downsample_ratio = np.array([0.25], dtype=np.float32) # dtype always FP32 + +for src in YOUR_VIDEO: # src is of [B, C, H, W] with dtype of the model. + fgr, pha, *rec = sess.run([], { + 'src': src, + 'r1i': rec[0], + 'r2i': rec[1], + 'r3i': rec[2], + 'r4i': rec[3], + 'downsample_ratio': downsample_ratio + }) +``` + +If you use GPU version of ONNX Runtime, the above naive implementation has recurrent states transferred between CPU and GPU on every frame. They could have just stayed on the GPU for better performance. Below is an example using `iobinding` to eliminate useless transfers. + +```python +import onnxruntime as ort +import numpy as np + +# Load model. +sess = ort.InferenceSession('rvm_mobilenetv3_fp16.onnx') + +# Create an io binding. +io = sess.io_binding() + +# Create tensors on CUDA. +rec = [ ort.OrtValue.ortvalue_from_numpy(np.zeros([1, 1, 1, 1], dtype=np.float16), 'cuda') ] * 4 +downsample_ratio = ort.OrtValue.ortvalue_from_numpy(np.asarray([0.25], dtype=np.float32), 'cuda') + +# Set output binding. +for name in ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o']: + io.bind_output(name, 'cuda') + +# Inference loop +for src in YOUR_VIDEO: + io.bind_cpu_input('src', src) + io.bind_ortvalue_input('r1i', rec[0]) + io.bind_ortvalue_input('r2i', rec[1]) + io.bind_ortvalue_input('r3i', rec[2]) + io.bind_ortvalue_input('r4i', rec[3]) + io.bind_ortvalue_input('downsample_ratio', downsample_ratio) + + sess.run_with_iobinding(io) + + fgr, pha, *rec = io.get_outputs() + + # Only transfer `fgr` and `pha` to CPU. + fgr = fgr.numpy() + pha = pha.numpy() +``` + +Note: depending on the inference tool you choose, it may not support all the operations in our official ONNX model. You are responsible for modifying the model code and exporting your own ONNX model. You can refer to our exporter code in the [onnx branch](https://github.com/PeterL1n/RobustVideoMatting/tree/onnx). + +
+ +### TensorFlow + +An example usage: + +```python +import tensorflow as tf + +model = tf.keras.models.load_model('rvm_mobilenetv3_tf') +model = tf.function(model) + +rec = [ tf.constant(0.) ] * 4 # Initial recurrent states. +downsample_ratio = tf.constant(0.25) # Adjust based on your video. + +for src in YOUR_VIDEO: # src is of shape [B, H, W, C], not [B, C, H, W]! + out = model([src, *rec, downsample_ratio]) + fgr, pha, *rec = out['fgr'], out['pha'], out['r1o'], out['r2o'], out['r3o'], out['r4o'] +``` + +Note the the tensors are all channel last. Otherwise, the inputs and outputs are exactly the same as PyTorch. + +We also provide the raw TensorFlow model code in the [tensorflow branch](https://github.com/PeterL1n/RobustVideoMatting/tree/tensorflow). You can transfer PyTorch checkpoint weights to TensorFlow models. + +
+ +### TensorFlow.js + +We provide a starter code in the [tfjs branch](https://github.com/PeterL1n/RobustVideoMatting/tree/tfjs). The example is very self-explanatory. It shows how to properly use the model. + +
+ +### CoreML + +We only show example usage of the CoreML models in Python API using `coremltools`. In production, the same logic can be applied in Swift. When processing the first frame, do not provide recurrent states. CoreML will internally construct zero tensors of the correct shapes as the initial recurrent states. + +```python +import coremltools as ct + +model = ct.models.model.MLModel('rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel') + +r1, r2, r3, r4 = None, None, None, None + +for src in YOUR_VIDEO: # src is PIL.Image. + + if r1 is None: + # Initial frame, do not provide recurrent states. + inputs = {'src': src} + else: + # Subsequent frames, provide recurrent states. + inputs = {'src': src, 'r1i': r1, 'r2i': r2, 'r3i': r3, 'r4i': r4} + + outputs = model.predict(inputs) + + fgr = outputs['fgr'] # PIL.Image. + pha = outputs['pha'] # PIL.Image. + + r1 = outputs['r1o'] # Numpy array. + r2 = outputs['r2o'] # Numpy array. + r3 = outputs['r3o'] # Numpy array. + r4 = outputs['r4o'] # Numpy array. + +``` + +Our CoreML models only support fixed resolutions. If you need other resolutions, you can export them yourself. See [coreml branch](https://github.com/PeterL1n/RobustVideoMatting/tree/coreml) for model export. \ No newline at end of file diff --git a/documentation/inference_zh_Hans.md b/documentation/inference_zh_Hans.md new file mode 100644 index 0000000..7995f40 --- /dev/null +++ b/documentation/inference_zh_Hans.md @@ -0,0 +1,353 @@ +# 推断文档 + +
English | 中文
+ +## 目录 + +* [概念](#概念) + * [下采样比](#下采样比) + * [循环记忆](#循环记忆) +* [PyTorch](#pytorch) +* [TorchHub](#torchhub) +* [TorchScript](#torchscript) +* [ONNX](#onnx) +* [TensorFlow](#tensorflow) +* [TensorFlow.js](#tensorflowjs) +* [CoreML](#coreml) + ++ + +## 概念 + +### 下采样比 + +该表仅供参考。可根据视频内容进行调节。 + +| 分辨率 | 人像 | 全身 | +| ------------- | ------------- | -------------- | +| <= 512x512 | 1 | 1 | +| 1280x720 | 0.375 | 0.6 | +| 1920x1080 | 0.25 | 0.4 | +| 3840x2160 | 0.125 | 0.2 | + +模型在内部将高分辨率输入缩小做初步的处理，然后再放大做细分处理。 + +建议设置 `downsample_ratio` 使缩小后的分辨率维持在 256 到 512 像素之间. 例如，`1920x1080` 的输入用 `downsample_ratio=0.25`，缩小后的分辨率 `480x270` 在 256 到 512 像素之间。 + +根据视频内容调整 `downsample_ratio`。若视频是上身人像，低 `downsample_ratio` 足矣。若视频是全身像，建议尝试更高的 `downsample_ratio`。但注意，过高的 `downsample_ratio` 反而会降低效果。 + + +
+ +### 循环记忆 +此模型是循环神经网络（Recurrent Neural Network）。必须按顺序处理视频每帧，并提供网络循环记忆。 + +**正确用法** + +循环记忆输出被传递到下一帧做输入。 + +```python +rec = [None] * 4 # 初始值设置为 None + +for frame in YOUR_VIDEO: + fgr, pha, *rec = model(frame, *rec, downsample_ratio) +``` + +**错误用法** + +没有使用循环记忆。此方法仅可用于处理单独的图片。 + +```python +for frame in YOUR_VIDEO: + fgr, pha = model(frame, downsample_ratio)[:2] +``` + +更多技术细节见[论文](https://peterl1n.github.io/RobustVideoMatting/)。 + +
+ + +## PyTorch + +载入模型： + +```python +import torch +from model import MattingNetwork + +model = MattingNetwork(variant='mobilenetv3').eval().cuda() # 或 variant="resnet50" +model.load_state_dict(torch.load('rvm_mobilenetv3.pth')) +``` + +推断循环： +```python +rec = [None] * 4 # 初始值设置为 None + +for src in YOUR_VIDEO: # src 可以是 [B, C, H, W] 或 [B, T, C, H, W] + fgr, pha, *rec = model(src, *rec, downsample_ratio=0.25) +``` + +* `src`: 输入帧（Source）。 + * 可以是 `[B, C, H, W]` 或 `[B, T, C, H, W]` 的张量。 + * 若是 `[B, T, C, H, W]`，可给模型一次 `T` 帧，做一小段一小段地处理，用于更好的并行计算。 + * RGB 通道输入，范围为 `0~1`。 + +* `fgr, pha`: 前景（Foreground）和透明度通道（Alpha）的预测。 + * 根据`src`，可为 `[B, C, H, W]` 或 `[B, T, C, H, W]` 的输出。 + * `fgr` 是 RGB 三通道，`pha` 为一通道。 + * 输出范围为 `0~1`。 +* `rec`: 循环记忆（Recurrent States）。 + * `List[Tensor, Tensor, Tensor, Tensor]` 类型。 + * 初始 `rec` 为 `List[None, None, None, None]`。 + * 有四个记忆，因为网络使用四个 `ConvGRU` 层。 + * 无论 `src` 的 Rank，所有记忆张量的 Rank 为 4。 + * 若一次给予 `T` 帧，只返回处理完最后一帧后的记忆。 + +完整的推断例子： + +```python +from torch.utils.data import DataLoader +from torchvision.transforms import ToTensor +from inference_utils import VideoReader, VideoWriter + +reader = VideoReader('input.mp4', transform=ToTensor()) +writer = VideoWriter('output.mp4', frame_rate=30) + +bgr = torch.tensor([.47, 1, .6]).view(3, 1, 1).cuda() # 绿背景 +rec = [None] * 4 # 初始记忆 + +with torch.no_grad(): + for src in DataLoader(reader): + fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio=0.25) # 将上一帧的记忆给下一帧 + writer.write(fgr * pha + bgr * (1 - pha)) +``` + +或者使用提供的视频转换 API： + +```python +from inference import convert_video + +convert_video( + model, # 模型，可以加载到任何设备（cpu 或 cuda） + input_source='input.mp4', # 视频文件，或图片序列文件夹 + input_resize=(1920, 1080), # [可选项] 缩放视频大小 + downsample_ratio=0.25, # [可选项] 下采样比，若 None，自动下采样至 512px + output_type='video', # 可选 "video"（视频）或 "png_sequence"（PNG 序列） + output_composition='com.mp4', # 若导出视频，提供文件路径。若导出 PNG 序列，提供文件夹路径 + output_alpha="pha.mp4", # [可选项] 输出透明度预测 + output_foreground="fgr.mp4", # [可选项] 输出前景预测 + output_video_mbps=4, # 若导出视频，提供视频码率 + seq_chunk=12, # 设置多帧并行计算 + num_workers=1, # 只适用于图片序列输入，读取线程 + progress=True # 显示进度条 +) +``` + +也可通过命令行调用转换 API： + +```sh +python inference.py \ + --variant mobilenetv3 \ + --checkpoint "CHECKPOINT" \ + --device cuda \ + --input-source "input.mp4" \ + --downsample-ratio 0.25 \ + --output-type video \ + --output-composition "composition.mp4" \ + --output-alpha "alpha.mp4" \ + --output-foreground "foreground.mp4" \ + --output-video-mbps 4 \ + --seq-chunk 12 +``` + +
+ +## TorchHub + +载入模型： + +```python +model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") # or "resnet50" +``` + +使用转换 API，具体请参考之前对 `convert_video` 的文档。 + +```python +convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter") + +convert_video(model, ...args...) +``` + +
+ +## TorchScript + +载入模型： + +```python +import torch +model = torch.jit.load('rvm_mobilenetv3.torchscript') +``` + +也可以可选的将模型固化（Freeze）。这会对模型进行优化，例如 BatchNorm Fusion 等。固化的模型更快。 + +```python +model = torch.jit.freeze(model) +``` + +然后，可以将 `model` 作为普通的 PyTorch 模型使用。但注意，若用固化模型调用转换 API，必须手动提供 `device` 和 `dtype`: + +```python +convert_video(frozen_model, ...args..., device='cuda', dtype=torch.float32) +``` + +
+ +## ONNX + +模型规格: +* 输入: [`src`, `r1i`, `r2i`, `r3i`, `r4i`, `downsample_ratio`]. + * `src`：输入帧，RGB 通道，形状为 `[B, C, H, W]`，范围为`0~1`。 + * `rXi`：记忆输入，初始值是是形状为 `[1, 1, 1, 1]` 的零张量。 + * `downsample_ratio` 下采样比，张量形状为 `[1]`。 + * 只有 `downsample_ratio` 必须是 `FP32`，其他输入必须和加载的模型使用一样的 `dtype`。 +* 输出: [`fgr`, `pha`, `r1o`, `r2o`, `r3o`, `r4o`] + * `fgr, pha`：前景和透明度通道输出，范围为 `0~1`。 + * `rXo`：记忆输出。 + +我们只展示用 ONNX Runtime CUDA Backend 在 Python 上的使用范例。 + +载入模型： + +```python +import onnxruntime as ort + +sess = ort.InferenceSession('rvm_mobilenetv3_fp16.onnx') +``` + +简单推断循环，但此方法不是最优化的： + +```python +import numpy as np + +rec = [ np.zeros([1, 1, 1, 1], dtype=np.float16) ] * 4 # 必须用模型一样的 dtype +downsample_ratio = np.array([0.25], dtype=np.float32) # 必须是 FP32 + +for src in YOUR_VIDEO: # src 张量是 [B, C, H, W] 形状，必须用模型一样的 dtype + fgr, pha, *rec = sess.run([], { + 'src': src, + 'r1i': rec[0], + 'r2i': rec[1], + 'r3i': rec[2], + 'r4i': rec[3], + 'downsample_ratio': downsample_ratio + }) +``` + +若使用 GPU，上例会将记忆输出传回到 CPU，再在下一帧时传回到 GPU。这种传输是无意义的，因为记忆值可以留在 GPU 上。下例使用 `iobinding` 来杜绝无用的传输。 + +```python +import onnxruntime as ort +import numpy as np + +# 载入模型 +sess = ort.InferenceSession('rvm_mobilenetv3_fp16.onnx') + +# 创建 io binding. +io = sess.io_binding() + +# 在 CUDA 上创建张量 +rec = [ ort.OrtValue.ortvalue_from_numpy(np.zeros([1, 1, 1, 1], dtype=np.float16), 'cuda') ] * 4 +downsample_ratio = ort.OrtValue.ortvalue_from_numpy(np.asarray([0.25], dtype=np.float32), 'cuda') + +# 设置输出项 +for name in ['fgr', 'pha', 'r1o', 'r2o', 'r3o', 'r4o']: + io.bind_output(name, 'cuda') + +# 推断 +for src in YOUR_VIDEO: + io.bind_cpu_input('src', src) + io.bind_ortvalue_input('r1i', rec[0]) + io.bind_ortvalue_input('r2i', rec[1]) + io.bind_ortvalue_input('r3i', rec[2]) + io.bind_ortvalue_input('r4i', rec[3]) + io.bind_ortvalue_input('downsample_ratio', downsample_ratio) + + sess.run_with_iobinding(io) + + fgr, pha, *rec = io.get_outputs() + + # 只将 `fgr` 和 `pha` 回传到 CPU + fgr = fgr.numpy() + pha = pha.numpy() +``` + +注：若你使用其他推断框架，可能有些 ONNX ops 不被支持，需被替换。可以参考 [onnx](https://github.com/PeterL1n/RobustVideoMatting/tree/onnx) 分支的代码做自行导出。 + +
+ +### TensorFlow + +范例: + +```python +import tensorflow as tf + +model = tf.keras.models.load_model('rvm_mobilenetv3_tf') +model = tf.function(model) + +rec = [ tf.constant(0.) ] * 4 # 初始记忆 +downsample_ratio = tf.constant(0.25) # 下采样率，根据视频调整 + +for src in YOUR_VIDEO: # src 张量是 [B, H, W, C] 的形状，而不是 [B, C, H, W]! + out = model([src, *rec, downsample_ratio]) + fgr, pha, *rec = out['fgr'], out['pha'], out['r1o'], out['r2o'], out['r3o'], out['r4o'] +``` + +注意，在 TensorFlow 上，所有张量都是 Channal Last 的格式。 + +我们提供 TensorFlow 的原始模型代码，请参考 [tensorflow](https://github.com/PeterL1n/RobustVideoMatting/tree/tensorflow) 分支。您可自行将 PyTorch 的权值转到 TensorFlow 模型上。 + + +
+ +### TensorFlow.js + +我们在 [tfjs](https://github.com/PeterL1n/RobustVideoMatting/tree/tfjs) 分支提供范例代码。代码简单易懂，解释如何正确使用模型。 + +
+ +### CoreML + +我们只展示在 Python 下通过 `coremltools` 使用 CoreML 模型。在部署时，同样逻辑可用于 Swift。模型的循环记忆输入不需要在处理第一帧时提供。CoreML 内部会自动创建零张量作为初始记忆。 + +```python +import coremltools as ct + +model = ct.models.model.MLModel('rvm_mobilenetv3_1920x1080_s0.25_int8.mlmodel') + +r1, r2, r3, r4 = None, None, None, None + +for src in YOUR_VIDEO: # src 是 PIL.Image. + + if r1 is None: + # 初始帧, 不用提供循环记忆 + inputs = {'src': src} + else: + # 剩余帧，提供循环记忆 + inputs = {'src': src, 'r1i': r1, 'r2i': r2, 'r3i': r3, 'r4i': r4} + + outputs = model.predict(inputs) + + fgr = outputs['fgr'] # PIL.Image + pha = outputs['pha'] # PIL.Image + + r1 = outputs['r1o'] # Numpy array + r2 = outputs['r2o'] # Numpy array + r3 = outputs['r3o'] # Numpy array + r4 = outputs['r4o'] # Numpy array + +``` + +我们的 CoreML 模型只支持固定分辨率。如果你需要其他分辨率，可自行导出。导出代码见 [coreml](https://github.com/PeterL1n/RobustVideoMatting/tree/coreml) 分支。 \ No newline at end of file diff --git a/documentation/misc/aim_test.txt b/documentation/misc/aim_test.txt new file mode 100644 index 0000000..088006b --- /dev/null +++ b/documentation/misc/aim_test.txt @@ -0,0 +1,11 @@ +boy-1518482_1920.png +girl-1219339_1920.png +girl-1467820_1280.png +girl-beautiful-young-face-53000.png +long-1245787_1920.png +model-600238_1920.png +pexels-photo-58463.png +sea-sunny-person-beach.png +wedding-dresses-1486260_1280.png +woman-952506_1920 (1).png +woman-morning-bathrobe-bathroom.png diff --git a/documentation/misc/d646_test.txt b/documentation/misc/d646_test.txt new file mode 100644 index 0000000..b5f02d5 --- /dev/null +++ b/documentation/misc/d646_test.txt @@ -0,0 +1,11 @@ +test_13.png +test_16.png +test_18.png +test_22.png +test_32.png +test_35.png +test_39.png +test_42.png +test_46.png +test_4.png +test_6.png diff --git a/documentation/misc/dvm_background_test_clips.txt b/documentation/misc/dvm_background_test_clips.txt new file mode 100644 index 0000000..a821c86 --- /dev/null +++ b/documentation/misc/dvm_background_test_clips.txt @@ -0,0 +1,162 @@ +0000 +0001 +0002 +0004 +0005 +0007 +0008 +0009 +0010 +0012 +0013 +0014 +0015 +0016 +0017 +0018 +0019 +0021 +0022 +0023 +0024 +0025 +0027 +0029 +0030 +0032 +0033 +0034 +0035 +0037 +0038 +0039 +0040 +0041 +0042 +0043 +0045 +0046 +0047 +0048 +0050 +0051 +0052 +0054 +0055 +0057 +0058 +0059 +0060 +0061 +0062 +0063 +0064 +0065 +0066 +0068 +0070 +0071 +0073 +0074 +0075 +0077 +0078 +0079 +0080 +0081 +0082 +0083 +0084 +0085 +0086 +0089 +0097 +0100 +0101 +0102 +0103 +0104 +0106 +0107 +0109 +0110 +0111 +0113 +0115 +0116 +0117 +0119 +0120 +0121 +0122 +0123 +0124 +0125 +0126 +0127 +0128 +0129 +0130 +0131 +0132 +0133 +0134 +0135 +0136 +0137 +0143 +0145 +0147 +0148 +0150 +0159 +0160 +0161 +0162 +0165 +0166 +0168 +0172 +0174 +0175 +0176 +0178 +0181 +0182 +0183 +0184 +0185 +0187 +0194 +0198 +0200 +0201 +0207 +0210 +0211 +0212 +0215 +0217 +0218 +0219 +0220 +0222 +0223 +0224 +0225 +0226 +0227 +0229 +0230 +0231 +0232 +0233 +0234 +0235 +0237 +0240 +0241 +0242 +0243 +0244 +0245 diff --git a/documentation/misc/dvm_background_train_clips.txt b/documentation/misc/dvm_background_train_clips.txt new file mode 100644 index 0000000..6bdfffe --- /dev/null +++ b/documentation/misc/dvm_background_train_clips.txt @@ -0,0 +1,3117 @@ +0000 +0002 +0003 +0004 +0005 +0006 +0007 +0009 +0010 +0012 +0013 +0014 +0015 +0016 +0019 +0021 +0022 +0023 +0024 +0025 +0028 +0029 +0030 +0031 +0032 +0034 +0035 +0036 +0037 +0039 +0040 +0041 +0042 +0043 +0044 +0046 +0047 +0048 +0049 +0050 +0051 +0052 +0053 +0054 +0055 +0056 +0057 +0058 +0060 +0061 +0062 +0063 +0064 +0065 +0066 +0067 +0068 +0069 +0070 +0071 +0073 +0074 +0075 +0076 +0077 +0078 +0079 +0081 +0082 +0087 +0088 +0099 +0100 +0101 +0104 +0105 +0107 +0108 +0109 +0110 +0111 +0112 +0113 +0114 +0115 +0117 +0118 +0119 +0120 +0122 +0123 +0124 +0125 +0127 +0128 +0129 +0130 +0131 +0132 +0133 +0134 +0135 +0136 +0137 +0138 +0139 +0140 +0141 +0142 +0144 +0146 +0147 +0148 +0150 +0151 +0152 +0153 +0154 +0155 +0156 +0157 +0158 +0159 +0160 +0161 +0163 +0164 +0165 +0167 +0168 +0169 +0170 +0171 +0172 +0174 +0175 +0176 +0177 +0178 +0180 +0181 +0182 +0184 +0185 +0187 +0188 +0189 +0190 +0192 +0193 +0194 +0195 +0196 +0197 +0198 +0199 +0200 +0202 +0203 +0204 +0205 +0206 +0207 +0208 +0209 +0210 +0211 +0212 +0213 +0214 +0215 +0217 +0218 +0219 +0220 +0221 +0222 +0223 +0224 +0225 +0226 +0227 +0229 +0230 +0231 +0233 +0234 +0235 +0236 +0237 +0238 +0240 +0241 +0242 +0243 +0244 +0245 +0246 +0247 +0248 +0249 +0250 +0251 +0252 +0253 +0254 +0255 +0256 +0257 +0258 +0259 +0260 +0261 +0262 +0263 +0264 +0265 +0266 +0267 +0268 +0269 +0270 +0271 +0272 +0273 +0274 +0275 +0276 +0277 +0278 +0279 +0280 +0281 +0282 +0283 +0284 +0285 +0286 +0287 +0288 +0289 +0290 +0291 +0292 +0293 +0294 +0297 +0298 +0299 +0300 +0301 +0302 +0303 +0304 +0305 +0306 +0308 +0309 +0310 +0311 +0312 +0313 +0314 +0315 +0316 +0317 +0319 +0320 +0321 +0322 +0323 +0324 +0325 +0326 +0327 +0328 +0329 +0330 +0331 +0332 +0333 +0335 +0336 +0337 +0338 +0339 +0341 +0342 +0344 +0345 +0346 +0348 +0349 +0352 +0353 +0356 +0357 +0358 +0359 +0360 +0361 +0362 +0363 +0364 +0365 +0366 +0368 +0369 +0370 +0371 +0372 +0373 +0374 +0375 +0376 +0377 +0378 +0379 +0380 +0381 +0382 +0383 +0384 +0385 +0386 +0387 +0388 +0389 +0391 +0392 +0393 +0394 +0395 +0397 +0398 +0399 +0400 +0401 +0402 +0403 +0404 +0405 +0406 +0407 +0408 +0409 +0410 +0411 +0413 +0414 +0415 +0416 +0417 +0419 +0420 +0421 +0422 +0423 +0424 +0425 +0426 +0427 +0428 +0429 +0431 +0433 +0434 +0435 +0436 +0437 +0438 +0439 +0440 +0441 +0442 +0443 +0445 +0446 +0447 +0448 +0449 +0450 +0451 +0452 +0453 +0454 +0456 +0457 +0458 +0459 +0462 +0463 +0464 +0465 +0466 +0467 +0468 +0469 +0470 +0471 +0472 +0473 +0474 +0475 +0476 +0477 +0478 +0479 +0480 +0481 +0482 +0483 +0484 +0485 +0486 +0487 +0488 +0489 +0490 +0491 +0492 +0493 +0494 +0499 +0501 +0502 +0503 +0504 +0505 +0506 +0507 +0509 +0510 +0511 +0512 +0513 +0514 +0515 +0517 +0518 +0519 +0520 +0521 +0522 +0524 +0526 +0527 +0529 +0530 +0534 +0535 +0536 +0538 +0539 +0541 +0542 +0543 +0544 +0545 +0546 +0548 +0549 +0550 +0552 +0554 +0555 +0556 +0557 +0558 +0559 +0560 +0561 +0562 +0563 +0564 +0565 +0566 +0567 +0568 +0571 +0572 +0573 +0574 +0575 +0576 +0577 +0578 +0579 +0580 +0581 +0582 +0583 +0584 +0586 +0587 +0589 +0590 +0591 +0592 +0594 +0595 +0596 +0597 +0598 +0600 +0601 +0602 +0603 +0604 +0605 +0606 +0608 +0609 +0610 +0611 +0612 +0613 +0614 +0615 +0616 +0617 +0618 +0619 +0620 +0624 +0625 +0626 +0627 +0628 +0629 +0630 +0631 +0634 +0635 +0636 +0637 +0638 +0639 +0640 +0641 +0642 +0643 +0644 +0645 +0646 +0647 +0648 +0650 +0651 +0652 +0654 +0655 +0656 +0658 +0659 +0660 +0661 +0662 +0663 +0664 +0665 +0666 +0667 +0669 +0670 +0671 +0672 +0673 +0674 +0675 +0676 +0677 +0678 +0679 +0680 +0681 +0682 +0683 +0684 +0685 +0686 +0687 +0689 +0690 +0691 +0692 +0693 +0694 +0695 +0696 +0697 +0698 +0699 +0700 +0701 +0702 +0703 +0704 +0705 +0706 +0707 +0708 +0709 +0710 +0711 +0712 +0713 +0714 +0715 +0716 +0717 +0718 +0719 +0720 +0721 +0723 +0724 +0725 +0726 +0727 +0729 +0730 +0731 +0732 +0733 +0734 +0735 +0736 +0738 +0740 +0741 +0742 +0743 +0744 +0746 +0747 +0748 +0749 +0750 +0752 +0753 +0754 +0755 +0756 +0757 +0758 +0759 +0760 +0762 +0763 +0764 +0765 +0766 +0767 +0768 +0770 +0771 +0772 +0773 +0774 +0775 +0776 +0777 +0778 +0779 +0780 +0781 +0782 +0783 +0784 +0786 +0787 +0788 +0789 +0790 +0791 +0792 +0793 +0794 +0795 +0796 +0797 +0798 +0800 +0801 +0804 +0806 +0808 +0809 +0811 +0812 +0813 +0814 +0815 +0816 +0817 +0819 +0823 +0824 +0825 +0827 +0828 +0829 +0830 +0831 +0832 +0833 +0834 +0835 +0836 +0837 +0840 +0841 +0842 +0847 +0848 +0850 +0851 +0852 +0853 +0854 +0855 +0856 +0857 +0858 +0859 +0860 +0861 +0862 +0864 +0867 +0868 +0869 +0870 +0871 +0872 +0873 +0874 +0876 +0877 +0878 +0879 +0880 +0881 +0882 +0883 +0885 +0886 +0887 +0889 +0890 +0891 +0892 +0893 +0894 +0895 +0896 +0899 +0900 +0901 +0902 +0903 +0904 +0905 +0906 +0907 +0908 +0909 +0910 +0911 +0912 +0913 +0914 +0915 +0916 +0917 +0918 +0919 +0921 +0922 +0923 +0924 +0925 +0926 +0927 +0929 +0930 +0931 +0932 +0933 +0934 +0935 +0936 +0937 +0939 +0940 +0941 +0942 +0945 +0946 +0947 +0948 +0949 +0950 +0951 +0952 +0953 +0954 +0955 +0956 +0957 +0958 +0960 +0962 +0963 +0964 +0965 +0966 +0967 +0968 +0969 +0970 +0971 +0973 +0974 +0976 +0977 +0978 +0979 +0980 +0981 +0982 +0983 +0984 +0985 +0986 +0987 +0988 +0989 +0990 +0991 +0992 +0993 +0994 +0995 +0996 +0997 +0998 +0999 +1000 +1001 +1002 +1003 +1005 +1006 +1008 +1009 +1010 +1011 +1012 +1013 +1014 +1016 +1017 +1018 +1019 +1020 +1021 +1022 +1023 +1024 +1025 +1026 +1028 +1029 +1030 +1032 +1035 +1036 +1037 +1038 +1039 +1040 +1041 +1042 +1043 +1044 +1045 +1046 +1047 +1048 +1049 +1050 +1052 +1053 +1054 +1055 +1056 +1057 +1058 +1061 +1062 +1063 +1064 +1065 +1066 +1067 +1068 +1069 +1072 +1075 +1076 +1077 +1078 +1079 +1081 +1082 +1083 +1084 +1087 +1088 +1089 +1090 +1096 +1097 +1098 +1099 +1101 +1102 +1103 +1104 +1105 +1106 +1107 +1108 +1109 +1110 +1111 +1112 +1113 +1115 +1116 +1117 +1118 +1119 +1120 +1121 +1122 +1123 +1124 +1128 +1129 +1130 +1131 +1132 +1134 +1137 +1138 +1139 +1140 +1141 +1142 +1143 +1144 +1145 +1146 +1147 +1148 +1149 +1150 +1151 +1152 +1153 +1154 +1155 +1156 +1157 +1158 +1159 +1160 +1161 +1162 +1163 +1165 +1168 +1169 +1170 +1171 +1172 +1173 +1174 +1175 +1176 +1177 +1178 +1179 +1180 +1181 +1182 +1183 +1184 +1185 +1186 +1187 +1188 +1189 +1190 +1191 +1192 +1193 +1194 +1195 +1196 +1197 +1199 +1200 +1201 +1202 +1203 +1204 +1206 +1207 +1208 +1211 +1212 +1213 +1215 +1216 +1217 +1219 +1220 +1221 +1222 +1223 +1224 +1225 +1226 +1227 +1228 +1229 +1230 +1231 +1232 +1233 +1234 +1235 +1237 +1238 +1239 +1240 +1241 +1242 +1245 +1246 +1248 +1249 +1252 +1253 +1255 +1256 +1257 +1258 +1259 +1260 +1261 +1262 +1263 +1264 +1265 +1266 +1267 +1268 +1269 +1270 +1271 +1272 +1273 +1274 +1275 +1277 +1278 +1279 +1280 +1281 +1282 +1283 +1284 +1287 +1288 +1289 +1290 +1291 +1293 +1294 +1295 +1296 +1297 +1299 +1300 +1301 +1302 +1303 +1304 +1305 +1306 +1307 +1308 +1309 +1310 +1311 +1312 +1313 +1316 +1317 +1318 +1319 +1320 +1321 +1322 +1323 +1324 +1325 +1326 +1327 +1328 +1329 +1331 +1332 +1333 +1334 +1335 +1336 +1337 +1338 +1339 +1340 +1341 +1342 +1343 +1344 +1345 +1346 +1347 +1348 +1350 +1351 +1352 +1353 +1354 +1355 +1356 +1357 +1359 +1360 +1361 +1362 +1363 +1364 +1365 +1366 +1367 +1368 +1369 +1370 +1371 +1372 +1373 +1374 +1375 +1376 +1378 +1379 +1380 +1381 +1382 +1383 +1385 +1386 +1387 +1389 +1390 +1391 +1392 +1393 +1394 +1395 +1396 +1397 +1398 +1399 +1400 +1402 +1403 +1405 +1406 +1409 +1410 +1411 +1412 +1413 +1414 +1415 +1416 +1417 +1418 +1420 +1421 +1422 +1423 +1424 +1425 +1426 +1427 +1428 +1429 +1430 +1431 +1432 +1433 +1434 +1435 +1436 +1437 +1438 +1439 +1441 +1442 +1443 +1444 +1445 +1446 +1447 +1448 +1449 +1450 +1451 +1452 +1453 +1454 +1455 +1456 +1457 +1458 +1459 +1460 +1461 +1462 +1465 +1466 +1467 +1468 +1469 +1470 +1471 +1472 +1473 +1474 +1475 +1476 +1477 +1478 +1479 +1480 +1481 +1482 +1483 +1484 +1485 +1486 +1487 +1488 +1489 +1490 +1491 +1493 +1494 +1495 +1496 +1497 +1499 +1500 +1501 +1502 +1503 +1504 +1505 +1506 +1507 +1508 +1509 +1510 +1511 +1512 +1513 +1514 +1515 +1516 +1519 +1520 +1521 +1522 +1523 +1524 +1525 +1526 +1527 +1528 +1529 +1530 +1531 +1532 +1534 +1535 +1536 +1537 +1538 +1539 +1540 +1541 +1542 +1543 +1544 +1545 +1546 +1547 +1548 +1549 +1550 +1551 +1552 +1553 +1554 +1555 +1556 +1557 +1558 +1560 +1561 +1562 +1563 +1565 +1566 +1567 +1568 +1569 +1570 +1571 +1572 +1573 +1574 +1575 +1576 +1577 +1578 +1579 +1580 +1581 +1582 +1583 +1584 +1585 +1586 +1587 +1588 +1589 +1590 +1591 +1592 +1593 +1594 +1595 +1596 +1597 +1598 +1599 +1600 +1601 +1602 +1603 +1604 +1605 +1606 +1608 +1609 +1611 +1612 +1613 +1614 +1615 +1616 +1617 +1618 +1619 +1620 +1621 +1622 +1623 +1624 +1625 +1626 +1627 +1628 +1629 +1630 +1631 +1632 +1633 +1634 +1635 +1636 +1637 +1638 +1639 +1640 +1641 +1642 +1643 +1644 +1645 +1646 +1649 +1650 +1651 +1652 +1653 +1654 +1655 +1656 +1657 +1658 +1659 +1660 +1661 +1662 +1663 +1664 +1665 +1666 +1667 +1668 +1669 +1670 +1671 +1672 +1673 +1675 +1676 +1677 +1678 +1679 +1680 +1681 +1682 +1683 +1684 +1685 +1686 +1687 +1688 +1689 +1690 +1691 +1692 +1693 +1694 +1695 +1696 +1697 +1698 +1699 +1700 +1701 +1702 +1703 +1704 +1705 +1706 +1707 +1708 +1709 +1710 +1712 +1713 +1714 +1715 +1716 +1717 +1718 +1719 +1720 +1721 +1722 +1723 +1724 +1725 +1726 +1727 +1728 +1729 +1730 +1731 +1732 +1733 +1734 +1735 +1736 +1737 +1738 +1739 +1740 +1741 +1742 +1743 +1744 +1745 +1746 +1747 +1748 +1749 +1750 +1751 +1752 +1753 +1754 +1755 +1756 +1757 +1758 +1759 +1760 +1761 +1762 +1763 +1764 +1765 +1766 +1767 +1768 +1769 +1770 +1771 +1772 +1773 +1774 +1775 +1776 +1778 +1779 +1780 +1781 +1782 +1784 +1785 +1786 +1787 +1788 +1789 +1790 +1791 +1792 +1793 +1794 +1795 +1796 +1797 +1798 +1799 +1800 +1804 +1806 +1807 +1808 +1809 +1811 +1812 +1813 +1814 +1815 +1816 +1817 +1818 +1819 +1820 +1821 +1822 +1823 +1825 +1826 +1827 +1828 +1829 +1831 +1833 +1834 +1835 +1836 +1837 +1838 +1839 +1840 +1842 +1843 +1844 +1845 +1846 +1847 +1848 +1849 +1850 +1851 +1852 +1853 +1854 +1855 +1856 +1857 +1858 +1859 +1861 +1862 +1863 +1864 +1865 +1866 +1867 +1868 +1869 +1870 +1871 +1872 +1873 +1874 +1875 +1876 +1877 +1879 +1880 +1881 +1882 +1886 +1887 +1889 +1891 +1892 +1893 +1894 +1896 +1897 +1898 +1899 +1900 +1901 +1902 +1903 +1904 +1905 +1906 +1907 +1908 +1909 +1910 +1911 +1912 +1913 +1914 +1915 +1916 +1917 +1918 +1919 +1920 +1921 +1922 +1923 +1924 +1925 +1926 +1927 +1928 +1929 +1930 +1931 +1932 +1933 +1934 +1935 +1936 +1937 +1938 +1939 +1940 +1941 +1942 +1943 +1944 +1945 +1946 +1947 +1948 +1949 +1950 +1952 +1953 +1954 +1955 +1956 +1958 +1959 +1960 +1961 +1962 +1963 +1964 +1965 +1966 +1967 +1968 +1969 +1970 +1971 +1972 +1973 +1974 +1975 +1976 +1977 +1978 +1980 +1981 +1982 +1983 +1984 +1985 +1986 +1988 +1989 +1990 +1991 +1992 +1993 +1994 +1995 +1997 +1998 +1999 +2000 +2002 +2003 +2004 +2005 +2007 +2008 +2010 +2011 +2012 +2013 +2014 +2015 +2016 +2017 +2018 +2019 +2020 +2021 +2022 +2024 +2025 +2026 +2027 +2028 +2029 +2030 +2031 +2032 +2035 +2036 +2037 +2038 +2039 +2040 +2041 +2042 +2043 +2045 +2046 +2047 +2049 +2050 +2052 +2053 +2054 +2056 +2057 +2059 +2060 +2061 +2064 +2066 +2068 +2069 +2070 +2071 +2072 +2073 +2074 +2075 +2077 +2078 +2079 +2080 +2081 +2082 +2083 +2084 +2085 +2086 +2087 +2088 +2089 +2090 +2091 +2092 +2093 +2094 +2096 +2097 +2098 +2099 +2100 +2101 +2102 +2103 +2104 +2105 +2107 +2108 +2109 +2111 +2112 +2113 +2114 +2115 +2116 +2117 +2119 +2120 +2121 +2122 +2123 +2124 +2125 +2126 +2127 +2128 +2129 +2130 +2131 +2132 +2133 +2134 +2135 +2136 +2137 +2138 +2139 +2140 +2141 +2143 +2144 +2145 +2146 +2147 +2148 +2149 +2152 +2153 +2155 +2158 +2159 +2160 +2161 +2162 +2163 +2164 +2165 +2166 +2167 +2168 +2169 +2170 +2171 +2174 +2176 +2177 +2178 +2179 +2180 +2181 +2182 +2183 +2184 +2186 +2187 +2188 +2189 +2190 +2191 +2192 +2193 +2195 +2197 +2198 +2199 +2200 +2201 +2202 +2203 +2204 +2205 +2206 +2207 +2208 +2209 +2210 +2211 +2212 +2213 +2214 +2215 +2216 +2217 +2218 +2219 +2220 +2221 +2222 +2223 +2225 +2226 +2227 +2228 +2229 +2230 +2231 +2232 +2233 +2234 +2235 +2236 +2238 +2239 +2240 +2241 +2242 +2243 +2244 +2245 +2246 +2247 +2248 +2249 +2250 +2251 +2252 +2253 +2255 +2256 +2257 +2258 +2259 +2260 +2261 +2263 +2264 +2265 +2267 +2268 +2269 +2270 +2271 +2272 +2273 +2274 +2275 +2276 +2277 +2278 +2279 +2280 +2281 +2282 +2283 +2284 +2285 +2286 +2287 +2288 +2290 +2291 +2292 +2293 +2294 +2295 +2297 +2299 +2300 +2301 +2302 +2303 +2304 +2305 +2311 +2312 +2313 +2314 +2315 +2316 +2317 +2318 +2319 +2320 +2322 +2324 +2325 +2326 +2329 +2331 +2332 +2334 +2335 +2336 +2337 +2338 +2339 +2340 +2341 +2342 +2343 +2344 +2345 +2347 +2349 +2350 +2351 +2352 +2353 +2355 +2356 +2358 +2359 +2360 +2361 +2362 +2363 +2364 +2365 +2367 +2368 +2369 +2370 +2372 +2373 +2374 +2375 +2376 +2377 +2378 +2379 +2380 +2381 +2382 +2383 +2384 +2386 +2389 +2390 +2391 +2392 +2393 +2394 +2395 +2396 +2397 +2398 +2399 +2400 +2401 +2402 +2403 +2404 +2406 +2407 +2408 +2409 +2410 +2411 +2412 +2413 +2414 +2416 +2417 +2418 +2419 +2420 +2421 +2422 +2423 +2424 +2425 +2426 +2427 +2428 +2429 +2431 +2432 +2433 +2434 +2437 +2439 +2440 +2441 +2442 +2443 +2444 +2445 +2446 +2448 +2449 +2450 +2451 +2452 +2453 +2454 +2455 +2456 +2457 +2460 +2461 +2462 +2463 +2464 +2465 +2466 +2467 +2468 +2469 +2470 +2472 +2474 +2475 +2479 +2480 +2481 +2482 +2483 +2484 +2485 +2486 +2487 +2488 +2489 +2490 +2491 +2492 +2493 +2494 +2495 +2496 +2497 +2499 +2500 +2501 +2502 +2503 +2504 +2505 +2506 +2507 +2508 +2509 +2510 +2511 +2512 +2513 +2514 +2515 +2516 +2517 +2518 +2519 +2520 +2521 +2522 +2523 +2524 +2525 +2530 +2534 +2536 +2537 +2539 +2540 +2541 +2542 +2543 +2544 +2545 +2546 +2548 +2549 +2550 +2551 +2553 +2554 +2555 +2556 +2559 +2560 +2562 +2567 +2570 +2573 +2575 +2576 +2578 +2579 +2582 +2585 +2588 +2590 +2593 +2594 +2595 +2596 +2597 +2598 +2600 +2601 +2602 +2603 +2604 +2605 +2606 +2607 +2612 +2613 +2615 +2616 +2617 +2618 +2619 +2620 +2622 +2623 +2624 +2625 +2629 +2630 +2633 +2634 +2635 +2637 +2638 +2639 +2641 +2643 +2644 +2646 +2648 +2649 +2650 +2652 +2654 +2656 +2659 +2661 +2662 +2663 +2664 +2665 +2667 +2669 +2670 +2671 +2672 +2674 +2675 +2677 +2679 +2680 +2682 +2684 +2685 +2687 +2689 +2691 +2692 +2693 +2694 +2695 +2696 +2697 +2698 +2699 +2700 +2701 +2702 +2703 +2705 +2706 +2707 +2710 +2713 +2714 +2715 +2717 +2718 +2720 +2721 +2722 +2726 +2727 +2729 +2733 +2734 +2737 +2739 +2740 +2741 +2742 +2743 +2745 +2747 +2748 +2749 +2755 +2756 +2757 +2758 +2759 +2761 +2762 +2763 +2765 +2767 +2768 +2769 +2770 +2773 +2775 +2776 +2780 +2781 +2782 +2785 +2787 +2790 +2791 +2793 +2794 +2795 +2796 +2797 +2798 +2800 +2801 +2802 +2803 +2806 +2808 +2810 +2812 +2813 +2814 +2815 +2816 +2818 +2819 +2820 +2821 +2822 +2823 +2825 +2827 +2828 +2829 +2830 +2831 +2832 +2833 +2834 +2835 +2837 +2838 +2842 +2843 +2844 +2845 +2847 +2848 +2849 +2850 +2851 +2852 +2853 +2854 +2856 +2857 +2858 +2860 +2861 +2862 +2863 +2864 +2865 +2866 +2868 +2869 +2871 +2874 +2875 +2877 +2879 +2881 +2882 +2884 +2885 +2887 +2888 +2889 +2890 +2891 +2893 +2894 +2895 +2896 +2900 +2902 +2903 +2904 +2905 +2906 +2907 +2908 +2911 +2912 +2914 +2915 +2916 +2917 +2918 +2919 +2921 +2923 +2924 +2925 +2926 +2928 +2929 +2930 +2931 +2932 +2933 +2934 +2935 +2936 +2938 +2939 +2940 +2941 +2944 +2945 +2946 +2947 +2952 +2954 +2957 +2958 +2960 +2962 +2963 +2966 +2967 +2968 +2969 +2970 +2971 +2972 +2973 +2974 +2975 +2976 +2981 +2982 +2984 +2985 +2986 +2988 +2993 +2994 +2996 +2998 +2999 +3000 +3001 +3002 +3004 +3006 +3007 +3008 +3010 +3015 +3016 +3017 +3018 +3019 +3020 +3021 +3022 +3024 +3026 +3027 +3029 +3033 +3034 +3035 +3036 +3037 +3040 +3041 +3042 +3043 +3044 +3045 +3046 +3047 +3048 +3049 +3050 +3051 +3053 +3056 +3057 +3058 +3059 +3060 +3061 +3063 +3065 +3066 +3068 +3069 +3070 +3074 +3077 +3078 +3079 +3080 +3081 +3082 +3083 +3084 +3085 +3086 +3087 +3094 +3095 +3097 +3098 +3100 +3101 +3102 +3103 +3105 +3106 +3108 +3109 +3110 +3111 +3112 +3113 +3114 +3115 +3116 +3117 +3118 +3119 +3121 +3123 +3124 +3125 +3127 +3128 +3130 +3131 +3133 +3134 +3136 +3137 +3140 +3142 +3143 +3144 +3145 +3147 +3148 +3152 +3153 +3155 +3156 +3157 +3159 +3160 +3162 +3164 +3165 +3166 +3168 +3170 +3171 +3173 +3176 +3178 +3179 +3180 +3181 +3184 +3187 +3188 +3190 +3191 +3194 +3197 +3199 +3200 +3201 +3202 +3204 +3209 +3210 +3211 +3212 +3214 +3215 +3217 +3218 +3219 +3221 +3222 +3223 +3224 +3226 +3228 +3229 +3230 +3232 +3237 +3238 +3239 +3240 +3241 +3242 +3243 +3245 +3247 +3248 +3250 +3251 +3252 +3253 +3254 +3257 +3258 +3259 +3260 +3262 +3264 +3266 +3267 +3269 +3270 +3275 +3278 +3279 +3280 +3282 +3284 +3285 +3286 +3287 +3288 +3289 +3292 +3293 +3295 +3296 +3298 +3299 +3300 +3301 +3302 +3303 +3304 +3309 +3311 +3312 +3315 +3316 +3317 +3318 +3319 +3322 +3325 +3328 +3332 +3334 +3339 +3340 +3342 +3346 +3348 +3349 +3350 +3351 +3352 +3354 +3355 +3357 +3358 +3361 +3363 +3364 +3365 +3366 +3367 +3368 +3369 +3370 +3373 +3374 +3377 +3378 +3380 +3381 +3382 +3383 +3384 +3385 +3386 +3391 +3393 +3395 +3396 +3397 +3398 +3399 +3400 +3401 +3402 +3403 +3404 +3405 +3407 +3408 +3409 +3410 +3411 +3415 +3416 +3418 +3420 +3422 +3423 +3424 +3427 +3428 +3431 +3433 +3435 +3437 +3438 +3439 +3440 +3441 +3442 +3443 +3444 +3446 +3449 +3450 +3451 +3452 +3453 +3454 +3455 +3456 +3457 +3460 +3462 +3463 +3464 +3465 +3466 +3467 +3468 +3470 +3475 +3476 +3477 +3483 +3484 +3486 +3487 +3489 +3492 +3496 +3497 +3498 +3500 +3501 +3502 +3505 +3507 +3508 +3509 +3510 +3511 +3512 +3513 +3514 +3515 +3517 +3518 +3519 +3521 +3524 +3525 +3526 +3528 +3529 +3532 +3536 +3541 +3542 +3543 +3544 +3545 +3546 +3549 +3550 +3551 +3552 +3554 +3556 +3557 +3558 +3559 +3560 +3562 +3563 +3564 +3566 +3567 +3568 +3571 +3572 +3573 +3574 +3575 +3576 +3578 +3579 +3580 +3582 +3584 +3585 +3588 +3590 +3592 +3593 +3594 +3599 +3602 +3605 +3606 +3608 +3611 +3612 +3615 +3617 +3620 +3621 +3622 +3623 +3624 +3625 +3626 +3629 +3630 +3632 +3633 +3636 +3637 +3638 +3641 +3644 +3646 +3647 +3648 +3649 +3654 +3655 +3656 +3657 +3660 +3662 +3667 +3671 +3672 +3673 +3674 +3675 +3676 +3678 +3679 +3680 +3681 +3682 +3683 +3685 +3687 +3691 +3692 +3694 +3695 +3697 +3698 +3699 +3700 +3701 +3703 +3704 +3705 +3707 +3709 +3711 +3712 +3715 +3717 +3718 +3719 +3720 +3721 +3723 +3724 +3725 +3726 +3728 +3729 +3733 +3734 +3735 +3737 +3738 +3739 +3741 +3743 +3746 +3748 +3750 +3753 +3754 +3756 +3757 +3759 +3760 +3762 +3764 +3765 +3766 +3768 +3772 +3773 +3774 +3776 +3777 +3779 +3780 +3782 +3783 +3784 +3785 +3786 +3790 +3792 +3793 +3794 +3798 +3799 +3800 +3801 +3802 +3803 +3804 +3807 +3809 +3813 +3814 +3816 +3819 +3822 +3823 +3824 +3826 +3827 +3828 +3829 +3830 +3831 +3832 +3833 +3834 +3836 +3837 +3838 +3839 +3840 +3841 +3842 +3844 +3845 +3847 +3848 +3849 +3850 +3852 +3854 +3855 +3856 +3857 +3858 +3861 +3862 +3863 +3865 +3869 +3872 +3873 +3874 +3879 +3880 +3883 +3885 +3886 +3887 +3888 +3889 +3890 +3891 +3893 +3894 +3895 +3897 +3898 +3899 +3900 +3901 +3903 +3904 +3906 +3914 +3915 +3916 +3920 +3923 +3924 +3925 +3926 +3928 +3929 +3931 +3932 +3934 +3935 +3937 +3939 +3942 +3943 +3945 +3949 +3950 +3952 +3955 +3956 +3963 +3965 +3966 +3967 +3969 +3970 +3971 +3974 +3976 +3977 +3978 +3980 +3981 +3982 +3983 +3984 +3987 +3988 +3992 +3995 +3996 +3997 +3999 diff --git a/documentation/misc/imagematte_train.txt b/documentation/misc/imagematte_train.txt new file mode 100644 index 0000000..314f673 --- /dev/null +++ b/documentation/misc/imagematte_train.txt @@ -0,0 +1,420 @@ +10743257206_18e7f44f2e_b.jpg +10845279884_d2d4c7b4d1_b.jpg +1-1252426161dfXY.jpg +1-1255621189mTnS.jpg +1-1259162624NMFK.jpg +1-1259245823Un3j.jpg +11363165393_05d7a21d76_b.jpg +131686738165901828.jpg +13564741125_753939e9ce_o.jpg +14731860273_5b40b19b51_o.jpg +16087-a-young-woman-showing-a-bitten-green-apple-pv.jpg +1609484818_b9bb12b.jpg +17620-a-beautiful-woman-in-a-bikini-pv.jpg +20672673163_20c8467827_b.jpg +3262986095_2d5afe583c_b.jpg +3588101233_f91aa5e3a3.jpg +3858897226_cae5b75963_o.jpg +4889657410_2d503ca287_o.jpg +4981835627_c4e6c4ffa8_o.jpg +5025666458_576b974455_o.jpg +5149410930_3a943dc43f_b.jpg +539641011387760661.jpg +5892503248_4b882863c7_o.jpg +604673748289192179.jpg +606189768665464996.jpg +624753897218113578.jpg +657454154710122500.jpg +664308724952072193.jpg +7669262460_e4be408343_b.jpg +8244818049_dfa59a3eb8_b.jpg +8688417335_01f3bafbe5_o.jpg +9434599749_e7ccfc7812_b.jpg +Aaron_Friedman_Headshot.jpg +arrgh___r___28_by_mjranum_stock.jpg +arrgh___r___29_by_mjranum_stock.jpg +arrgh___r___30_by_mjranum_stock.jpg +a-single-person-1084191_960_720.jpg +ballerina-855652_1920.jpg +beautiful-19075_960_720.jpg +boy-454633_1920.jpg +bride-2819673_1920.jpg +bride-442894_1920.jpg +face-1223346_960_720.jpg +fashion-model-portrait.jpg +fashion-model-pose.jpg +girl-1535859_1920.jpg +Girl_in_front_of_a_green_background.jpg +goth_by_bugidifino-d4w7zms.jpg +h_0.jpg +h_100.jpg +h_101.jpg +h_102.jpg +h_103.jpg +h_104.jpg +h_105.jpg +h_106.jpg +h_107.jpg +h_108.jpg +h_109.jpg +h_10.jpg +h_111.jpg +h_112.jpg +h_113.jpg +h_114.jpg +h_115.jpg +h_116.jpg +h_117.jpg +h_118.jpg +h_119.jpg +h_11.jpg +h_120.jpg +h_121.jpg +h_122.jpg +h_123.jpg +h_124.jpg +h_125.jpg +h_126.jpg +h_127.jpg +h_128.jpg +h_129.jpg +h_12.jpg +h_130.jpg +h_131.jpg +h_132.jpg +h_133.jpg +h_134.jpg +h_135.jpg +h_136.jpg +h_137.jpg +h_138.jpg +h_139.jpg +h_13.jpg +h_140.jpg +h_141.jpg +h_142.jpg +h_143.jpg +h_144.jpg +h_145.jpg +h_146.jpg +h_147.jpg +h_148.jpg +h_149.jpg +h_14.jpg +h_151.jpg +h_152.jpg +h_153.jpg +h_154.jpg +h_155.jpg +h_156.jpg +h_157.jpg +h_158.jpg +h_159.jpg +h_15.jpg +h_160.jpg +h_161.jpg +h_162.jpg +h_163.jpg +h_164.jpg +h_165.jpg +h_166.jpg +h_167.jpg +h_168.jpg +h_169.jpg +h_170.jpg +h_171.jpg +h_172.jpg +h_173.jpg +h_174.jpg +h_175.jpg +h_176.jpg +h_177.jpg +h_178.jpg +h_179.jpg +h_17.jpg +h_180.jpg +h_181.jpg +h_182.jpg +h_183.jpg +h_184.jpg +h_185.jpg +h_186.jpg +h_187.jpg +h_188.jpg +h_189.jpg +h_18.jpg +h_190.jpg +h_191.jpg +h_192.jpg +h_193.jpg +h_194.jpg +h_195.jpg +h_196.jpg +h_197.jpg +h_198.jpg +h_199.jpg +h_19.jpg +h_1.jpg +h_200.jpg +h_201.jpg +h_202.jpg +h_204.jpg +h_205.jpg +h_206.jpg +h_207.jpg +h_208.jpg +h_209.jpg +h_20.jpg +h_210.jpg +h_211.jpg +h_212.jpg +h_213.jpg +h_214.jpg +h_215.jpg +h_216.jpg +h_217.jpg +h_218.jpg +h_219.jpg +h_21.jpg +h_220.jpg +h_221.jpg +h_222.jpg +h_223.jpg +h_224.jpg +h_225.jpg +h_226.jpg +h_227.jpg +h_228.jpg +h_229.jpg +h_22.jpg +h_230.jpg +h_231.jpg +h_232.jpg +h_233.jpg +h_234.jpg +h_235.jpg +h_236.jpg +h_237.jpg +h_238.jpg +h_239.jpg +h_23.jpg +h_240.jpg +h_241.jpg +h_242.jpg +h_243.jpg +h_244.jpg +h_245.jpg +h_247.jpg +h_248.jpg +h_249.jpg +h_24.jpg +h_250.jpg +h_251.jpg +h_252.jpg +h_253.jpg +h_254.jpg +h_255.jpg +h_256.jpg +h_257.jpg +h_258.jpg +h_259.jpg +h_25.jpg +h_260.jpg +h_261.jpg +h_262.jpg +h_263.jpg +h_264.jpg +h_265.jpg +h_266.jpg +h_268.jpg +h_269.jpg +h_26.jpg +h_270.jpg +h_271.jpg +h_272.jpg +h_273.jpg +h_274.jpg +h_276.jpg +h_277.jpg +h_278.jpg +h_279.jpg +h_27.jpg +h_280.jpg +h_281.jpg +h_282.jpg +h_283.jpg +h_284.jpg +h_285.jpg +h_286.jpg +h_287.jpg +h_288.jpg +h_289.jpg +h_28.jpg +h_290.jpg +h_291.jpg +h_292.jpg +h_293.jpg +h_294.jpg +h_295.jpg +h_296.jpg +h_297.jpg +h_298.jpg +h_299.jpg +h_29.jpg +h_300.jpg +h_301.jpg +h_302.jpg +h_303.jpg +h_304.jpg +h_305.jpg +h_307.jpg +h_308.jpg +h_309.jpg +h_30.jpg +h_310.jpg +h_311.jpg +h_312.jpg +h_313.jpg +h_314.jpg +h_315.jpg +h_316.jpg +h_317.jpg +h_318.jpg +h_319.jpg +h_31.jpg +h_320.jpg +h_321.jpg +h_322.jpg +h_323.jpg +h_324.jpg +h_325.jpg +h_326.jpg +h_327.jpg +h_329.jpg +h_32.jpg +h_33.jpg +h_34.jpg +h_35.jpg +h_36.jpg +h_37.jpg +h_38.jpg +h_39.jpg +h_3.jpg +h_40.jpg +h_41.jpg +h_42.jpg +h_43.jpg +h_44.jpg +h_45.jpg +h_46.jpg +h_47.jpg +h_48.jpg +h_49.jpg +h_4.jpg +h_50.jpg +h_51.jpg +h_52.jpg +h_53.jpg +h_54.jpg +h_55.jpg +h_56.jpg +h_57.jpg +h_58.jpg +h_59.jpg +h_5.jpg +h_60.jpg +h_61.jpg +h_62.jpg +h_63.jpg +h_65.jpg +h_67.jpg +h_68.jpg +h_69.jpg +h_6.jpg +h_70.jpg +h_71.jpg +h_72.jpg +h_73.jpg +h_74.jpg +h_75.jpg +h_76.jpg +h_77.jpg +h_78.jpg +h_79.jpg +h_7.jpg +h_80.jpg +h_81.jpg +h_82.jpg +h_83.jpg +h_84.jpg +h_85.jpg +h_86.jpg +h_87.jpg +h_88.jpg +h_89.jpg +h_8.jpg +h_90.jpg +h_91.jpg +h_92.jpg +h_93.jpg +h_94.jpg +h_95.jpg +h_96.jpg +h_97.jpg +h_98.jpg +h_99.jpg +h_9.jpg +hair-flying-142210_1920.jpg +headshotid_by_bokogreat_stock-d355xf3.jpg +lil_white_goth_grl___23_by_mjranum_stock.jpg +lil_white_goth_grl___26_by_mjranum_stock.jpg +man-388104_960_720.jpg +man_headshot.jpg +MFettes-headshot.jpg +model-429733_960_720.jpg +model-610352_960_720.jpg +model-858753_960_720.jpg +model-858755_960_720.jpg +model-873675_960_720.jpg +model-873678_960_720.jpg +model-873690_960_720.jpg +model-881425_960_720.jpg +model-881431_960_720.jpg +model-female-girl-beautiful-51969.jpg +Model_in_green_dress_3.jpg +Modern_shingle_bob_haircut.jpg +Motivate_(Fitness_model).jpg +Official_portrait_of_Barack_Obama.jpg +person-woman-eyes-face.jpg +pink-hair-855660_960_720.jpg +portrait-750774_1920.jpg +Professor_Steven_Chu_ForMemRS_headshot.jpg +sailor_flying_4_by_senshistock-d4k2wmr.jpg +skin-care-937667_960_720.jpg +sorcery___8_by_mjranum_stock.jpg +t_62.jpg +t_65.jpg +test_32.jpg +test_8.jpg +train_245.jpg +train_246.jpg +train_255.jpg +train_304.jpg +train_333.jpg +train_361.jpg +train_395.jpg +train_480.jpg +train_488.jpg +train_539.jpg +wedding-846926_1920.jpg +Wild_hair.jpg +with_wings___pose_reference_by_senshistock-d6by42n_2.jpg +with_wings___pose_reference_by_senshistock-d6by42n.jpg +woman-1138435_960_720.jpg +woman1.jpg +woman2.jpg +woman-659354_960_720.jpg +woman-804072_960_720.jpg +woman-868519_960_720.jpg +Woman_in_white_shirt_on_August_2009_02.jpg +women-878869_1920.jpg diff --git a/documentation/misc/imagematte_valid.txt b/documentation/misc/imagematte_valid.txt new file mode 100644 index 0000000..0444046 --- /dev/null +++ b/documentation/misc/imagematte_valid.txt @@ -0,0 +1,15 @@ +13564741125_753939e9ce_o.jpg +3858897226_cae5b75963_o.jpg +538724499685900405.jpg +ballerina-855652_1920.jpg +boy-454633_1920.jpg +h_110.jpg +h_150.jpg +h_16.jpg +h_246.jpg +h_267.jpg +h_275.jpg +h_306.jpg +h_328.jpg +model-610352_960_720.jpg +t_66.jpg diff --git a/documentation/misc/spd_preprocess.py b/documentation/misc/spd_preprocess.py new file mode 100644 index 0000000..b9a309a --- /dev/null +++ b/documentation/misc/spd_preprocess.py @@ -0,0 +1,45 @@ +# pip install supervisely +import supervisely_lib as sly +import numpy as np +import os +from PIL import Image +from tqdm import tqdm + +# Download dataset from
+ +## Evaluation + +We synthetically composite test samples to both image and video backgrounds. Image samples (from D646, AIM) are augmented with synthetic motion. + +We only provide the composited VideoMatte240K test set. They are used in our paper evaluation. For D646 and AIM, you need to acquire the data from their authors and composite them yourself. The composition scripts we used are saved in `/evaluation` folder as reference backup. You need to modify them based on your setup. + +* [videomatte_512x512.tar (PNG 1.8G)](https://robustvideomatting.s3.us-west-2.amazonaws.com/eval/videomatte_512x288.tar) +* [videomatte_1920x1080.tar (JPG 2.2G)](https://robustvideomatting.s3.us-west-2.amazonaws.com/eval/videomatte_1920x1080.tar) + +Evaluation scripts are provided in `/evaluation` folder. \ No newline at end of file diff --git a/evaluation/evaluate_hr.py b/evaluation/evaluate_hr.py new file mode 100644 index 0000000..cb5d146 --- /dev/null +++ b/evaluation/evaluate_hr.py @@ -0,0 +1,216 @@ +""" +HR (High-Resolution) evaluation. We found using numpy is very slow for high resolution, so we moved it to PyTorch using CUDA. + +Note, the script only does evaluation. You will need to first inference yourself and save the results to disk +Expected directory format for both prediction and ground-truth is: + + videomatte_1920x1080 + ├── videomatte_motion + ├── pha + ├── 0000 + ├── 0000.png + ├── fgr + ├── 0000 + ├── 0000.png + ├── videomatte_static + ├── pha + ├── 0000 + ├── 0000.png + ├── fgr + ├── 0000 + ├── 0000.png + +Prediction must have the exact file structure and file name as the ground-truth, +meaning that if the ground-truth is png/jpg, prediction should be png/jpg. + +Example usage: + +python evaluate.py \ + --pred-dir pred/videomatte_1920x1080 \ + --true-dir true/videomatte_1920x1080 + +An excel sheet with evaluation results will be written to "pred/videomatte_1920x1080/videomatte_1920x1080.xlsx" +""" + + +import argparse +import os +import cv2 +import kornia +import numpy as np +import xlsxwriter +import torch +from concurrent.futures import ThreadPoolExecutor +from tqdm import tqdm + + +class Evaluator: + def __init__(self): + self.parse_args() + self.init_metrics() + self.evaluate() + self.write_excel() + + def parse_args(self): + parser = argparse.ArgumentParser() + parser.add_argument('--pred-dir', type=str, required=True) + parser.add_argument('--true-dir', type=str, required=True) + parser.add_argument('--num-workers', type=int, default=48) + parser.add_argument('--metrics', type=str, nargs='+', default=[ + 'pha_mad', 'pha_mse', 'pha_grad', 'pha_dtssd', 'fgr_mse']) + self.args = parser.parse_args() + + def init_metrics(self): + self.mad = MetricMAD() + self.mse = MetricMSE() + self.grad = MetricGRAD() + self.dtssd = MetricDTSSD() + + def evaluate(self): + tasks = [] + position = 0 + + with ThreadPoolExecutor(max_workers=self.args.num_workers) as executor: + for dataset in sorted(os.listdir(self.args.pred_dir)): + if os.path.isdir(os.path.join(self.args.pred_dir, dataset)): + for clip in sorted(os.listdir(os.path.join(self.args.pred_dir, dataset))): + future = executor.submit(self.evaluate_worker, dataset, clip, position) + tasks.append((dataset, clip, future)) + position += 1 + + self.results = [(dataset, clip, future.result()) for dataset, clip, future in tasks] + + def write_excel(self): + workbook = xlsxwriter.Workbook(os.path.join(self.args.pred_dir, f'{os.path.basename(self.args.pred_dir)}.xlsx')) + summarysheet = workbook.add_worksheet('summary') + metricsheets = [workbook.add_worksheet(metric) for metric in self.results[0][2].keys()] + + for i, metric in enumerate(self.results[0][2].keys()): + summarysheet.write(i, 0, metric) + summarysheet.write(i, 1, f'={metric}!B2') + + for row, (dataset, clip, metrics) in enumerate(self.results): + for metricsheet, metric in zip(metricsheets, metrics.values()): + # Write the header + if row == 0: + metricsheet.write(1, 0, 'Average') + metricsheet.write(1, 1, f'=AVERAGE(C2:ZZ2)') + for col in range(len(metric)): + metricsheet.write(0, col + 2, col) + colname = xlsxwriter.utility.xl_col_to_name(col + 2) + metricsheet.write(1, col + 2, f'=AVERAGE({colname}3:{colname}9999)') + + metricsheet.write(row + 2, 0, dataset) + metricsheet.write(row + 2, 1, clip) + metricsheet.write_row(row + 2, 2, metric) + + workbook.close() + + def evaluate_worker(self, dataset, clip, position): + framenames = sorted(os.listdir(os.path.join(self.args.pred_dir, dataset, clip, 'pha'))) + metrics = {metric_name : [] for metric_name in self.args.metrics} + + pred_pha_tm1 = None + true_pha_tm1 = None + + for i, framename in enumerate(tqdm(framenames, desc=f'{dataset} {clip}', position=position, dynamic_ncols=True)): + true_pha = cv2.imread(os.path.join(self.args.true_dir, dataset, clip, 'pha', framename), cv2.IMREAD_GRAYSCALE) + pred_pha = cv2.imread(os.path.join(self.args.pred_dir, dataset, clip, 'pha', framename), cv2.IMREAD_GRAYSCALE) + + true_pha = torch.from_numpy(true_pha).cuda(non_blocking=True).float().div_(255) + pred_pha = torch.from_numpy(pred_pha).cuda(non_blocking=True).float().div_(255) + + if 'pha_mad' in self.args.metrics: + metrics['pha_mad'].append(self.mad(pred_pha, true_pha)) + if 'pha_mse' in self.args.metrics: + metrics['pha_mse'].append(self.mse(pred_pha, true_pha)) + if 'pha_grad' in self.args.metrics: + metrics['pha_grad'].append(self.grad(pred_pha, true_pha)) + if 'pha_conn' in self.args.metrics: + metrics['pha_conn'].append(self.conn(pred_pha, true_pha)) + if 'pha_dtssd' in self.args.metrics: + if i == 0: + metrics['pha_dtssd'].append(0) + else: + metrics['pha_dtssd'].append(self.dtssd(pred_pha, pred_pha_tm1, true_pha, true_pha_tm1)) + + pred_pha_tm1 = pred_pha + true_pha_tm1 = true_pha + + if 'fgr_mse' in self.args.metrics: + true_fgr = cv2.imread(os.path.join(self.args.true_dir, dataset, clip, 'fgr', framename), cv2.IMREAD_COLOR) + pred_fgr = cv2.imread(os.path.join(self.args.pred_dir, dataset, clip, 'fgr', framename), cv2.IMREAD_COLOR) + + true_fgr = torch.from_numpy(true_fgr).float().div_(255) + pred_fgr = torch.from_numpy(pred_fgr).float().div_(255) + + true_msk = true_pha > 0 + metrics['fgr_mse'].append(self.mse(pred_fgr[true_msk], true_fgr[true_msk])) + + return metrics + + +class MetricMAD: + def __call__(self, pred, true): + return (pred - true).abs_().mean() * 1e3 + + +class MetricMSE: + def __call__(self, pred, true): + return ((pred - true) ** 2).mean() * 1e3 + + +class MetricGRAD: + def __init__(self, sigma=1.4): + self.filter_x, self.filter_y = self.gauss_filter(sigma) + self.filter_x = torch.from_numpy(self.filter_x).unsqueeze(0).cuda() + self.filter_y = torch.from_numpy(self.filter_y).unsqueeze(0).cuda() + + def __call__(self, pred, true): + true_grad = self.gauss_gradient(true) + pred_grad = self.gauss_gradient(pred) + return ((true_grad - pred_grad) ** 2).sum() / 1000 + + def gauss_gradient(self, img): + img_filtered_x = kornia.filters.filter2D(img[None, None, :, :], self.filter_x, border_type='replicate')[0, 0] + img_filtered_y = kornia.filters.filter2D(img[None, None, :, :], self.filter_y, border_type='replicate')[0, 0] + return (img_filtered_x**2 + img_filtered_y**2).sqrt() + + @staticmethod + def gauss_filter(sigma, epsilon=1e-2): + half_size = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon))) + size = np.int(2 * half_size + 1) + + # create filter in x axis + filter_x = np.zeros((size, size)) + for i in range(size): + for j in range(size): + filter_x[i, j] = MetricGRAD.gaussian(i - half_size, sigma) * MetricGRAD.dgaussian( + j - half_size, sigma) + + # normalize filter + norm = np.sqrt((filter_x**2).sum()) + filter_x = filter_x / norm + filter_y = np.transpose(filter_x) + + return filter_x, filter_y + + @staticmethod + def gaussian(x, sigma): + return np.exp(-x**2 / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi)) + + @staticmethod + def dgaussian(x, sigma): + return -x * MetricGRAD.gaussian(x, sigma) / sigma**2 + + +class MetricDTSSD: + def __call__(self, pred_t, pred_tm1, true_t, true_tm1): + dtSSD = ((pred_t - pred_tm1) - (true_t - true_tm1)) ** 2 + dtSSD = dtSSD.sum() / true_t.numel() + dtSSD = dtSSD.sqrt() + return dtSSD * 1e2 + + +if __name__ == '__main__': + Evaluator() \ No newline at end of file diff --git a/evaluation/evaluate_lr.py b/evaluation/evaluate_lr.py new file mode 100644 index 0000000..92d671f --- /dev/null +++ b/evaluation/evaluate_lr.py @@ -0,0 +1,254 @@ +""" +LR (Low-Resolution) evaluation. + +Note, the script only does evaluation. You will need to first inference yourself and save the results to disk +Expected directory format for both prediction and ground-truth is: + + videomatte_512x288 + ├── videomatte_motion + ├── pha + ├── 0000 + ├── 0000.png + ├── fgr + ├── 0000 + ├── 0000.png + ├── videomatte_static + ├── pha + ├── 0000 + ├── 0000.png + ├── fgr + ├── 0000 + ├── 0000.png + +Prediction must have the exact file structure and file name as the ground-truth, +meaning that if the ground-truth is png/jpg, prediction should be png/jpg. + +Example usage: + +python evaluate.py \ + --pred-dir PATH_TO_PREDICTIONS/videomatte_512x288 \ + --true-dir PATH_TO_GROUNDTURTH/videomatte_512x288 + +An excel sheet with evaluation results will be written to "PATH_TO_PREDICTIONS/videomatte_512x288/videomatte_512x288.xlsx" +""" + + +import argparse +import os +import cv2 +import numpy as np +import xlsxwriter +from concurrent.futures import ThreadPoolExecutor +from tqdm import tqdm + + +class Evaluator: + def __init__(self): + self.parse_args() + self.init_metrics() + self.evaluate() + self.write_excel() + + def parse_args(self): + parser = argparse.ArgumentParser() + parser.add_argument('--pred-dir', type=str, required=True) + parser.add_argument('--true-dir', type=str, required=True) + parser.add_argument('--num-workers', type=int, default=48) + parser.add_argument('--metrics', type=str, nargs='+', default=[ + 'pha_mad', 'pha_mse', 'pha_grad', 'pha_conn', 'pha_dtssd', 'fgr_mad', 'fgr_mse']) + self.args = parser.parse_args() + + def init_metrics(self): + self.mad = MetricMAD() + self.mse = MetricMSE() + self.grad = MetricGRAD() + self.conn = MetricCONN() + self.dtssd = MetricDTSSD() + + def evaluate(self): + tasks = [] + position = 0 + + with ThreadPoolExecutor(max_workers=self.args.num_workers) as executor: + for dataset in sorted(os.listdir(self.args.pred_dir)): + if os.path.isdir(os.path.join(self.args.pred_dir, dataset)): + for clip in sorted(os.listdir(os.path.join(self.args.pred_dir, dataset))): + future = executor.submit(self.evaluate_worker, dataset, clip, position) + tasks.append((dataset, clip, future)) + position += 1 + + self.results = [(dataset, clip, future.result()) for dataset, clip, future in tasks] + + def write_excel(self): + workbook = xlsxwriter.Workbook(os.path.join(self.args.pred_dir, f'{os.path.basename(self.args.pred_dir)}.xlsx')) + summarysheet = workbook.add_worksheet('summary') + metricsheets = [workbook.add_worksheet(metric) for metric in self.results[0][2].keys()] + + for i, metric in enumerate(self.results[0][2].keys()): + summarysheet.write(i, 0, metric) + summarysheet.write(i, 1, f'={metric}!B2') + + for row, (dataset, clip, metrics) in enumerate(self.results): + for metricsheet, metric in zip(metricsheets, metrics.values()): + # Write the header + if row == 0: + metricsheet.write(1, 0, 'Average') + metricsheet.write(1, 1, f'=AVERAGE(C2:ZZ2)') + for col in range(len(metric)): + metricsheet.write(0, col + 2, col) + colname = xlsxwriter.utility.xl_col_to_name(col + 2) + metricsheet.write(1, col + 2, f'=AVERAGE({colname}3:{colname}9999)') + + metricsheet.write(row + 2, 0, dataset) + metricsheet.write(row + 2, 1, clip) + metricsheet.write_row(row + 2, 2, metric) + + workbook.close() + + def evaluate_worker(self, dataset, clip, position): + framenames = sorted(os.listdir(os.path.join(self.args.pred_dir, dataset, clip, 'pha'))) + metrics = {metric_name : [] for metric_name in self.args.metrics} + + pred_pha_tm1 = None + true_pha_tm1 = None + + for i, framename in enumerate(tqdm(framenames, desc=f'{dataset} {clip}', position=position, dynamic_ncols=True)): + true_pha = cv2.imread(os.path.join(self.args.true_dir, dataset, clip, 'pha', framename), cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255 + pred_pha = cv2.imread(os.path.join(self.args.pred_dir, dataset, clip, 'pha', framename), cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255 + if 'pha_mad' in self.args.metrics: + metrics['pha_mad'].append(self.mad(pred_pha, true_pha)) + if 'pha_mse' in self.args.metrics: + metrics['pha_mse'].append(self.mse(pred_pha, true_pha)) + if 'pha_grad' in self.args.metrics: + metrics['pha_grad'].append(self.grad(pred_pha, true_pha)) + if 'pha_conn' in self.args.metrics: + metrics['pha_conn'].append(self.conn(pred_pha, true_pha)) + if 'pha_dtssd' in self.args.metrics: + if i == 0: + metrics['pha_dtssd'].append(0) + else: + metrics['pha_dtssd'].append(self.dtssd(pred_pha, pred_pha_tm1, true_pha, true_pha_tm1)) + + pred_pha_tm1 = pred_pha + true_pha_tm1 = true_pha + + if 'fgr_mse' in self.args.metrics or 'fgr_mad' in self.args.metrics: + true_fgr = cv2.imread(os.path.join(self.args.true_dir, dataset, clip, 'fgr', framename), cv2.IMREAD_COLOR).astype(np.float32) / 255 + pred_fgr = cv2.imread(os.path.join(self.args.pred_dir, dataset, clip, 'fgr', framename), cv2.IMREAD_COLOR).astype(np.float32) / 255 + true_msk = true_pha > 0 + + if 'fgr_mse' in self.args.metrics: + metrics['fgr_mse'].append(self.mse(pred_fgr[true_msk], true_fgr[true_msk])) + if 'fgr_mad' in self.args.metrics: + metrics['fgr_mad'].append(self.mad(pred_fgr[true_msk], true_fgr[true_msk])) + + return metrics + + +class MetricMAD: + def __call__(self, pred, true): + return np.abs(pred - true).mean() * 1e3 + + +class MetricMSE: + def __call__(self, pred, true): + return ((pred - true) ** 2).mean() * 1e3 + + +class MetricGRAD: + def __init__(self, sigma=1.4): + self.filter_x, self.filter_y = self.gauss_filter(sigma) + + def __call__(self, pred, true): + pred_normed = np.zeros_like(pred) + true_normed = np.zeros_like(true) + cv2.normalize(pred, pred_normed, 1., 0., cv2.NORM_MINMAX) + cv2.normalize(true, true_normed, 1., 0., cv2.NORM_MINMAX) + + true_grad = self.gauss_gradient(true_normed).astype(np.float32) + pred_grad = self.gauss_gradient(pred_normed).astype(np.float32) + + grad_loss = ((true_grad - pred_grad) ** 2).sum() + return grad_loss / 1000 + + def gauss_gradient(self, img): + img_filtered_x = cv2.filter2D(img, -1, self.filter_x, borderType=cv2.BORDER_REPLICATE) + img_filtered_y = cv2.filter2D(img, -1, self.filter_y, borderType=cv2.BORDER_REPLICATE) + return np.sqrt(img_filtered_x**2 + img_filtered_y**2) + + @staticmethod + def gauss_filter(sigma, epsilon=1e-2): + half_size = np.ceil(sigma * np.sqrt(-2 * np.log(np.sqrt(2 * np.pi) * sigma * epsilon))) + size = np.int(2 * half_size + 1) + + # create filter in x axis + filter_x = np.zeros((size, size)) + for i in range(size): + for j in range(size): + filter_x[i, j] = MetricGRAD.gaussian(i - half_size, sigma) * MetricGRAD.dgaussian( + j - half_size, sigma) + + # normalize filter + norm = np.sqrt((filter_x**2).sum()) + filter_x = filter_x / norm + filter_y = np.transpose(filter_x) + + return filter_x, filter_y + + @staticmethod + def gaussian(x, sigma): + return np.exp(-x**2 / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi)) + + @staticmethod + def dgaussian(x, sigma): + return -x * MetricGRAD.gaussian(x, sigma) / sigma**2 + + +class MetricCONN: + def __call__(self, pred, true): + step=0.1 + thresh_steps = np.arange(0, 1 + step, step) + round_down_map = -np.ones_like(true) + for i in range(1, len(thresh_steps)): + true_thresh = true >= thresh_steps[i] + pred_thresh = pred >= thresh_steps[i] + intersection = (true_thresh & pred_thresh).astype(np.uint8) + + # connected components + _, output, stats, _ = cv2.connectedComponentsWithStats( + intersection, connectivity=4) + # start from 1 in dim 0 to exclude background + size = stats[1:, -1] + + # largest connected component of the intersection + omega = np.zeros_like(true) + if len(size) != 0: + max_id = np.argmax(size) + # plus one to include background + omega[output == max_id + 1] = 1 + + mask = (round_down_map == -1) & (omega == 0) + round_down_map[mask] = thresh_steps[i - 1] + round_down_map[round_down_map == -1] = 1 + + true_diff = true - round_down_map + pred_diff = pred - round_down_map + # only calculate difference larger than or equal to 0.15 + true_phi = 1 - true_diff * (true_diff >= 0.15) + pred_phi = 1 - pred_diff * (pred_diff >= 0.15) + + connectivity_error = np.sum(np.abs(true_phi - pred_phi)) + return connectivity_error / 1000 + + +class MetricDTSSD: + def __call__(self, pred_t, pred_tm1, true_t, true_tm1): + dtSSD = ((pred_t - pred_tm1) - (true_t - true_tm1)) ** 2 + dtSSD = np.sum(dtSSD) / true_t.size + dtSSD = np.sqrt(dtSSD) + return dtSSD * 1e2 + + + +if __name__ == '__main__': + Evaluator() \ No newline at end of file diff --git a/evaluation/generate_imagematte_with_background_image.py b/evaluation/generate_imagematte_with_background_image.py new file mode 100644 index 0000000..9d836f1 --- /dev/null +++ b/evaluation/generate_imagematte_with_background_image.py @@ -0,0 +1,146 @@ +""" +python generate_imagematte_with_background_image.py \ + --imagematte-dir ../matting-data/Distinctions/test \ + --background-dir ../matting-data/Backgrounds/valid \ + --resolution 512 \ + --out-dir ../matting-data/evaluation/distinction_static_sd/ \ + --random-seed 10 + +Seed: + 10 - distinction-static + 11 - distinction-motion + 12 - adobe-static + 13 - adobe-motion + +""" + +import argparse +import os +import pims +import numpy as np +import random +from PIL import Image +from tqdm import tqdm +from tqdm.contrib.concurrent import process_map +from torchvision import transforms +from torchvision.transforms import functional as F + +parser = argparse.ArgumentParser() +parser.add_argument('--imagematte-dir', type=str, required=True) +parser.add_argument('--background-dir', type=str, required=True) +parser.add_argument('--num-samples', type=int, default=20) +parser.add_argument('--num-frames', type=int, default=100) +parser.add_argument('--resolution', type=int, required=True) +parser.add_argument('--out-dir', type=str, required=True) +parser.add_argument('--random-seed', type=int) +parser.add_argument('--extension', type=str, default='.png') +args = parser.parse_args() + +random.seed(args.random_seed) + +imagematte_filenames = os.listdir(os.path.join(args.imagematte_dir, 'fgr')) +background_filenames = os.listdir(args.background_dir) +random.shuffle(imagematte_filenames) +random.shuffle(background_filenames) + + +def lerp(a, b, percentage): + return a * (1 - percentage) + b * percentage + +def motion_affine(*imgs): + config = dict(degrees=(-10, 10), translate=(0.1, 0.1), + scale_ranges=(0.9, 1.1), shears=(-5, 5), img_size=imgs[0][0].size) + angleA, (transXA, transYA), scaleA, (shearXA, shearYA) = transforms.RandomAffine.get_params(**config) + angleB, (transXB, transYB), scaleB, (shearXB, shearYB) = transforms.RandomAffine.get_params(**config) + + T = len(imgs[0]) + variation_over_time = random.random() + for t in range(T): + percentage = (t / (T - 1)) * variation_over_time + angle = lerp(angleA, angleB, percentage) + transX = lerp(transXA, transXB, percentage) + transY = lerp(transYA, transYB, percentage) + scale = lerp(scaleA, scaleB, percentage) + shearX = lerp(shearXA, shearXB, percentage) + shearY = lerp(shearYA, shearYB, percentage) + for img in imgs: + img[t] = F.affine(img[t], angle, (transX, transY), scale, (shearX, shearY), F.InterpolationMode.BILINEAR) + return imgs + + + +def process(i): + imagematte_filename = imagematte_filenames[i % len(imagematte_filenames)] + background_filename = background_filenames[i % len(background_filenames)] + + out_path = os.path.join(args.out_dir, str(i).zfill(4)) + os.makedirs(os.path.join(out_path, 'fgr'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'pha'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'com'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'bgr'), exist_ok=True) + + with Image.open(os.path.join(args.background_dir, background_filename)) as bgr: + bgr = bgr.convert('RGB') + + w, h = bgr.size + scale = args.resolution / min(h, w) + w, h = int(w * scale), int(h * scale) + bgr = bgr.resize((w, h)) + bgr = F.center_crop(bgr, (args.resolution, args.resolution)) + + with Image.open(os.path.join(args.imagematte_dir, 'fgr', imagematte_filename)) as fgr, \ + Image.open(os.path.join(args.imagematte_dir, 'pha', imagematte_filename)) as pha: + fgr = fgr.convert('RGB') + pha = pha.convert('L') + + fgrs = [fgr] * args.num_frames + phas = [pha] * args.num_frames + fgrs, phas = motion_affine(fgrs, phas) + + for t in tqdm(range(args.num_frames), desc=str(i).zfill(4)): + fgr = fgrs[t] + pha = phas[t] + + w, h = fgr.size + scale = args.resolution / max(h, w) + w, h = int(w * scale), int(h * scale) + + fgr = fgr.resize((w, h)) + pha = pha.resize((w, h)) + + if h < args.resolution: + pt = (args.resolution - h) // 2 + pb = args.resolution - h - pt + else: + pt = 0 + pb = 0 + + if w < args.resolution: + pl = (args.resolution - w) // 2 + pr = args.resolution - w - pl + else: + pl = 0 + pr = 0 + + fgr = F.pad(fgr, [pl, pt, pr, pb]) + pha = F.pad(pha, [pl, pt, pr, pb]) + + if i // len(imagematte_filenames) % 2 == 1: + fgr = fgr.transpose(Image.FLIP_LEFT_RIGHT) + pha = pha.transpose(Image.FLIP_LEFT_RIGHT) + + fgr.save(os.path.join(out_path, 'fgr', str(t).zfill(4) + args.extension)) + pha.save(os.path.join(out_path, 'pha', str(t).zfill(4) + args.extension)) + + if t == 0: + bgr.save(os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension)) + else: + os.symlink(str(0).zfill(4) + args.extension, os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension)) + + pha = np.asarray(pha).astype(float)[:, :, None] / 255 + com = Image.fromarray(np.uint8(np.asarray(fgr) * pha + np.asarray(bgr) * (1 - pha))) + com.save(os.path.join(out_path, 'com', str(t).zfill(4) + args.extension)) + + +if __name__ == '__main__': + r = process_map(process, range(args.num_samples), max_workers=32) \ No newline at end of file diff --git a/evaluation/generate_imagematte_with_background_video.py b/evaluation/generate_imagematte_with_background_video.py new file mode 100644 index 0000000..6c67c7b --- /dev/null +++ b/evaluation/generate_imagematte_with_background_video.py @@ -0,0 +1,189 @@ +""" +python generate_imagematte_with_background_video.py \ + --imagematte-dir ../matting-data/Distinctions/test \ + --background-dir ../matting-data/BackgroundVideos_mp4/test \ + --resolution 512 \ + --out-dir ../matting-data/evaluation/distinction_motion_sd/ \ + --random-seed 11 + +Seed: + 10 - distinction-static + 11 - distinction-motion + 12 - adobe-static + 13 - adobe-motion + +""" + +import argparse +import os +import pims +import numpy as np +import random +from multiprocessing import Pool +from PIL import Image +# from tqdm import tqdm +from tqdm.contrib.concurrent import process_map +from torchvision import transforms +from torchvision.transforms import functional as F + +parser = argparse.ArgumentParser() +parser.add_argument('--imagematte-dir', type=str, required=True) +parser.add_argument('--background-dir', type=str, required=True) +parser.add_argument('--num-samples', type=int, default=20) +parser.add_argument('--num-frames', type=int, default=100) +parser.add_argument('--resolution', type=int, required=True) +parser.add_argument('--out-dir', type=str, required=True) +parser.add_argument('--random-seed', type=int) +parser.add_argument('--extension', type=str, default='.png') +args = parser.parse_args() + +random.seed(args.random_seed) + +imagematte_filenames = os.listdir(os.path.join(args.imagematte_dir, 'fgr')) +random.shuffle(imagematte_filenames) + +background_filenames = [ + "0000.mp4", + "0007.mp4", + "0008.mp4", + "0010.mp4", + "0013.mp4", + "0015.mp4", + "0016.mp4", + "0018.mp4", + "0021.mp4", + "0029.mp4", + "0033.mp4", + "0035.mp4", + "0039.mp4", + "0050.mp4", + "0052.mp4", + "0055.mp4", + "0060.mp4", + "0063.mp4", + "0087.mp4", + "0086.mp4", + "0090.mp4", + "0101.mp4", + "0110.mp4", + "0117.mp4", + "0120.mp4", + "0122.mp4", + "0123.mp4", + "0125.mp4", + "0128.mp4", + "0131.mp4", + "0172.mp4", + "0176.mp4", + "0181.mp4", + "0187.mp4", + "0193.mp4", + "0198.mp4", + "0220.mp4", + "0221.mp4", + "0224.mp4", + "0229.mp4", + "0233.mp4", + "0238.mp4", + "0241.mp4", + "0245.mp4", + "0246.mp4" +] + +random.shuffle(background_filenames) + +def lerp(a, b, percentage): + return a * (1 - percentage) + b * percentage + +def motion_affine(*imgs): + config = dict(degrees=(-10, 10), translate=(0.1, 0.1), + scale_ranges=(0.9, 1.1), shears=(-5, 5), img_size=imgs[0][0].size) + angleA, (transXA, transYA), scaleA, (shearXA, shearYA) = transforms.RandomAffine.get_params(**config) + angleB, (transXB, transYB), scaleB, (shearXB, shearYB) = transforms.RandomAffine.get_params(**config) + + T = len(imgs[0]) + variation_over_time = random.random() + for t in range(T): + percentage = (t / (T - 1)) * variation_over_time + angle = lerp(angleA, angleB, percentage) + transX = lerp(transXA, transXB, percentage) + transY = lerp(transYA, transYB, percentage) + scale = lerp(scaleA, scaleB, percentage) + shearX = lerp(shearXA, shearXB, percentage) + shearY = lerp(shearYA, shearYB, percentage) + for img in imgs: + img[t] = F.affine(img[t], angle, (transX, transY), scale, (shearX, shearY), F.InterpolationMode.BILINEAR) + return imgs + + +def process(i): + imagematte_filename = imagematte_filenames[i % len(imagematte_filenames)] + background_filename = background_filenames[i % len(background_filenames)] + + bgrs = pims.PyAVVideoReader(os.path.join(args.background_dir, background_filename)) + + out_path = os.path.join(args.out_dir, str(i).zfill(4)) + os.makedirs(os.path.join(out_path, 'fgr'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'pha'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'com'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'bgr'), exist_ok=True) + + with Image.open(os.path.join(args.imagematte_dir, 'fgr', imagematte_filename)) as fgr, \ + Image.open(os.path.join(args.imagematte_dir, 'pha', imagematte_filename)) as pha: + fgr = fgr.convert('RGB') + pha = pha.convert('L') + + fgrs = [fgr] * args.num_frames + phas = [pha] * args.num_frames + fgrs, phas = motion_affine(fgrs, phas) + + for t in range(args.num_frames): + fgr = fgrs[t] + pha = phas[t] + + w, h = fgr.size + scale = args.resolution / max(h, w) + w, h = int(w * scale), int(h * scale) + + fgr = fgr.resize((w, h)) + pha = pha.resize((w, h)) + + if h < args.resolution: + pt = (args.resolution - h) // 2 + pb = args.resolution - h - pt + else: + pt = 0 + pb = 0 + + if w < args.resolution: + pl = (args.resolution - w) // 2 + pr = args.resolution - w - pl + else: + pl = 0 + pr = 0 + + fgr = F.pad(fgr, [pl, pt, pr, pb]) + pha = F.pad(pha, [pl, pt, pr, pb]) + + if i // len(imagematte_filenames) % 2 == 1: + fgr = fgr.transpose(Image.FLIP_LEFT_RIGHT) + pha = pha.transpose(Image.FLIP_LEFT_RIGHT) + + fgr.save(os.path.join(out_path, 'fgr', str(t).zfill(4) + args.extension)) + pha.save(os.path.join(out_path, 'pha', str(t).zfill(4) + args.extension)) + + bgr = Image.fromarray(bgrs[t]).convert('RGB') + w, h = bgr.size + scale = args.resolution / min(h, w) + w, h = int(w * scale), int(h * scale) + bgr = bgr.resize((w, h)) + bgr = F.center_crop(bgr, (args.resolution, args.resolution)) + bgr.save(os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension)) + + pha = np.asarray(pha).astype(float)[:, :, None] / 255 + com = Image.fromarray(np.uint8(np.asarray(fgr) * pha + np.asarray(bgr) * (1 - pha))) + com.save(os.path.join(out_path, 'com', str(t).zfill(4) + args.extension)) + +if __name__ == '__main__': + r = process_map(process, range(args.num_samples), max_workers=10) + diff --git a/evaluation/generate_videomatte_with_background_image.py b/evaluation/generate_videomatte_with_background_image.py new file mode 100644 index 0000000..1c9f349 --- /dev/null +++ b/evaluation/generate_videomatte_with_background_image.py @@ -0,0 +1,78 @@ +""" +python generate_videomatte_with_background_image.py \ + --videomatte-dir ../matting-data/VideoMatte240K_JPEG_HD/test \ + --background-dir ../matting-data/Backgrounds/valid \ + --num-samples 25 \ + --resize 512 288 \ + --out-dir ../matting-data/evaluation/vidematte_static_sd/ +""" + +import argparse +import os +import pims +import numpy as np +import random +from PIL import Image +from tqdm import tqdm + +parser = argparse.ArgumentParser() +parser.add_argument('--videomatte-dir', type=str, required=True) +parser.add_argument('--background-dir', type=str, required=True) +parser.add_argument('--num-samples', type=int, default=20) +parser.add_argument('--num-frames', type=int, default=100) +parser.add_argument('--resize', type=int, default=None, nargs=2) +parser.add_argument('--out-dir', type=str, required=True) +parser.add_argument('--extension', type=str, default='.png') +args = parser.parse_args() + +random.seed(10) + +videomatte_filenames = [(clipname, sorted(os.listdir(os.path.join(args.videomatte_dir, 'fgr', clipname)))) + for clipname in sorted(os.listdir(os.path.join(args.videomatte_dir, 'fgr')))] + +background_filenames = os.listdir(args.background_dir) +random.shuffle(background_filenames) + +for i in range(args.num_samples): + + clipname, framenames = videomatte_filenames[i % len(videomatte_filenames)] + + out_path = os.path.join(args.out_dir, str(i).zfill(4)) + os.makedirs(os.path.join(out_path, 'fgr'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'pha'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'com'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'bgr'), exist_ok=True) + + with Image.open(os.path.join(args.background_dir, background_filenames[i])) as bgr: + bgr = bgr.convert('RGB') + + + base_t = random.choice(range(len(framenames) - args.num_frames)) + + for t in tqdm(range(args.num_frames), desc=str(i).zfill(4)): + with Image.open(os.path.join(args.videomatte_dir, 'fgr', clipname, framenames[base_t + t])) as fgr, \ + Image.open(os.path.join(args.videomatte_dir, 'pha', clipname, framenames[base_t + t])) as pha: + fgr = fgr.convert('RGB') + pha = pha.convert('L') + + if args.resize is not None: + fgr = fgr.resize(args.resize, Image.BILINEAR) + pha = pha.resize(args.resize, Image.BILINEAR) + + + if i // len(videomatte_filenames) % 2 == 1: + fgr = fgr.transpose(Image.FLIP_LEFT_RIGHT) + pha = pha.transpose(Image.FLIP_LEFT_RIGHT) + + fgr.save(os.path.join(out_path, 'fgr', str(t).zfill(4) + args.extension)) + pha.save(os.path.join(out_path, 'pha', str(t).zfill(4) + args.extension)) + + if t == 0: + bgr = bgr.resize(fgr.size, Image.BILINEAR) + bgr.save(os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension)) + else: + os.symlink(str(0).zfill(4) + args.extension, os.path.join(out_path, 'bgr', str(t).zfill(4) + args.extension)) + + pha = np.asarray(pha).astype(float)[:, :, None] / 255 + com = Image.fromarray(np.uint8(np.asarray(fgr) * pha + np.asarray(bgr) * (1 - pha))) + com.save(os.path.join(out_path, 'com', str(t).zfill(4) + args.extension)) diff --git a/evaluation/generate_videomatte_with_background_video.py b/evaluation/generate_videomatte_with_background_video.py new file mode 100644 index 0000000..63c24ce --- /dev/null +++ b/evaluation/generate_videomatte_with_background_video.py @@ -0,0 +1,118 @@ +""" +python generate_videomatte_with_background_video.py \ + --videomatte-dir ../matting-data/VideoMatte240K_JPEG_HD/test \ + --background-dir ../matting-data/BackgroundVideos_mp4/test \ + --resize 512 288 \ + --out-dir ../matting-data/evaluation/vidematte_motion_sd/ +""" + +import argparse +import os +import pims +import numpy as np +import random +from PIL import Image +from tqdm import tqdm + +parser = argparse.ArgumentParser() +parser.add_argument('--videomatte-dir', type=str, required=True) +parser.add_argument('--background-dir', type=str, required=True) +parser.add_argument('--num-samples', type=int, default=20) +parser.add_argument('--num-frames', type=int, default=100) +parser.add_argument('--resize', type=int, default=None, nargs=2) +parser.add_argument('--out-dir', type=str, required=True) +args = parser.parse_args() + +# Hand selected a list of videos +background_filenames = [ + "0000.mp4", + "0007.mp4", + "0008.mp4", + "0010.mp4", + "0013.mp4", + "0015.mp4", + "0016.mp4", + "0018.mp4", + "0021.mp4", + "0029.mp4", + "0033.mp4", + "0035.mp4", + "0039.mp4", + "0050.mp4", + "0052.mp4", + "0055.mp4", + "0060.mp4", + "0063.mp4", + "0087.mp4", + "0086.mp4", + "0090.mp4", + "0101.mp4", + "0110.mp4", + "0117.mp4", + "0120.mp4", + "0122.mp4", + "0123.mp4", + "0125.mp4", + "0128.mp4", + "0131.mp4", + "0172.mp4", + "0176.mp4", + "0181.mp4", + "0187.mp4", + "0193.mp4", + "0198.mp4", + "0220.mp4", + "0221.mp4", + "0224.mp4", + "0229.mp4", + "0233.mp4", + "0238.mp4", + "0241.mp4", + "0245.mp4", + "0246.mp4" +] + +random.seed(10) + +videomatte_filenames = [(clipname, sorted(os.listdir(os.path.join(args.videomatte_dir, 'fgr', clipname)))) + for clipname in sorted(os.listdir(os.path.join(args.videomatte_dir, 'fgr')))] + +random.shuffle(background_filenames) + +for i in range(args.num_samples): + bgrs = pims.PyAVVideoReader(os.path.join(args.background_dir, background_filenames[i % len(background_filenames)])) + clipname, framenames = videomatte_filenames[i % len(videomatte_filenames)] + + out_path = os.path.join(args.out_dir, str(i).zfill(4)) + os.makedirs(os.path.join(out_path, 'fgr'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'pha'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'com'), exist_ok=True) + os.makedirs(os.path.join(out_path, 'bgr'), exist_ok=True) + + base_t = random.choice(range(len(framenames) - args.num_frames)) + + for t in tqdm(range(args.num_frames), desc=str(i).zfill(4)): + with Image.open(os.path.join(args.videomatte_dir, 'fgr', clipname, framenames[base_t + t])) as fgr, \ + Image.open(os.path.join(args.videomatte_dir, 'pha', clipname, framenames[base_t + t])) as pha: + fgr = fgr.convert('RGB') + pha = pha.convert('L') + + if args.resize is not None: + fgr = fgr.resize(args.resize, Image.BILINEAR) + pha = pha.resize(args.resize, Image.BILINEAR) + + + if i // len(videomatte_filenames) % 2 == 1: + fgr = fgr.transpose(Image.FLIP_LEFT_RIGHT) + pha = pha.transpose(Image.FLIP_LEFT_RIGHT) + + fgr.save(os.path.join(out_path, 'fgr', str(t).zfill(4) + '.png')) + pha.save(os.path.join(out_path, 'pha', str(t).zfill(4) + '.png')) + + bgr = Image.fromarray(bgrs[t]) + bgr = bgr.resize(fgr.size, Image.BILINEAR) + bgr.save(os.path.join(out_path, 'bgr', str(t).zfill(4) + '.png')) + + pha = np.asarray(pha).astype(float)[:, :, None] / 255 + com = Image.fromarray(np.uint8(np.asarray(fgr) * pha + np.asarray(bgr) * (1 - pha))) + com.save(os.path.join(out_path, 'com', str(t).zfill(4) + '.png')) diff --git a/hubconf.py b/hubconf.py new file mode 100644 index 0000000..9b1ceba --- /dev/null +++ b/hubconf.py @@ -0,0 +1,39 @@ +""" +Loading model + model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3") + model = torch.hub.load("PeterL1n/RobustVideoMatting", "resnet50") + +Converter API + convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter") +""" + + +dependencies = ['torch', 'torchvision'] + +import torch +from model import MattingNetwork + + +def mobilenetv3(pretrained: bool = True, progress: bool = True): + model = MattingNetwork('mobilenetv3') + if pretrained: + url = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3.pth' + model.load_state_dict(torch.hub.load_state_dict_from_url(url, map_location='cpu', progress=progress)) + return model + + +def resnet50(pretrained: bool = True, progress: bool = True): + model = MattingNetwork('resnet50') + if pretrained: + url = 'https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_resnet50.pth' + model.load_state_dict(torch.hub.load_state_dict_from_url(url, map_location='cpu', progress=progress)) + return model + + +def converter(): + try: + from inference import convert_video + return convert_video + except ModuleNotFoundError as error: + print(error) + print('Please run "pip install av tqdm pims"') diff --git a/inference.py b/inference.py new file mode 100644 index 0000000..acf2a74 --- /dev/null +++ b/inference.py @@ -0,0 +1,208 @@ +""" +python inference.py \ + --variant mobilenetv3 \ + --checkpoint "CHECKPOINT" \ + --device cuda \ + --input-source "input.mp4" \ + --output-type video \ + --output-composition "composition.mp4" \ + --output-alpha "alpha.mp4" \ + --output-foreground "foreground.mp4" \ + --output-video-mbps 4 \ + --seq-chunk 1 +""" + +import torch +import os +from torch.utils.data import DataLoader +from torchvision import transforms +from typing import Optional, Tuple +from tqdm.auto import tqdm + +from inference_utils import VideoReader, VideoWriter, ImageSequenceReader, ImageSequenceWriter + +def convert_video(model, + input_source: str, + input_resize: Optional[Tuple[int, int]] = None, + downsample_ratio: Optional[float] = None, + output_type: str = 'video', + output_composition: Optional[str] = None, + output_alpha: Optional[str] = None, + output_foreground: Optional[str] = None, + output_video_mbps: Optional[float] = None, + seq_chunk: int = 1, + num_workers: int = 0, + progress: bool = True, + device: Optional[str] = None, + dtype: Optional[torch.dtype] = None): + + """ + Args: + input_source:A video file, or an image sequence directory. Images must be sorted in accending order, support png and jpg. + input_resize: If provided, the input are first resized to (w, h). + downsample_ratio: The model's downsample_ratio hyperparameter. If not provided, model automatically set one. + output_type: Options: ["video", "png_sequence"]. + output_composition: + The composition output path. File path if output_type == 'video'. Directory path if output_type == 'png_sequence'. + If output_type == 'video', the composition has green screen background. + If output_type == 'png_sequence'. the composition is RGBA png images. + output_alpha: The alpha output from the model. + output_foreground: The foreground output from the model. + seq_chunk: Number of frames to process at once. Increase it for better parallelism. + num_workers: PyTorch's DataLoader workers. Only use >0 for image input. + progress: Show progress bar. + device: Only need to manually provide if model is a TorchScript freezed model. + dtype: Only need to manually provide if model is a TorchScript freezed model. + """ + + assert downsample_ratio is None or (downsample_ratio > 0 and downsample_ratio <= 1), 'Downsample ratio must be between 0 (exclusive) and 1 (inclusive).' + assert any([output_composition, output_alpha, output_foreground]), 'Must provide at least one output.' + assert output_type in ['video', 'png_sequence'], 'Only support "video" and "png_sequence" output modes.' + assert seq_chunk >= 1, 'Sequence chunk must be >= 1' + assert num_workers >= 0, 'Number of workers must be >= 0' + assert output_video_mbps == None or output_type == 'video', 'Mbps is not available for png_sequence output.' + + # Initialize transform + if input_resize is not None: + transform = transforms.Compose([ + transforms.Resize(input_resize[::-1]), + transforms.ToTensor() + ]) + else: + transform = transforms.ToTensor() + + # Initialize reader + if os.path.isfile(input_source): + source = VideoReader(input_source, transform) + else: + source = ImageSequenceReader(input_source, transform) + reader = DataLoader(source, batch_size=seq_chunk, pin_memory=True, num_workers=num_workers) + + # Initialize writers + if output_type == 'video': + frame_rate = source.frame_rate if isinstance(source, VideoReader) else 30 + output_video_mbps = 1 if output_video_mbps is None else output_video_mbps + if output_composition is not None: + writer_com = VideoWriter( + path=output_composition, + frame_rate=frame_rate, + bit_rate=int(output_video_mbps * 1000000)) + if output_alpha is not None: + writer_pha = VideoWriter( + path=output_alpha, + frame_rate=frame_rate, + bit_rate=int(output_video_mbps * 1000000)) + if output_foreground is not None: + writer_fgr = VideoWriter( + path=output_foreground, + frame_rate=frame_rate, + bit_rate=int(output_video_mbps * 1000000)) + else: + if output_composition is not None: + writer_com = ImageSequenceWriter(output_composition, 'png') + if output_alpha is not None: + writer_pha = VideoWriter(output_alpha, 'png') + if output_foreground is not None: + writer_fgr = VideoWriter(output_foreground, 'png') + + # Inference + model = model.eval() + if device is None or dtype is None: + param = next(model.parameters()) + dtype = param.dtype + device = param.device + + if (output_composition is not None) and (output_type == 'video'): + bgr = torch.tensor([120, 255, 155], device=device, dtype=dtype).div(255).view(1, 1, 3, 1, 1) + + try: + with torch.no_grad(): + bar = tqdm(total=len(source), disable=not progress, dynamic_ncols=True) + rec = [None] * 4 + for src in reader: + + if downsample_ratio is None: + downsample_ratio = auto_downsample_ratio(*src.shape[2:]) + + src = src.to(device, dtype, non_blocking=True).unsqueeze(0) # [B, T, C, H, W] + fgr, pha, *rec = model(src, *rec, downsample_ratio) + + if output_foreground is not None: + writer_fgr.write(fgr[0]) + if output_alpha is not None: + writer_pha.write(pha[0]) + if output_composition is not None: + if output_type == 'video': + com = fgr * pha + bgr * (1 - pha) + else: + fgr = fgr * pha.gt(0) + com = torch.cat([fgr, pha], dim=-3) + writer_com.write(com[0]) + + bar.update(src.size(1)) + + finally: + # Clean up + if output_composition is not None: + writer_com.close() + if output_alpha is not None: + writer_pha.close() + if output_foreground is not None: + writer_fgr.close() + + +def auto_downsample_ratio(h, w): + """ + Automatically find a downsample ratio so that the largest side of the resolution be 512px. + """ + return min(512 / max(h, w), 1) + + +class Converter: + def __init__(self, variant: str, checkpoint: str, device: str): + self.model = MattingNetwork(variant).eval().to(device) + self.model.load_state_dict(torch.load(checkpoint, map_location=device)) + self.model = torch.jit.script(self.model) + self.model = torch.jit.freeze(self.model) + self.device = device + + def convert(self, *args, **kwargs): + convert_video(self.model, device=self.device, dtype=torch.float32, *args, **kwargs) + +if __name__ == '__main__': + import argparse + from model import MattingNetwork + + parser = argparse.ArgumentParser() + parser.add_argument('--variant', type=str, required=True, choices=['mobilenetv3', 'resnet50']) + parser.add_argument('--checkpoint', type=str, required=True) + parser.add_argument('--device', type=str, required=True) + parser.add_argument('--input-source', type=str, required=True) + parser.add_argument('--input-resize', type=int, default=None, nargs=2) + parser.add_argument('--downsample-ratio', type=float) + parser.add_argument('--output-composition', type=str) + parser.add_argument('--output-alpha', type=str) + parser.add_argument('--output-foreground', type=str) + parser.add_argument('--output-type', type=str, required=True, choices=['video', 'png_sequence']) + parser.add_argument('--output-video-mbps', type=int, default=1) + parser.add_argument('--seq-chunk', type=int, default=1) + parser.add_argument('--num-workers', type=int, default=0) + parser.add_argument('--disable-progress', action='store_true') + args = parser.parse_args() + + converter = Converter(args.variant, args.checkpoint, args.device) + converter.convert( + input_source=args.input_source, + input_resize=args.input_resize, + downsample_ratio=args.downsample_ratio, + output_type=args.output_type, + output_composition=args.output_composition, + output_alpha=args.output_alpha, + output_foreground=args.output_foreground, + output_video_mbps=args.output_video_mbps, + seq_chunk=args.seq_chunk, + num_workers=args.num_workers, + progress=not args.disable_progress + ) + + \ No newline at end of file diff --git a/inference_speed_test.py b/inference_speed_test.py new file mode 100644 index 0000000..2f6808f --- /dev/null +++ b/inference_speed_test.py @@ -0,0 +1,50 @@ +""" +python inference_speed_test.py \ + --model-variant mobilenetv3 \ + --resolution 1920 1080 \ + --downsample-ratio 0.25 \ + --precision float32 +""" + +import argparse +import torch +from tqdm import tqdm + +from model.model import MattingNetwork + +torch.backends.cudnn.benchmark = True + +class InferenceSpeedTest: + def __init__(self): + self.parse_args() + self.init_model() + self.loop() + + def parse_args(self): + parser = argparse.ArgumentParser() + parser.add_argument('--model-variant', type=str, required=True) + parser.add_argument('--resolution', type=int, required=True, nargs=2) + parser.add_argument('--downsample-ratio', type=float, required=True) + parser.add_argument('--precision', type=str, default='float32') + parser.add_argument('--disable-refiner', action='store_true') + self.args = parser.parse_args() + + def init_model(self): + self.device = 'cuda' + self.precision = {'float32': torch.float32, 'float16': torch.float16}[self.args.precision] + self.model = MattingNetwork(self.args.model_variant) + self.model = self.model.to(device=self.device, dtype=self.precision).eval() + self.model = torch.jit.script(self.model) + self.model = torch.jit.freeze(self.model) + + def loop(self): + w, h = self.args.resolution + src = torch.randn((1, 3, h, w), device=self.device, dtype=self.precision) + with torch.no_grad(): + rec = None, None, None, None + for _ in tqdm(range(1000)): + fgr, pha, *rec = self.model(src, *rec, self.args.downsample_ratio) + torch.cuda.synchronize() + +if __name__ == '__main__': + InferenceSpeedTest() \ No newline at end of file diff --git a/inference_utils.py b/inference_utils.py new file mode 100644 index 0000000..1e92fd4 --- /dev/null +++ b/inference_utils.py @@ -0,0 +1,88 @@ +import av +import os +import pims +import numpy as np +from torch.utils.data import Dataset +from torchvision.transforms.functional import to_pil_image +from PIL import Image + + +class VideoReader(Dataset): + def __init__(self, path, transform=None): + self.video = pims.PyAVVideoReader(path) + self.rate = self.video.frame_rate + self.transform = transform + + @property + def frame_rate(self): + return self.rate + + def __len__(self): + return len(self.video) + + def __getitem__(self, idx): + frame = self.video[idx] + frame = Image.fromarray(np.asarray(frame)) + if self.transform is not None: + frame = self.transform(frame) + return frame + + +class VideoWriter: + def __init__(self, path, frame_rate, bit_rate=1000000): + self.container = av.open(path, mode='w') + self.stream = self.container.add_stream('h264', rate=round(frame_rate)) + self.stream.pix_fmt = 'yuv420p' + self.stream.bit_rate = bit_rate + + def write(self, frames): + # frames: [T, C, H, W] + self.stream.width = frames.size(3) + self.stream.height = frames.size(2) + if frames.size(1) == 1: + frames = frames.repeat(1, 3, 1, 1) # convert grayscale to RGB + frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy() + for t in range(frames.shape[0]): + frame = frames[t] + frame = av.VideoFrame.from_ndarray(frame, format='rgb24') + self.container.mux(self.stream.encode(frame)) + + def close(self): + self.container.mux(self.stream.encode()) + self.container.close() + + +class ImageSequenceReader(Dataset): + def __init__(self, path, transform=None): + self.path = path + self.files = sorted(os.listdir(path)) + self.transform = transform + + def __len__(self): + return len(self.files) + + def __getitem__(self, idx): + with Image.open(os.path.join(self.path, self.files[idx])) as img: + img.load() + if self.transform is not None: + return self.transform(img) + return img + + +class ImageSequenceWriter: + def __init__(self, path, extension='jpg'): + self.path = path + self.extension = extension + self.counter = 0 + os.makedirs(path, exist_ok=True) + + def write(self, frames): + # frames: [T, C, H, W] + for t in range(frames.shape[0]): + to_pil_image(frames[t]).save(os.path.join( + self.path, str(self.counter).zfill(4) + '.' + self.extension)) + self.counter += 1 + + def close(self): + pass + \ No newline at end of file diff --git a/model/__init__.py b/model/__init__.py new file mode 100644 index 0000000..ac047a1 --- /dev/null +++ b/model/__init__.py @@ -0,0 +1 @@ +from .model import MattingNetwork \ No newline at end of file diff --git a/model/decoder.py b/model/decoder.py new file mode 100644 index 0000000..7307435 --- /dev/null +++ b/model/decoder.py @@ -0,0 +1,210 @@ +import torch +from torch import Tensor +from torch import nn +from torch.nn import functional as F +from typing import Tuple, Optional + +class RecurrentDecoder(nn.Module): + def __init__(self, feature_channels, decoder_channels): + super().__init__() + self.avgpool = AvgPool() + self.decode4 = BottleneckBlock(feature_channels[3]) + self.decode3 = UpsamplingBlock(feature_channels[3], feature_channels[2], 3, decoder_channels[0]) + self.decode2 = UpsamplingBlock(decoder_channels[0], feature_channels[1], 3, decoder_channels[1]) + self.decode1 = UpsamplingBlock(decoder_channels[1], feature_channels[0], 3, decoder_channels[2]) + self.decode0 = OutputBlock(decoder_channels[2], 3, decoder_channels[3]) + + def forward(self, + s0: Tensor, f1: Tensor, f2: Tensor, f3: Tensor, f4: Tensor, + r1: Optional[Tensor], r2: Optional[Tensor], + r3: Optional[Tensor], r4: Optional[Tensor]): + s1, s2, s3 = self.avgpool(s0) + x4, r4 = self.decode4(f4, r4) + x3, r3 = self.decode3(x4, f3, s3, r3) + x2, r2 = self.decode2(x3, f2, s2, r2) + x1, r1 = self.decode1(x2, f1, s1, r1) + x0 = self.decode0(x1, s0) + return x0, r1, r2, r3, r4 + + +class AvgPool(nn.Module): + def __init__(self): + super().__init__() + self.avgpool = nn.AvgPool2d(2, 2, count_include_pad=False, ceil_mode=True) + + def forward_single_frame(self, s0): + s1 = self.avgpool(s0) + s2 = self.avgpool(s1) + s3 = self.avgpool(s2) + return s1, s2, s3 + + def forward_time_series(self, s0): + B, T = s0.shape[:2] + s0 = s0.flatten(0, 1) + s1, s2, s3 = self.forward_single_frame(s0) + s1 = s1.unflatten(0, (B, T)) + s2 = s2.unflatten(0, (B, T)) + s3 = s3.unflatten(0, (B, T)) + return s1, s2, s3 + + def forward(self, s0): + if s0.ndim == 5: + return self.forward_time_series(s0) + else: + return self.forward_single_frame(s0) + + +class BottleneckBlock(nn.Module): + def __init__(self, channels): + super().__init__() + self.channels = channels + self.gru = ConvGRU(channels // 2) + + def forward(self, x, r: Optional[Tensor]): + a, b = x.split(self.channels // 2, dim=-3) + b, r = self.gru(b, r) + x = torch.cat([a, b], dim=-3) + return x, r + + +class UpsamplingBlock(nn.Module): + def __init__(self, in_channels, skip_channels, src_channels, out_channels): + super().__init__() + self.out_channels = out_channels + self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) + self.conv = nn.Sequential( + nn.Conv2d(in_channels + skip_channels + src_channels, out_channels, 3, 1, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(True), + ) + self.gru = ConvGRU(out_channels // 2) + + def forward_single_frame(self, x, f, s, r: Optional[Tensor]): + x = self.upsample(x) + x = x[:, :, :s.size(2), :s.size(3)] + x = torch.cat([x, f, s], dim=1) + x = self.conv(x) + a, b = x.split(self.out_channels // 2, dim=1) + b, r = self.gru(b, r) + x = torch.cat([a, b], dim=1) + return x, r + + def forward_time_series(self, x, f, s, r: Optional[Tensor]): + B, T, _, H, W = s.shape + x = x.flatten(0, 1) + f = f.flatten(0, 1) + s = s.flatten(0, 1) + x = self.upsample(x) + x = x[:, :, :H, :W] + x = torch.cat([x, f, s], dim=1) + x = self.conv(x) + x = x.unflatten(0, (B, T)) + a, b = x.split(self.out_channels // 2, dim=2) + b, r = self.gru(b, r) + x = torch.cat([a, b], dim=2) + return x, r + + def forward(self, x, f, s, r: Optional[Tensor]): + if x.ndim == 5: + return self.forward_time_series(x, f, s, r) + else: + return self.forward_single_frame(x, f, s, r) + + +class OutputBlock(nn.Module): + def __init__(self, in_channels, src_channels, out_channels): + super().__init__() + self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False) + self.conv = nn.Sequential( + nn.Conv2d(in_channels + src_channels, out_channels, 3, 1, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(True), + nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU(True), + ) + + def forward_single_frame(self, x, s): + x = self.upsample(x) + x = x[:, :, :s.size(2), :s.size(3)] + x = torch.cat([x, s], dim=1) + x = self.conv(x) + return x + + def forward_time_series(self, x, s): + B, T, _, H, W = s.shape + x = x.flatten(0, 1) + s = s.flatten(0, 1) + x = self.upsample(x) + x = x[:, :, :H, :W] + x = torch.cat([x, s], dim=1) + x = self.conv(x) + x = x.unflatten(0, (B, T)) + return x + + def forward(self, x, s): + if x.ndim == 5: + return self.forward_time_series(x, s) + else: + return self.forward_single_frame(x, s) + + +class ConvGRU(nn.Module): + def __init__(self, + channels: int, + kernel_size: int = 3, + padding: int = 1): + super().__init__() + self.channels = channels + self.ih = nn.Sequential( + nn.Conv2d(channels * 2, channels * 2, kernel_size, padding=padding), + nn.Sigmoid() + ) + self.hh = nn.Sequential( + nn.Conv2d(channels * 2, channels, kernel_size, padding=padding), + nn.Tanh() + ) + + def forward_single_frame(self, x, h): + r, z = self.ih(torch.cat([x, h], dim=1)).split(self.channels, dim=1) + c = self.hh(torch.cat([x, r * h], dim=1)) + h = (1 - z) * h + z * c + return h, h + + def forward_time_series(self, x, h): + o = [] + for xt in x.unbind(dim=1): + ot, h = self.forward_single_frame(xt, h) + o.append(ot) + o = torch.stack(o, dim=1) + return o, h + + def forward(self, x, h: Optional[Tensor]): + if h is None: + h = torch.zeros((x.size(0), x.size(-3), x.size(-2), x.size(-1)), + device=x.device, dtype=x.dtype) + + if x.ndim == 5: + return self.forward_time_series(x, h) + else: + return self.forward_single_frame(x, h) + + +class Projection(nn.Module): + def __init__(self, in_channels, out_channels): + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, 1) + + def forward_single_frame(self, x): + return self.conv(x) + + def forward_time_series(self, x): + B, T = x.shape[:2] + return self.conv(x.flatten(0, 1)).unflatten(0, (B, T)) + + def forward(self, x): + if x.ndim == 5: + return self.forward_time_series(x) + else: + return self.forward_single_frame(x) + \ No newline at end of file diff --git a/model/deep_guided_filter.py b/model/deep_guided_filter.py new file mode 100644 index 0000000..a24b8c5 --- /dev/null +++ b/model/deep_guided_filter.py @@ -0,0 +1,61 @@ +import torch +from torch import nn +from torch.nn import functional as F + +""" +Adopted from