mirror
/
MODNet
mirror of https://github.com/ZHKKKe/MODNet.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix the model forward bug caused by old PyTorch version

pull/35/head
ZHKKKe 2020-12-17 16:46:22 +08:00
parent 60760cea59
commit fea4e3a2ee
2 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
demo/image_matting/colab/inference.py
View File

@ -90,7 +90,7 @@ if __name__ == '__main__':
im = F.interpolate(im, size=(im_rh, im_rw), mode='area')
# inference
_, _, matte = modnet(im.cuda(), inference=False)
_, _, matte = modnet(im.cuda(), True)
# resize and save matte
matte = F.interpolate(matte, size=(im_h, im_w), mode='area')

2
demo/video_matting/webcam/run.py
View File

@ -52,7 +52,7 @@ while(True):
frame_tensor = frame_tensor.cuda()
with torch.no_grad():
_, _, matte_tensor = modnet(frame_tensor, inference=True)
_, _, matte_tensor = modnet(frame_tensor, True)
matte_tensor = matte_tensor.repeat(1, 3, 1, 1)
matte_np = matte_tensor[0].data.cpu().numpy().transpose(1, 2, 0)