diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..9481906b --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__ +*/__pycache__ +**/__pycache__ +saved_models/ +.vscode \ No newline at end of file diff --git a/model/__pycache__/__init__.cpython-36.pyc b/model/__pycache__/__init__.cpython-36.pyc deleted file mode 100644 index ddc2097a..00000000 Binary files a/model/__pycache__/__init__.cpython-36.pyc and /dev/null differ diff --git a/model/__pycache__/u2net.cpython-36.pyc b/model/__pycache__/u2net.cpython-36.pyc deleted file mode 100644 index fae9df08..00000000 Binary files a/model/__pycache__/u2net.cpython-36.pyc and /dev/null differ diff --git a/u2net_test.py b/u2net_test.py index 8b501c6c..66424562 100644 --- a/u2net_test.py +++ b/u2net_test.py @@ -1,99 +1,107 @@ import os +import glob +import time + +import numpy as np +from PIL import Image from skimage import io, transform + import torch -import torchvision -from torch.autograd import Variable import torch.nn as nn import torch.nn.functional as F +from torch.autograd import Variable from torch.utils.data import Dataset, DataLoader -from torchvision import transforms#, utils -# import torch.optim as optim -import numpy as np -from PIL import Image -import glob +# import torch.optim as optim +import torchvision +from torchvision import transforms # , utils from data_loader import RescaleT from data_loader import ToTensor from data_loader import ToTensorLab from data_loader import SalObjDataset -from model import U2NET # full size version 173.6 MB -from model import U2NETP # small version u2net 4.7 MB +from model import U2NET # full size version 173.6 MB +from model import U2NETP # small version u2net 4.7 MB + # normalize the predicted SOD probability map def normPRED(d): ma = torch.max(d) mi = torch.min(d) - dn = (d-mi)/(ma-mi) + dn = (d - mi) / (ma - mi) return dn -def save_output(image_name,pred,d_dir): + +def save_output(image_name, pred, d_dir): predict = pred predict = predict.squeeze() predict_np = predict.cpu().data.numpy() - im = Image.fromarray(predict_np*255).convert('RGB') + im = Image.fromarray(predict_np * 255).convert("RGB") img_name = image_name.split("/")[-1] image = io.imread(image_name) - imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR) + imo = im.resize((image.shape[1], image.shape[0]), resample=Image.BILINEAR) pb_np = np.array(imo) aaa = img_name.split(".") bbb = aaa[0:-1] imidx = bbb[0] - for i in range(1,len(bbb)): + for i in range(1, len(bbb)): imidx = imidx + "." + bbb[i] - imo.save(d_dir+imidx+'.png') + imo.save(d_dir + imidx + ".png") + def main(): # --------- 1. get image path and name --------- - model_name='u2net'#u2netp + model_name = "u2net" # u2netp + image_dir = "./test_data/test_images/" + prediction_dir = "./test_data/" + model_name + "_results/" + model_dir = "./saved_models/" + model_name + "/" + model_name + ".pth" - image_dir = './test_data/test_images/' - prediction_dir = './test_data/' + model_name + '_results/' - model_dir = './saved_models/'+ model_name + '/' + model_name + '.pth' - - img_name_list = glob.glob(image_dir + '*') + img_name_list = glob.glob(image_dir + "*") print(img_name_list) # --------- 2. dataloader --------- - #1. dataloader - test_salobj_dataset = SalObjDataset(img_name_list = img_name_list, - lbl_name_list = [], - transform=transforms.Compose([RescaleT(320), - ToTensorLab(flag=0)]) - ) - test_salobj_dataloader = DataLoader(test_salobj_dataset, - batch_size=1, - shuffle=False, - num_workers=1) + # 1. dataloader + test_salobj_dataset = SalObjDataset( + img_name_list=img_name_list, + lbl_name_list=[], + transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)]), + ) + test_salobj_dataloader = DataLoader( + test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1 + ) # --------- 3. model define --------- - if(model_name=='u2net'): + if model_name == "u2net": print("...load U2NET---173.6 MB") - net = U2NET(3,1) - elif(model_name=='u2netp'): + net = U2NET(3, 1) + elif model_name == "u2netp": print("...load U2NEP---4.7 MB") - net = U2NETP(3,1) - net.load_state_dict(torch.load(model_dir)) + net = U2NETP(3, 1) + if torch.cuda.is_available(): + net.load_state_dict(torch.load(model_dir)) net.cuda() + else: + net.load_state_dict(torch.load(model_dir, map_location="cpu")) + net.eval() # --------- 4. inference for each image --------- for i_test, data_test in enumerate(test_salobj_dataloader): - print("inferencing:",img_name_list[i_test].split("/")[-1]) + start = time.time() - inputs_test = data_test['image'] + inputs_test = data_test["image"] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): @@ -101,16 +109,21 @@ def main(): else: inputs_test = Variable(inputs_test) - d1,d2,d3,d4,d5,d6,d7= net(inputs_test) + d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) + + print( + f"Predicted {os.path.basename(img_name_list[i_test])} in {time.time() - start:.2f}s" + ) # normalization - pred = d1[:,0,:,:] + pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder - save_output(img_name_list[i_test],pred,prediction_dir) + save_output(img_name_list[i_test], pred, prediction_dir) + + del d1, d2, d3, d4, d5, d6, d7 - del d1,d2,d3,d4,d5,d6,d7 if __name__ == "__main__": main() diff --git a/u2net_train.py b/u2net_train.py index 43d35a95..c91983d7 100644 --- a/u2net_train.py +++ b/u2net_train.py @@ -1,21 +1,22 @@ +import glob + import torch -import torchvision -from torch.autograd import Variable import torch.nn as nn +import torch.optim as optim import torch.nn.functional as F - +from torch.autograd import Variable from torch.utils.data import Dataset, DataLoader + +import torchvision from torchvision import transforms, utils -import torch.optim as optim import torchvision.transforms as standard_transforms import numpy as np -import glob from data_loader import Rescale from data_loader import RescaleT -from data_loader import RandomCrop from data_loader import ToTensor +from data_loader import RandomCrop from data_loader import ToTensorLab from data_loader import SalObjDataset @@ -26,139 +27,182 @@ bce_loss = nn.BCELoss(size_average=True) -def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v): - - loss0 = bce_loss(d0,labels_v) - loss1 = bce_loss(d1,labels_v) - loss2 = bce_loss(d2,labels_v) - loss3 = bce_loss(d3,labels_v) - loss4 = bce_loss(d4,labels_v) - loss5 = bce_loss(d5,labels_v) - loss6 = bce_loss(d6,labels_v) - - loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6 - print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"%(loss0.data[0],loss1.data[0],loss2.data[0],loss3.data[0],loss4.data[0],loss5.data[0],loss6.data[0])) - - return loss0, loss - - -# ------- 2. set the directory of training dataset -------- - -model_name = 'u2net' #'u2netp' - -data_dir = './train_data/' -tra_image_dir = 'DUTS/DUTS-TR/DUTS-TR/im_aug/' -tra_label_dir = 'DUTS/DUTS-TR/DUTS-TR/gt_aug/' - -image_ext = '.jpg' -label_ext = '.png' -model_dir = './saved_models/' + model_name +'/' - -epoch_num = 100000 -batch_size_train = 12 -batch_size_val = 1 -train_num = 0 -val_num = 0 - -tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext) - -tra_lbl_name_list = [] -for img_path in tra_img_name_list: - img_name = img_path.split("/")[-1] - - aaa = img_name.split(".") - bbb = aaa[0:-1] - imidx = bbb[0] - for i in range(1,len(bbb)): - imidx = imidx + "." + bbb[i] - - tra_lbl_name_list.append(data_dir + tra_label_dir + imidx + label_ext) - -print("---") -print("train images: ", len(tra_img_name_list)) -print("train labels: ", len(tra_lbl_name_list)) -print("---") - -train_num = len(tra_img_name_list) - -salobj_dataset = SalObjDataset( - img_name_list=tra_img_name_list, - lbl_name_list=tra_lbl_name_list, - transform=transforms.Compose([ - RescaleT(320), - RandomCrop(288), - ToTensorLab(flag=0)])) -salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=1) - -# ------- 3. define model -------- -# define the net -if(model_name=='u2net'): - net = U2NET(3, 1) -elif(model_name=='u2netp'): - net = U2NETP(3,1) - -if torch.cuda.is_available(): - net.cuda() - -# ------- 4. define optimizer -------- -print("---define optimizer...") -optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) - -# ------- 5. training process -------- -print("---start training...") -ite_num = 0 -running_loss = 0.0 -running_tar_loss = 0.0 -ite_num4val = 0 -save_frq = 2000 # save the model every 2000 iterations - -for epoch in range(0, epoch_num): - net.train() - - for i, data in enumerate(salobj_dataloader): - ite_num = ite_num + 1 - ite_num4val = ite_num4val + 1 - - inputs, labels = data['image'], data['label'] - - inputs = inputs.type(torch.FloatTensor) - labels = labels.type(torch.FloatTensor) - - # wrap them in Variable - if torch.cuda.is_available(): - inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(), - requires_grad=False) - else: - inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False) - - # y zero the parameter gradients - optimizer.zero_grad() - - # forward + backward + optimize - d0, d1, d2, d3, d4, d5, d6 = net(inputs_v) - loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v) - - loss.backward() - optimizer.step() - - # # print statistics - running_loss += loss.data[0] - running_tar_loss += loss2.data[0] +def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v): - # del temporary outputs and loss - del d0, d1, d2, d3, d4, d5, d6, loss2, loss + loss0 = bce_loss(d0, labels_v) + loss1 = bce_loss(d1, labels_v) + loss2 = bce_loss(d2, labels_v) + loss3 = bce_loss(d3, labels_v) + loss4 = bce_loss(d4, labels_v) + loss5 = bce_loss(d5, labels_v) + loss6 = bce_loss(d6, labels_v) + + loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6 + print( + "l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n" + % ( + loss0.data[0], + loss1.data[0], + loss2.data[0], + loss3.data[0], + loss4.data[0], + loss5.data[0], + loss6.data[0], + ) + ) + + return loss0, loss + + +def main(): + + # ------- 2. set the directory of training dataset -------- + model_name = "u2net" #'u2netp' + + data_dir = "./train_data/" + tra_image_dir = "DUTS/DUTS-TR/DUTS-TR/im_aug/" + tra_label_dir = "DUTS/DUTS-TR/DUTS-TR/gt_aug/" + + image_ext = ".jpg" + label_ext = ".png" + + model_dir = "./saved_models/" + model_name + "/" + + epoch_num = 100000 + batch_size_train = 12 + batch_size_val = 1 + train_num = 0 + val_num = 0 + + tra_img_name_list = glob.glob(data_dir + tra_image_dir + "*" + image_ext) + + tra_lbl_name_list = [] + for img_path in tra_img_name_list: + img_name = img_path.split("/")[-1] + + aaa = img_name.split(".") + bbb = aaa[0:-1] + imidx = bbb[0] + for i in range(1, len(bbb)): + imidx = imidx + "." + bbb[i] + + tra_lbl_name_list.append(data_dir + tra_label_dir + imidx + label_ext) + + print("---") + print("train images: ", len(tra_img_name_list)) + print("train labels: ", len(tra_lbl_name_list)) + print("---") + + train_num = len(tra_img_name_list) + + salobj_dataset = SalObjDataset( + img_name_list=tra_img_name_list, + lbl_name_list=tra_lbl_name_list, + transform=transforms.Compose( + [RescaleT(320), RandomCrop(288), ToTensorLab(flag=0)] + ), + ) + salobj_dataloader = DataLoader( + salobj_dataset, batch_size=batch_size_train, shuffle=True, num_workers=1 + ) + + # ------- 3. define model -------- + # define the net + if model_name == "u2net": + net = U2NET(3, 1) + elif model_name == "u2netp": + net = U2NETP(3, 1) + + if torch.cuda.is_available(): + net.cuda() + + # ------- 4. define optimizer -------- + print("---define optimizer...") + optimizer = optim.Adam( + net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0 + ) + + # ------- 5. training process -------- + print("---start training...") + ite_num = 0 + running_loss = 0.0 + running_tar_loss = 0.0 + ite_num4val = 0 + save_frq = 2000 # save the model every 2000 iterations + + for epoch in range(0, epoch_num): + net.train() + + for i, data in enumerate(salobj_dataloader): + ite_num = ite_num + 1 + ite_num4val = ite_num4val + 1 + + inputs, labels = data["image"], data["label"] + + inputs = inputs.type(torch.FloatTensor) + labels = labels.type(torch.FloatTensor) + + # wrap them in Variable + if torch.cuda.is_available(): + inputs_v, labels_v = ( + Variable(inputs.cuda(), requires_grad=False), + Variable(labels.cuda(), requires_grad=False), + ) + else: + inputs_v, labels_v = ( + Variable(inputs, requires_grad=False), + Variable(labels, requires_grad=False), + ) + + # y zero the parameter gradients + optimizer.zero_grad() + + # forward + backward + optimize + d0, d1, d2, d3, d4, d5, d6 = net(inputs_v) + loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v) + + loss.backward() + optimizer.step() + + # # print statistics + running_loss += loss.data[0] + running_tar_loss += loss2.data[0] + + # del temporary outputs and loss + del d0, d1, d2, d3, d4, d5, d6, loss2, loss + + print( + "[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f " + % ( + epoch + 1, + epoch_num, + (i + 1) * batch_size_train, + train_num, + ite_num, + running_loss / ite_num4val, + running_tar_loss / ite_num4val, + ) + ) + + if ite_num % save_frq == 0: + + torch.save( + net.state_dict(), + model_dir + + model_name + + "_bce_itr_%d_train_%3f_tar_%3f.pth" + % ( + ite_num, + running_loss / ite_num4val, + running_tar_loss / ite_num4val, + ), + ) + running_loss = 0.0 + running_tar_loss = 0.0 + net.train() # resume train + ite_num4val = 0 - print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f " % ( - epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val)) - - if ite_num % save_frq == 0: - - torch.save(net.state_dict(), model_dir + model_name+"_bce_itr_%d_train_%3f_tar_%3f.pth" % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val)) - running_loss = 0.0 - running_tar_loss = 0.0 - net.train() # resume train - ite_num4val = 0 if __name__ == "__main__": main()