From 0e59ff09c9b77a778f705f7193d620044c4648d1 Mon Sep 17 00:00:00 2001 From: chflame163 Date: Tue, 22 Oct 2024 20:00:03 +0800 Subject: [PATCH] fix image2mask function --- py/imagefunc.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/py/imagefunc.py b/py/imagefunc.py index 09d21bb..2168a02 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -138,12 +138,9 @@ def tensor2cv2(image:torch.Tensor) -> np.array: return cv2.cvtColor(cv2image, cv2.COLOR_RGB2BGR) def image2mask(image:Image) -> torch.Tensor: - _image = image.convert('RGBA') - alpha = _image.split() [0] - bg = Image.new("L", _image.size) - _image = Image.merge('RGBA', (bg, bg, bg, alpha)) - ret_mask = torch.tensor([pil2tensor(_image)[0, :, :, 3].tolist()]) - return ret_mask + if image.mode != 'L': + image = image.convert('L') + return torch.tensor([pil2tensor(image)[0, :, :].tolist()]) def mask2image(mask:torch.Tensor) -> Image: masks = tensor2np(mask)