Predicted U-Net mask is all filled with value 1

Viewed 22

I am trying to predict a mask from a Unet model. However, the prediction is all values of 1. My code is below. The mask is shown but is just blank.

#define augmentations 
inference_transform = A.Compose([
    A.Resize(256, 256, always_apply=True),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)), 
    ToTensorV2()
])

#define function for predictions
def predict(model, img, device):
    model.eval()
    with torch.no_grad():
        images = img.to(device)
        output = model(images)
        predicted_masks = (output.squeeze() >= 0.5).float().cpu().numpy()
        
    return(predicted_masks)

#define function to load image and output mask
def get_mask(img_path):
    image = cv2.imread(img_path)
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    original_height, original_width = tuple(image.shape[:2])
    
    image_trans = inference_transform(image = image)
    image_trans = image_trans["image"]
    image_trans = image_trans.unsqueeze(0)
    
    image_mask = predict(unet, image_trans, device)
    image_mask = cv2.resize(image_mask, (original_height, original_width), 
                          interpolation=cv2.INTER_NEAREST)
        
    return(image_mask)


#image example
example_path = '../input/test-image/10078.tiff'
image = cv2.imread(example_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

mask = get_mask(example_path)

masked_img = image*np.expand_dims(mask, 2).astype("uint8")

#plot the image, mask and multiplied together
fig, (ax1, ax2, ax3) = plt.subplots(3)

ax1.imshow(image)
ax2.imshow(mask)
ax3.imshow(masked_img)

Output:


np_img_tmp = np.uint8(Image.open('./' + list_masks))
plt.imshow(np_img_tmp)
print('min:%s, max: %s' % (np_img_tmp.min(), np_img_tmp.max()))
np_img_tmp

Mask output:

enter image description here

min:1, max: 1
array([[1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.],
       ...,
       [1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.],
       [1., 1., 1., ..., 1., 1., 1.]], dtype=float32)
0 Answers
Related