[Transfer to CPU version] Refence BackgroundMattingV2 Colab you posted

[Charmve]

Hi, thanks for your awesome work!

description

I have tried code you posed in colab. However, those code are running in GPU not CPU, I, unfortunatly, have not GPU. Then I wirite a CPU version according to your code.

Could you please help me to address this issue?

virtual-machine:~/BackgroundMattingV2$ python3 test.py
Traceback (most recent call last):
  File "test.py", line 25, in <module>
    pha, fgr = model(src, bgr)[:2]
TypeError: 'collections.OrderedDict' object is not callable

attachment code: python3

import torch
from torchvision.transforms.functional import to_tensor, to_pil_image
from PIL import Image

# model = torch.jit.load('model/pytorch_resnet50.pth').cuda().eval() 
"""
RuntimeError: Attempting to deserialize object on a CUDA device but 
torch.cuda.is_available() is False. If you are running on a CPU-only machine, 
please use torch.load with map_location=torch.device('cpu') to map your 
storages to the CPU.
"""

model = torch.load('model/pytorch_resnet50.pth',map_location ='cpu')

src = Image.open('images/img/12.png')
bgr = Image.open('images/bgr/12.png')

src = to_tensor(src)
bgr = to_tensor(bgr)

if src.size(1) <= 2048 and src.size(2) <= 2048:
  model.backbone_scale = 1/4
  model.refine_sample_pixels = 80_000
else:
  model.backbone_scale = 1/8
  model.refine_sample_pixels = 320_000

pha, fgr = model(src, bgr)[:2]

com = pha * fgr + (1 - pha) * torch.tensor([120/255, 255/255, 155/255], device='cuda').view(1, 3, 1, 1)

to_pil_image(com[0].cpu())

to_pil_image(pha[0].cpu()).save('pha.png')
to_pil_image(fgr[0].cpu()).save('fgr.png')
to_pil_image(com[0].cpu()).save('com.png')

Looking forward to your reply!

Best wishes, @Charmve

[PeterL1n]

You are loading the model wrong.

Use the original line that you currently comment out. Just remove .cuda(). Also, to use that line, you need to download the TorchScript weights, not the PyTorch weights.

[Charmve]

Thanks for your reply！

The model is loaded successful.

But, when I use if src.size(2) <= 2048 and src.size(3) <= 2048: line, it goes wrong.

virtual-machine:~/BackgroundMattingV2$ python3 test.py
Traceback (most recent call last):
  File "test.py", line 18, in <module>
    if src.size(2) <= 2048 and src.size(3) <= 2048:
IndexError: Dimension out of range (expected to be in range of [-3, 2], but got 3)

After I modified it to if src.size(1) <= 2048 and src.size(2) <= 2048:, it report that

virtual-machine:~/BackgroundMattingV2$ python3 test.py
Traceback (most recent call last):
  File "test.py", line 25, in <module>
    pha, fgr = model(src, bgr)[:2]
  File "/home/charmve/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/model/model.py", line 28, in forward
    _5 = torch.eq(torch.mul(_4, 4), torch.size(src, 2))
    if _5:
      _7 = torch.floordiv(torch.size(src, 3), 4)
                          ~~~~~~~~~~ <--- HERE
      _8 = torch.eq(torch.mul(_7, 4), torch.size(src, 3))
      _6 = _8

Traceback of TorchScript, original code (most recent call last):
  File "/projects/grail/linsh/matting/model/model.py", line 160, in forward
    def forward(self, src, bgr):
        assert src.size() == bgr.size(), 'src and bgr must have the same shape'
        assert src.size(2) // 4 * 4 == src.size(2) and src.size(3) // 4 * 4 == src.size(3), \
                                                       ~~~~~~~~ <--- HERE
            'src and bgr must have width and height that are divisible by 4'

RuntimeError: Dimension out of range (expected to be in range of [-3, 2], but got 3)

attachment: original code python3

import torch
from torchvision.transforms.functional import to_tensor, to_pil_image
from PIL import Image

model = torch.jit.load('model/TorchScript/torchscript_resnet50_fp32.pth').eval() 

src = Image.open('images/img/12.png')
bgr = Image.open('images/bgr/12.png')

src = to_tensor(src)
bgr = to_tensor(bgr)

if src.size(2) <= 2048 and src.size(3) <= 2048:
  model.backbone_scale = 1/4
  model.refine_sample_pixels = 80_000
else:
  model.backbone_scale = 1/8
  model.refine_sample_pixels = 320_000

pha, fgr = model(src, bgr)[:2]

com = pha * fgr + (1 - pha) * torch.tensor([120/255, 255/255, 155/255], device='cpu').view(1, 3, 1, 1)

to_pil_image(com[0].cpu())

to_pil_image(pha[0].cpu()).save('pha.png')
to_pil_image(fgr[0].cpu()).save('fgr.png')
to_pil_image(com[0].cpu()).save('com.png')

note: com = pha * fgr + (1 - pha) * torch.tensor([120/255, 255/255, 155/255], device='cpu').view(1, 3, 1, 1) instead of device='cuda'.

[PeterL1n]

Change to

src = to_tensor(src).unsqueeze(0)
bgr = to_tensor(bgr).unsqueeze(0)

Original tensor has [Channel, Height, Width] or [C, H, W] shape. Unsqueeze adds another dimension at index 0, making it [1, C, H, W]. The first dimension is the batch dimension, or [B, C, H, W]. It is a convention for neural networks to have a batch dimension, such that you can process multiple images concurrently by stacking them in the batch dimension.

[Charmve]

Thank you very much！