问题:今天跑完模型代码分割出来的图片不是二值图,而是如下问题,蒙上了一层雾气一样,不是我想要的二值图
错误代码:
import cv2
output=model(img)
pred = torch.sigmoid(output)
img = pred[0].mul(255).byte()
img = img.cpu().numpy().transpose((1, 2, 0))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
cv2.imwrite("./result/{}.png".format(epoch_num), img)
原因是未进行二值化操作,记录正确使用代码保存预测图片,代码如下:
from PIL import Image
def mask_to_image(mask):
return Image.fromarray((mask * 255).astype(np.uint8))
output=model(img)
pred = torch.sigmoid(output)
probs = pred.squeeze(0)
tf = transforms.Compose(
[
transforms.ToPILImage(),
transforms.ToTensor()
]
)
probs = tf(probs.cpu())
full_mask = probs.squeeze().cpu().numpy()
full_mask=(full_mask > 0.5)
result = mask_to_image(full_mask)
result.save("./result/Kvasir/lstm_4skip_150/{}".format(name[0]))