import matplotlib.pyplot as plt

from models.segmentation_models import DeepLab

from data.hyperkvasir import KvasirSyntheticDataset

import torch

from torch.utils.data import DataLoader

def test():

    model = DeepLab().to("cuda")

    model.load_state_dict(torch.load("Predictors/Augmented/DeepLab/inpainter_augmentation_1"))

    for x, y, _ in DataLoader(KvasirSyntheticDataset("Datasets/HyperKvasir")):

        x = x.to("cuda")

        y = y.to("cuda")

        out = model.predict(x)

        plt.imshow(x[0].cpu().T)

        plt.axis("off")

        plt.show(bbox_inches='tight', pad_inches=0)

        plt.imshow(y[0].cpu().T, alpha=0.5)

        plt.show()

if __name__ == '__main__':

    test()