"""
读取图像的ground truth和每个round重建结果, 并且绘制error map.
"""
import os
import numpy as np
from PIL import Image
from skimage import io
from matplotlib import pyplot as plt


def normalize_image(image):
    # return (image - image.min())/(image.max() - image.min())
    image = image[40:200, 55:215]
    # image = image[80:160, 95:175]
    print("image shape:", image.shape)
    return image/255.0


def viz_diff_img(image, test_outputdir, image_name):
    print("image range:", image.max(), image.min())
    plt.imshow(image, cmap='jet')
    plt.savefig(os.path.join(test_outputdir, f'{image_name}'),
                bbox_inches='tight')


root_dir = "/data/xiaohan/BRATS_dataset/image_100patients_unimodal/"
image_name = "BraTS20_Training_042_60_t1"

dst_dir = "./recon_image_visualization"


img_gt = normalize_image(np.array(Image.open(root_dir + image_name + ".png")))
img_in = normalize_image(np.array(Image.open(root_dir + image_name + "_10dB.png")))

img_round1 = normalize_image(np.array(Image.open(root_dir + image_name + "_10dB_krecon_round1.png")))

print(img_gt.max(), img_gt.min())
print(img_in.max(), img_in.min())
print(img_round1.max(), img_round1.min())

io.imsave(os.path.join(dst_dir, image_name + ".png"), img_gt)
io.imsave(os.path.join(dst_dir, image_name + "_10dB.png"), img_in)
io.imsave(os.path.join(dst_dir, image_name + "_10dB_round1.png"), img_round1)

viz_diff_img(np.abs(img_gt - img_in)*255, dst_dir, image_name + "_input_error.png")
viz_diff_img(np.abs(img_gt - img_round1)*255, dst_dir, image_name + "_round1_error.png")

print("input error:", np.mean(np.abs(img_gt - img_in)))
print("round1 error:", np.mean(np.abs(img_gt - img_round1)))