fix image2mask issue

This commit is contained in:
chflame163
2024-10-23 18:18:34 +08:00
parent fbdfff3555
commit 2df6cf23ed
2 changed files with 6 additions and 4 deletions
+5 -3
View File
@@ -138,9 +138,11 @@ def tensor2cv2(image:torch.Tensor) -> np.array:
return cv2.cvtColor(cv2image, cv2.COLOR_RGB2BGR)
def image2mask(image:Image) -> torch.Tensor:
if image.mode != 'L':
image = image.convert('L')
return torch.tensor([pil2tensor(image)[0, :, :].tolist()])
if image.mode == 'L':
return torch.tensor([pil2tensor(image)[0, :, :].tolist()])
else:
image = image.convert('RGB').split()[0]
return torch.tensor([pil2tensor(image)[0, :, :].tolist()])
def mask2image(mask:torch.Tensor) -> Image:
masks = tensor2np(mask)