Visualise the test images after training the model on segmentation task

vision
3

I did this then I am pretty much lost what to do.

def predict_fundus_image(img_path,model):
    img = Image.open(img_path)
    
    transform = transforms.Compose([transforms.Resize((512,512)),
                                    transforms.ToTensor()])
    img_transformed = transform(img)
    img_transformed = transform(img).unsqueeze_(0)
    if train_on_gpu:
        img_transformed = img_transformed.cuda()
    if train_on_gpu:
        model = model.cuda()
    output = model(img_transformed)
    output = torch.argmax(output,1)
    return output


img = predict_fundus_image(img_path,model)
mapping = {(0, 0, 0): 0, (0, 0, 255): 1, (255, 0, 0): 2, (255, 255, 255): 3}
rev_map = {v: k for k, v in mapping.items()}