fix image2mask function

2024-10-22 20:00:03 +08:00
parent f657006b74
commit 0e59ff09c9
1 changed files with 3 additions and 6 deletions
@@ -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)