MODNet: Trimap-Free Portrait Matting in Real Time

i am making the train dataset, it needs 3 folders--original, trimap, matter. so the size of image must be 512,? and the image need to do other operations, like change the color and so on ? what else should I pay attention to do in the dataset ?

Hello Thanks for contributing to MODNet project. I want to know if the code about OFD strategy is in the inference file? How to apply real-time inference about OFD strategy?

MODNET is an impressive work in solving image and video matting problems. This work has a non-commercial license. Could you please provide a more convenient license like MIT or Apache license that allows for commercial use. This would make it easier for this work to be used in open source projects. Non-commercial License work is not compatible for use in open source projects. You can read more about the incompatibility of non-commercial work from here. I am unable to make use of this work in my open source project because of the non-commercial license.

Hi,

Thanks for this great project. I tried your colab for image matting, it looks like the boundary is not clear enough for some inputs(also the one in Github readme).

Is there anyway to improve image matting accuracy?

Hello, I was just wondering, what dataset are you using for your pre-trained model?I'm currently training with Adobe Matting datasets, but I'm not getting the results I want.

Hello How are you? Thanks for contributing to MODNet project. I have a question about the SOC process. I know that the SOC process is to refine alpha matte by using the predicted segmentation on unlabeled samples.

If so, I think that this is based on the supposition that the predicted segmentation result on even unlabeled(unseen) samples are always good but the alpha mattes are NOT good relatively. Then what is the guarantee that the segmentation result on unlabeled(unseen) images by MODNet is always good?

Thanks for your sharing.Nice work~ here is a question about the SOC in your paper. The self-supervied stage is used in the new domain datasets, so the new or the target datasets are which we will test later?

And another question is when i try to train MODNet, the prediction of dp if just boundary which is not same as your paper .

Dear @ZHKKKe I hope you can help me and others to come,

I'm pretty new to this (also I'm not a programmer) but I'm very fascinated about your great research and I'm trying to get higher quality, more accurate results like in your personal website's DEMO.

I tried the Collab version and the Local version as I'm using Windows 10 + Anaconda.

At first I wasn't sure but then I noticed that there is a HUGE different with the current pretrained model compare to YOUR website demo.

So I'm not sure if there is a way to download and experiment with the same pretrained model you use on your website demo, that will be very interesting to try and do comparisons on local machine.

I don't really know how to TRAIN my own dataset, I tried to look but it is not very clear to me, I wish it was simple as running the demo on anaconda. (maybe it is not hard, but I didn't find a step-by-step guide or tutorial which would be VERY helpful)

Here are just couple of examples between the ModNet pretrained model and your online website which is much more accurate:

Modnet (collab) - 1

ZHKE (personal website) - 1

Modnet (collab) - 2

ZHKE (personal website) - 2

I'm enjoying doing these comparisons, I must mention so far... Your @ZHKKKe website demo are much more accurate and cleaner, not in all areas of course probably I would love to keep training it even more to make it more accurate!

Some interesting numbers I've tested, roughly: 16 out of 20 (portrait human) images where very accurate (ZHKKKe's demo website) compare to a messy results on the pretrained model (about 25 MB) used in Local and Collab ModNet. 4 images pretty much failed on hair/background and added some not-needed objects in the background.

I'm now very curious about the pretrained model you're using on your demo website and also about training my own dataset (if I'll understand how to do it of course...)

Please consider to share the pretrained model you used on your website demo. and if there is a guide or tutorial explain how we can Train our own dataset, it will be VERY helpful as well.

I would like to experiment and help in your research by showing comparisons or improve training if I'll know how to do it on my local machine.

I hope you can help in this, thanks ahead and please keep up the wonderful work, YOU ROCKS! 😎

When I export the TorchScript version of MODNet, error occurs:

Traceback (most recent call last):
  File "/usr/local/lib/python3.6/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/local/lib/python3.6/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/home/sk49/new_workspace/cz/MODNet-master/torchscript/export_torchscript.py", line 46, in <module>
    scripted_model = torch.jit.script(modnet.module)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/__init__.py", line 1203, in script
    return torch.jit.torch.jit._recursive.recursive_script(obj)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 173, in recursive_script
    return copy_to_script_module(mod, overload_stubs + stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 95, in copy_to_script_module
    torch.jit._create_methods_from_stubs(script_module, stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/__init__.py", line 1423, in _create_methods_from_stubs
    self._c._create_methods(self, defs, rcbs, defaults)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 195, in make_strong_submodule
    new_strong_submodule = recursive_script(module)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 173, in recursive_script
    return copy_to_script_module(mod, overload_stubs + stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 95, in copy_to_script_module
    torch.jit._create_methods_from_stubs(script_module, stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/__init__.py", line 1423, in _create_methods_from_stubs
    self._c._create_methods(self, defs, rcbs, defaults)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 195, in make_strong_submodule
    new_strong_submodule = recursive_script(module)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 173, in recursive_script
    return copy_to_script_module(mod, overload_stubs + stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 95, in copy_to_script_module
    torch.jit._create_methods_from_stubs(script_module, stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/__init__.py", line 1423, in _create_methods_from_stubs
    self._c._create_methods(self, defs, rcbs, defaults)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 195, in make_strong_submodule
    new_strong_submodule = recursive_script(module)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 173, in recursive_script
    return copy_to_script_module(mod, overload_stubs + stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/_recursive.py", line 95, in copy_to_script_module
    torch.jit._create_methods_from_stubs(script_module, stubs)
  File "/home/sk49/.local/lib/python3.6/site-packages/torch/jit/__init__.py", line 1423, in _create_methods_from_stubs
    self._c._create_methods(self, defs, rcbs, defaults)
RuntimeError: 
'module' object is not subscriptable:
at /home/sk49/new_workspace/cz/MODNet-master/src/models/backbones/mobilenetv2.py:141:6
	def forward(self, x):
		# Stage1
		x = self.features[0](x)
      ~~~~~~~~~~~~~~~ <--- HERE
		x = self.features[1](x)
		# Stage2
		x = self.features[2](x)
		x = self.features[3](x)
		# Stage3
		x = self.features[4](x)
		x = self.features[5](x)
		x = self.features[6](x)
		# Stage4
'__torch__.src.models.backbones.mobilenetv2.MobileNetV2.forward' is being compiled since it was called from '__torch__.src.models.backbones.wrapper.MobileNetV2Backbone.forward'
at /home/sk49/new_workspace/cz/MODNet-master/src/models/backbones/wrapper.py:40:8
    def forward(self, x):
        # x = reduce(lambda x, n: self.model.features[n](x), list(range(0, 2)), x)
        x = self.model.features[0](x)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        x = self.model.features[1](x)
        enc2x = x

        # x = reduce(lambda x, n: self.model.features[n](x), list(range(2, 4)), x)
        x = self.model.features[2](x)
        x = self.model.features[3](x)
        enc4x = x

        # x = reduce(lambda x, n: self.model.features[n](x), list(range(4, 7)), x)
'__torch__.src.models.backbones.wrapper.MobileNetV2Backbone.forward' is being compiled since it was called from '__torch__.torchscript.modnet_torchscript.LRBranch.forward'
at /home/sk49/new_workspace/cz/MODNet-master/torchscript/modnet_torchscript.py:108:8
    def forward(self, img):
        enc_features = self.backbone.forward(img)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        enc2x, enc4x, enc32x = enc_features[0], enc_features[1], enc_features[4]

        enc32x = self.se_block(enc32x)
        lr16x = F.interpolate(enc32x, scale_factor=2.0, mode='bilinear', align_corners=False)
        lr16x = self.conv_lr16x(lr16x)
        lr8x = F.interpolate(lr16x, scale_factor=2.0, mode='bilinear', align_corners=False)
        lr8x = self.conv_lr8x(lr8x)

        return lr8x, enc2x, enc4x
'__torch__.torchscript.modnet_torchscript.LRBranch.forward' is being compiled since it was called from '__torch__.torchscript.modnet_torchscript.MODNet.forward'
at /home/sk49/new_workspace/cz/MODNet-master/torchscript/modnet_torchscript.py:230:8
    def forward(self, img):
        # NOTE
        lr_out = self.lr_branch(img)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        lr8x = lr_out[0]
        enc2x = lr_out[1]
        enc4x = lr_out[2]

        hr2x = self.hr_branch(img, enc2x, enc4x, lr8x)
        
        pred_matte = self.f_branch(img, lr8x, hr2x)

        return pred_matte

Environment: Centos 7 torch 1.3.1 gpu

 view_np = matte_np * frame_np + (1 - matte_np) * frame_dy

只这步就得25~28ms，我们就拿26来算，也就38帧，是怎么得到一秒钟63帧的，而且分辨率好低，如果用HD或者4K的视频，

吓死个人好吧

_, _, matte_tensor = modnet(frame_tensor, True)

这步确实是快，14ms，只按照这个算的话，一秒钟可以71帧，但是不能这么算啊

我就不说modnet = MODNet(backbone_pretrained=False)

modnet.load_state_dict(torch.load(pretrained_ckpt))这两部分了

求作者给个指点呀

我微信nangongqiuhan003

@ZHKKKe ，您好！ 我训练了40个epoch，matte_loss从3.587减少到： epoch = 39, semantic_loss=0.0039, detail_loss=0.0332, matte_loss=2.2689 epoch = 39, semantic_loss=0.0053, detail_loss=0.0327, matte_loss=2.4930 epoch = 39, semantic_loss=0.0067, detail_loss=0.0446, matte_loss=2.5472 epoch = 39, semantic_loss=0.0034, detail_loss=0.0272, matte_loss=2.2228 epoch = 39, semantic_loss=0.0050, detail_loss=0.0378, matte_loss=2.6175 epoch = 39, semantic_loss=0.0046, detail_loss=0.0338, matte_loss=2.3804 epoch = 39, semantic_loss=0.0061, detail_loss=0.0410, matte_loss=2.2642 epoch = 39, semantic_loss=0.0054, detail_loss=0.0399, matte_loss=2.4980

我使用的爱分割公开的数据集，共34000张，matte_loss减小幅度没有detail_loss那么大，甚至每一个epoch中都会出现3.1左右的值，不止一次。我的bs设置为16.

我想请教一下，这个是不是因为样本的原因导致的。另外，我第一个训练抽取其中10000张样本，但是融合结果细节损失很严重。谢谢！

Hello, I am a student of Tsinghua University. From your paper, I know that modnet includes segmentation, detail prediction and fusion. I want to ask whether the backbone of the segmentation can be replaced with other networks, such as switchtransformer. If so, can it be replaced completely or only partially? In addition, how to deal with the interaction between the segmentation and detail, In particular, the paper mentioned that the characteristics of the middle layer of the segmentation module are used as the input of the detail prediction module. Thank you very much

按照论文说法， modnet包含分割，细节预测和融合三个部分，我想问的是， 问题1，如果把分割部分的backbone换成其他网络，是否可行？例如换成efficintnet，是否可以无缝替换？ 问题2，分割部分和其他部分之间的交互问题，例如，其中细节预测部分用到了部分分割部分的中间层的特征信息，如果把分割的backbone换成其他网络，例如efficinet网络，那么中间层的特征交互问题该如何处理呢？ 谢谢，感激不尽

I have tried training MODNet on the publicly available 30k-image dataset and the performance was poor due to the unclean dataset. Now, I've changed the dataset and chosen the P3M-10k dataset that uses 10k images with good-quality of segmentation. Now I'm fine-tuning the model using this code where I'm training on top of the existing modnet_photographic_portrait_matting.ckpt with backbone_pretrained = True.

My losses are however drastically high. I've seen losses to be as low as this but my losses are: Semantic Loss: 9.234, Detail Loss: 0.334, Matte Loss: 6.392,

and seem to roll around in the same vicinity. Could you please check the code and see if there's any mistake in my data preparation/augmentation or is my training code somewhat wildly incorrect?

I had to review a lot of documentation and issues to implement the training code. So here is the code you'll be needing for training.

Initially, you'll need to download the pretrained model files from https://drive.google.com/drive/folders/1umYmlCulvIFNaqPjwod1SayFmSRHziyR?usp=sharing and move it to MODNet/pretrained. In case you need to fine-tune the model to your own dataset, download - modnet_photographic_portrait_matting.ckpt. In case you need to use the backbone mobilenetv2 model, download that too.

For preparing the dataset, I prepared a pandas dataframe which had 2 columns - ["image", "matte"] "image" had the absolute path to the images' location and "matte" had that respective image's matte image location.

After downloading, for preprocessing, the code is:

class ModNetDataLoader(Dataset):
    def __init__(self, annotations_file, resize_dim, transform=None):
        self.img_labels =annotations_file
        self.transform=transform
        self.resize_dim=resize_dim

    def __len__(self):
        #return the total number of images
        return len(self.img_labels)

    def __getitem__(self, idx):
        img_path = self.img_labels.iloc[idx,0]
        mask_path = self.img_labels.iloc[idx,1]

        img = np.asarray(Image.open(img_path))

        in_image = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
        mask = in_image[:,:,3]

        if len(img.shape)==2:
            img = img[:,:,None]
        if img.shape[2]==1:
            img = np.repeat(img, 3, axis=2)
        elif img.shape[2]==4:
            img = img[:,:,0:3]

        if len(mask.shape)==3:
            mask = mask[:,:, 0]

        #convert Image to pytorch tensor
        img = Image.fromarray(img)
        mask = Image.fromarray(mask)
        if self.transform:
            img = self.transform(img)
            trimap = self.get_trimap(mask)
            mask = self.transform(mask)

        img = self._resize(img)
        mask = self._resize(mask)
        trimap = self._resize(trimap, trimap=True)

        img = torch.squeeze(img, 0)
        mask = torch.squeeze(mask, 0)
        trimap = torch.squeeze(trimap, 1)

        return img, trimap, mask

    def get_trimap(self, alpha):
        # alpha \in [0, 1] should be taken into account
        # be careful when dealing with regions of alpha=0 and alpha=1
        fg = np.array(np.equal(alpha, 255).astype(np.float32))
        unknown = np.array(np.not_equal(alpha, 0).astype(np.float32)) # unknown = alpha > 0
        unknown = unknown - fg
        # image dilation implemented by Euclidean distance transform
        unknown = morphology.distance_transform_edt(unknown==0) <= np.random.randint(1, 20)
        trimap = fg
        trimap[unknown] = 0.5
        return torch.unsqueeze(torch.from_numpy(trimap), dim=0)#.astype(np.uint8)

    def _resize(self, img, trimap=False):
        im = img[None, :, :, :]
        ref_size = self.resize_dim

        # resize image for input
        im_b, im_c, im_h, im_w = im.shape
        if max(im_h, im_w) < ref_size or min(im_h, im_w) > ref_size:
            if im_w >= im_h:
                im_rh = ref_size
                im_rw = int(im_w / im_h * ref_size)
            elif im_w < im_h:
                im_rw = ref_size
                im_rh = int(im_h / im_w * ref_size)
        else:
            im_rh = im_h
            im_rw = im_w

        im_rw = im_rw - im_rw % 32
        im_rh = im_rh - im_rh % 32
        if trimap == True:
            im = F.interpolate(im, size=(im_rh, im_rw), mode='nearest')
        else:
            im = F.interpolate(im, size=(im_rh, im_rw), mode='area')
        return im

You might need to change the above code in methods, get_trimap and get_item according to your dataset You would need to verify if your data is proper in the next to next step

Finally, create your dataset using the code below:

transformer = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize((0.5), (0.5)
            )
        ]
    )
data = ModNetDataLoader(data_csv, 512, transform=transformer)

After your dataset has been created, 1st verify it by printing the first row of data and verifying if the shapes of image, matte, trimap are equal (only the channels can be different). IMPORTANT: Try printing the first of the trimaps. The only values in the numpy array should be 0, 0.5 and 1. Use the dataloader function to prepare your data for training:

train_dataloader = DataLoader(data, batch_size=8, shuffle=True)

After this, the code for training is available in the trainer.py file:

import torch
from src.models.modnet import MODNet
from src.trainer import supervised_training_iter
bs = 16         # batch size
lr = 0.01       # learn rate
epochs = 40     # total epochs

modnet = torch.nn.DataParallel(MODNet()).cuda()
optimizer = torch.optim.SGD(modnet.parameters(), lr=lr, momentum=0.9)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(0.25 * epochs), gamma=0.1)

for epoch in range(0, epochs):
    for idx, (image, trimap, gt_matte) in enumerate(dataloader):
        semantic_loss, detail_loss, matte_loss = \
            supervised_training_iter(modnet, optimizer, image.cuda(), trimap.cuda(), gt_matte.cuda())
    lr_scheduler.step()

For using the backbone - Change modnet = torch.nn.DataParallel(MODNet()).cuda() to modnet = torch.nn.DataParallel(MODNet(backbone_pretrained=True)).cuda() in case you have the mobilenetv2 in your pretrained directory.

For fine-tuning the existing MODNet model, use this snippet before the optimizer line:

modnet = torch.nn.DataParallel(MODNet()).cuda()
state_dict = torch.load("path_to_torchscript_model.ckpt")
modnet.load_state_dict(state_dict)
modnet.train()