from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

pred = "model/inference/monotemporalL2/export/epoch_1/test/img-0_pred.npy"
target = "model/inference/monotemporalL2/export/epoch_1/test/img-0_target.npy"
var = "model/inference/monotemporalL2/export/epoch_1/test/img-0_var.npy"

pred = np.load(pred)
target = np.load(target)
var = np.load(var)

print(pred.shape) # (13, 256, 256)
print(target.shape) # (13, 256, 256)
print(var.shape) # (0, 256, 256)

print(pred.dtype) # float32
print(target.dtype) # float32
print(var.dtype) # float32

print(pred.min(), pred.max()) # 0.0 1.0
print(target.min(), target.max()) # 0.0 1.0
# print(var.min(), var.max()) # nan nan

rgb = [3, 2, 1]

fig, ax = plt.subplots()
# get discrete colormap
cmap = plt.get_cmap("gray", 13)
ax.matshow(np.transpose(pred, (1, 2, 0)), cmap=cmap, vmin=0, vmax=1)
ax.axis("off")
fig.tight_layout()
plt.savefig("pred.png")