import torch
from torch import optim
from torch.utils.data import DataLoader
from data import *
from net import *

from torchvision.utils import save_image

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
<<<<<<< HEAD
weight_path = r'D:\\MasterThesis\\UNet\\params\\unet.pth'
data_path = r'D:\\MasterThesis\\data\\VOCdevkit\\VOC2007'
save_path = r'D:\\MasterThesis\\UNet\\train_image'
=======
weight_path = r'/home/stud/hanzhang/MasterThesis/UNet/params/unet.pth'
data_path = r'/home/stud/hanzhang/MasterThesis/VOCdevkit/VOCdevkit/VOC2007'
save_path = r'/home/stud/hanzhang/MasterThesis/UNet/train_image'
>>>>>>> 37fdde8e83ce6de72d8d7226f22343e79b8a56d0

if __name__ == '__main__':
    data_loader = DataLoader(MyDataset(data_path), batch_size= 4, shuffle=True)

    net = UNet().to(device)
    if os.path.exists(weight_path):
        net.load_state_dict(torch.load(weight_path))
        print('successful load weight!')
    else:
        print('Failed on load weight!')

    opt = optim.Adam(net.parameters())
    loss_fun = nn.BCELoss()

    epoch=1

    while True:
        for i,(image,segment_image) in enumerate(data_loader):
            image, segment_image = image.to(device), segment_image.to(device)

            out_image = net(image)
            train_loss = loss_fun(out_image, segment_image)

            opt.zero_grad()
            train_loss.backward()
            opt.step() # 更新梯度

            if i%5 ==0 :
                print(f'{epoch} -- {i} -- train loss ===>> {train_loss.item()}')

            if i % 50 == 0:
                torch.save(net.state_dict(), weight_path)

            _image = image[0]
            _segment_image = segment_image[0]
            _out_image = out_image[0]

            img = torch.stack([_image, _segment_image, _out_image], dim=0)
            save_image(img, f'{save_path}/{i}.png')

        epoch += 1