from model.unet_model import UNet
from utils.dataset import ISBI_Loader
from torch import optim
import torch.nn as nn
import torch
from tqdm import tqdm
from PIL import Image
import os
import numpy as np



def train_net(net, device, data_path, epochs=50, batch_size=4, lr=0.00001): # 0.00001
    # 加載訓練集
    isbi_dataset = ISBI_Loader(data_path)
    per_epoch_num = len(isbi_dataset) / batch_size
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=batch_size,
                                               shuffle=True)
    # 定義RMSprop演算法
    #optimizer = torch.optim.Adam(net.parameters(), lr=lr, betas=(0.9, 0.99))
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
    # 定義Loss演算法 損失函數
    criterion = nn.BCEWithLogitsLoss()
    # best_loss統計，初始化為正無窮
    best_loss = float('inf')
    # 訓練epochs次
    with tqdm(total=epochs*per_epoch_num) as pbar:
        for epoch in range(epochs):
            # 訓練模式
            net.train()
            # 按照batch_size開始訓練
            for i,(image, label) in enumerate(train_loader):
                optimizer.zero_grad()
                image = image.to(device=device, dtype=torch.float32)
                label = label.to(device=device, dtype=torch.float32)
                pred = net(image)
                # 計算loss
                loss = criterion(pred, label)
                # print('{}/{}：Loss/train'.format(epoch + 1, epochs), loss.item())
                # 保存loss值最小的網絡參數
                if loss < best_loss:
                    best_loss = loss
                    torch.save(net.state_dict(), 'best_model_30_1.pth')
                # 更新參數
                loss.backward()
                optimizer.step()
                pbar.update(1)
             

if __name__ == "__main__":
    # 選擇設備，有cuda用cuda，沒有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加載網絡，圖片單通道1，分類為1
    net = UNet(n_channels=1, n_classes=1)  # todo 編輯 input_channels 和 n_classes
    # 將網絡拷貝到device中
    net.to(device=device)
    # 指定訓練集地址，開始訓練
    data_path = r"C:\Users\USER\Desktop\segment\train"# todo 修改為你本地的資料庫位置
    print("進度條出現卡住不是程式問題，是它正在計算，請耐心等待")
    train_net(net, device, data_path, epochs=50, batch_size=4)
