I am having trouble visualizing the predictions of a multi-class (4 classes) instance segmentation.
I can visualize the masks before the modeling like this but how could I do the same with the predictions?
CLASSES = ['sky',  'building', 'pole', 'road']

dataset = Dataset(x_train_dir, y_train_dir, classes= ['sky', 'building', 'pole', 'road'])

image, mask = dataset[4] # get some sample
visualize(
    image=image, 
    #cars_mask=mask.squeeze(),
    sky_mask = mask[...,0].squeeze(),
    cars_mask = mask[...,1] .squeeze(),
    pole_mask = mask[...,2].squeeze(),
    road_mask = mask[...,3] .squeeze(),
)

Example output:

Visualizing predictions
for i in range(5):
    n = np.random.choice(len(test_dataset))
    
    image_vis = test_dataset_vis[n][0].astype('uint8')
    image, gt_mask = test_dataset[n]
    
    gt_mask = gt_mask.squeeze()
    
    x_tensor = torch.from_numpy(image).to(DEVICE).unsqueeze(0)
    pr_mask = model.predict(x_tensor)
    pr_mask = (pr_mask.squeeze().cpu().numpy().round())
        
    visualize(
        image=image_vis, 
        ground_truth_mask=gt_mask, 
        predicted_mask=pr_mask
    )