import torch

from Unet.model import UNET

from Utils import utils

import torchvision

from Preprocessing.preprocessing import Preprocessor

import cv2

import os

DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = UNET(in_channels=1, out_channels=1).to(DEVICE).float()

utils.load_checkpoint(torch.load('tmp/checkpoint.pth.tar', map_location=torch.device('cpu')), model)

model.eval()

prep = Preprocessor(128, 128)

for fileA in os.listdir('Dataset/zdim/valimagesbayes'):

    with torch.no_grad():

        x = cv2.imread('Dataset/zdim/valimagesbayes/'+fileA, 0)

        x = prep.preprocess(x)

        x = torch.from_numpy(x).float()

        x = x.unsqueeze(0).unsqueeze(0)

        x = x.to(DEVICE)

        preds = model(x)

        # preds[preds<0.5]=0

        # preds = torch.sigmoid(preds)

        torchvision.utils.save_image(preds, './resultsz/'+fileA)