0
点赞
收藏
分享

微信扫一扫

医学图像分割模型预测保存图片不是二值图

DT_M 2022-04-14 阅读 21
python

问题:今天跑完模型代码分割出来的图片不是二值图,而是如下问题,蒙上了一层雾气一样,不是我想要的二值图

 错误代码:

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]))
举报

相关推荐

0 条评论