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:
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)
