def __init__(self, model_name='u2net', cuda_mode=True, output_format='np'): self.model_name = model_name self.cuda_mode = cuda_mode and torch.cuda.is_available() # Fallback to CPU mode, if cuda is not available self.trans = transforms.Compose([RescaleT(320), ToTensorLab(flag=0)]) self.output_format = output_format # Validate if output_format not in self.FORMATS: raise AssertionError('Invalid "output_format"', 'Use "np" or "pil"') if model_name not in self.MODEL_NAMES: raise AssertionError('Invalid "model_name"', 'Use "u2net" or "u2netp"') if model_name == 'u2net': print("Model: U2NET (173.6 MB)") self.net = U2NET(3, 1) # 173.6 MB elif model_name == 'u2netp': print("Model: U2NetP (4.7 MB)") self.net = U2NETP(3, 1) # 4.7 MB else: raise AssertionError('Invalid "model_name"', 'Use "u2net" or "u2netp"') # Load network model_file = os.path.join(os.path.dirname(__file__), 'saved_models', model_name + '.pth') print("model_file:", model_file) if cuda_mode: print("CUDA mode") self.net.load_state_dict(torch.load(model_file)) self.net.cuda() else: print("CPU mode") self.net.load_state_dict(torch.load(model_file, map_location=torch.device('cpu'))) self.net.eval()
def main(): # --------- 1. get image path and name --------- 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' 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) # --------- 3. model define --------- if(model_name=='u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() 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]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1,d2,d3,d4,d5,d6,d7= net(inputs_test) # normalization pred = d1[:,0,:,:] pred = normPRED(pred) # save results to test_results folder save_output(img_name_list[i_test],pred,prediction_dir) del d1,d2,d3,d4,d5,d6,d7
def main(): model_name = 'u2net' # u2netp model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() cap = cv2.VideoCapture(0) import time while True: ret, frame = cap.read() if ret: t0 = time.time() img = cv2.resize(frame, (320, 320)) img = img.transpose((2, 0, 1)) img = img[None, ...] / 255. inputs_test = torch.from_numpy(img) inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) pred = d1[:, 0, :, :] predict = normalize(pred) predict = predict.squeeze() predict_np = predict.cpu().data.numpy() h, w = frame.shape[:2] pred_resized = cv2.resize(predict_np, (w, h)) img = (frame.astype(np.float32) * np.dstack( (pred_resized, pred_resized, pred_resized))).astype(np.uint8) cv2.imshow("out", img) cv2.waitKey(1) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'u2netp' # u2netp image_dir = '../train2014' prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir = '../models/' + model_name + '.pth' img_name_list = glob.glob(image_dir + os.sep + '*') # --------- 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) if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) # if torch.cuda.is_available(): # net.cuda() net.eval() all_out = {} for i_test, data_test in tqdm(enumerate(test_salobj_dataloader)): sep_ = img_name_list[i_test].split(os.sep)[-1] inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) # # if torch.cuda.is_available(): # inputs_test = Variable(inputs_test.cuda()) # else: inputs_test = Variable(inputs_test) d = net(inputs_test) pred = normPRED(d) all_out[sep_] = pred pickle.dump(all_out, open("../data/coco_train_u2net.pik", "wb"), protocol=2)
def model(model_name='u2net'): model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() return net
def model(model_name="u2net"): model_dir = os.path.join(os.getcwd(), "saved_models", model_name, model_name + ".pth") if model_name == "u2net": print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() return net
def __init__(self, model_name): self.model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') if model_name == 'u2net': print("...load U2NET---173.6 MB") net = U2NET(3, 1) elif model_name == 'u2netp': print("...load U2NEP---4.7 MB") net = U2NETP(3, 1) net.load_state_dict( torch.load(self.model_dir, map_location=torch.device('cpu'))) if torch.cuda.is_available(): net.cuda() net.eval() self.net = net
def main(): model_name='u2netp' model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') img_name_list = glob.glob(image_dir + os.sep + '*') 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) if(model_name=='u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() for data_test in test_salobj_dataloader: inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) dummy_input = inputs_test torch.onnx.export(net, dummy_input,"exported/onnx/{}.onnx".format(model_name), opset_version=10) break
def main(): # --------- 1. get image path and name --------- model_name = 'u2net' # u2netp image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') print(img_name_list) # --------- 2. dataloader --------- test_salobj_dataset = SalObjDataset( img_name_list=img_name_list, lbl_name_list=[], transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)])) if (model_name == 'u2net'): print("...load U2NET---173.6 MB") net = U2NET(3, 1) elif (model_name == 'u2netp'): print("...load U2NEP---4.7 MB") net = U2NETP(3, 1) for i_test, data_test in enumerate(len(img_name_list)): print("inferencing:", img_name_list[i_test].split(os.sep)[-1]) d1, d2, d3, d4, d5, d6, d7 = net.predict(test_salobj_dataset, steps=1) pred = d1[:, 0, :, :] pred = normPRED(pred) if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'u2netp' # u2net # image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') image_dir = '/home/hha/dataset/circle/circle' prediction_dir = '/home/hha/dataset/circle/circle_pred' model_dir = '/home/hha/pytorch_code/U-2-Net-master/saved_models/u2netp/u2netp.pthu2netp_bce_itr_2000_train_0.077763_tar_0.006976.pth' # model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') img_name_list = list(filter(lambda f: f.find('_mask') < 0, img_name_list)) # 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), # 320 ToTensorLab(flag=0) ])) test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=0) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() 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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = inputs_test.cuda() else: inputs_test = inputs_test d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] # pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
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
def run_model(): # --------- 1. get image path and name --------- model_name = 'u2net' image_dir = os.path.join(os.getcwd(), 'static', 'uploads') prediction_dir = os.path.join(os.getcwd(), 'static', "final" + os.sep) model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') #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) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") net = U2NET(3, 1) elif (model_name == 'u2netp'): print("...load U2NEP---4.7 MB") 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 --------- final_dict = {} for i_test, data_test in enumerate(test_salobj_dataloader): print("inferencing:", img_name_list[i_test].split(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) final_dict[img_name_list[i_test].split(os.sep)[-1]] = convert_img( img_name_list[i_test], pred, prediction_dir, img_name_list[i_test]) del d1, d2, d3, d4, d5, d6, d7 # returns dict mapping file name to [original image, nobg image, checker, whiteg, blackg] return final_dict
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 #shuffle=False , num_workers=0) # ------- 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 每迭代两千次存一次模型,这个可以在下面改成每个epoch存,很好改,这里我就不改了
def main(): # --------- 1. get image path and name --------- #model_name='u2net' model_name = 'u2netp' image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') 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) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() 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(os.sep)[-1]) inputs_test = data_test['image'] #print("test", inputs_test.shape) inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) #print("pred",pred.shape) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'u2net' #u2netp image_dir = "/home/vybt/Downloads/U2_Net_Test" prediction_dir = "/home/vybt/Downloads/u-2--bps-net-prediction" model_dir = '/media/vybt/DATA/SmartFashion/deep-learning-projects/U-2-Net/saved_models/u2net/_bps_bce_itr_300000_train_0.107041_tar_0.011690.pth' img_name_list = glob.glob(image_dir + os.sep + '*') # 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) # --------- 3. model define --------- if model_name == 'u2net': print("...load U2NET---173.6 MB") net = U2NET(3, 8) elif model_name == 'u2netp': print("...load U2NEP---4.7 MB") net = U2NETP(3, 8) net.load_state_dict(torch.load(model_dir)) if torch.cuda.is_available(): net.cuda() net.eval() # --------- 4. inference for each image --------- for i_test, data_test in enumerate(test_salobj_dataloader): print("Inference: ", img_name_list[i_test].split(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) # print("inputs test: {}".format(inputs_test.shape)) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1 pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): name = 'test' # please input the height unit:M body_height = 1.63 #"Input image" path image_dir = os.path.join(os.getcwd(), 'input') #"Output model" path outbody_filenames = './output/{}.obj'.format(name) ######################################################### #this code used for image segmentation to remove the background to get Silhouettes # --------- 1. get image path and name --------- model_name='u2net'#u2net or u2netp #set orignal silhouette images path prediction_dir1 = os.path.join(os.getcwd(), 'Silhouette' + os.sep) #set the path of silhouette images after horizontal flippath prediction_dir = os.path.join(os.getcwd(), 'test_data' + os.sep) model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') print(img_name_list) # --------- 2. 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. silhouette cutting model define --------- if(model_name=='u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() 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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1,d2,d3,d4,d5,d6,d7= net(inputs_test) # normalization pred = d1[:,0,:,:] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test],pred,prediction_dir,prediction_dir1) del d1,d2,d3,d4,d5,d6,d7 ########################################################### #this code used for reconstruct 3d model: #--------5.get the silhouette images -------- img_filenames = ['./test_data/front.png', './test_data/side.png'] # img = cv2.imread(img_filenames[1]) # cv2.flip(img,1) # -----------6.load input data--------- sampling_num = 648 data = np.zeros([2, 2, sampling_num]) for i in np.arange(len(img_filenames)): img = img_filenames[i] im = getBinaryimage(img, 600) # deal with white-black image simply sample_points = getSamplePoints(im, sampling_num, i) center_p = np.mean(sample_points, axis=0) sample_points = sample_points - center_p data[i, :, :] = sample_points.T data = repeat_data(data) #--------7 load CNN model----reconstruct 3d body shape print('==> begining...') len_out = 22 model_name = './Models/model.ckpt' ourModel = RegressionPCA(len_out) ourModel.load_state_dict(torch.load(model_name)) ourModel.eval() #----------8 output results-------------- save_obj(outbody_filenames, ourModel, body_height, data)
def main(): # --------- 1. get image path and name --------- model_name = 'u2netp' #u2netp image_dir = './data/workbench/' prediction_dir = './data/workbench_out/' model_dir = './saved_models/' + model_name + '/' + model_name + '.pth' img_name_list = glob.glob(image_dir + '*') print(img_name_list) # TODO: consider data loader over sets of videos # --------- 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)])) batch_size = 1 test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=batch_size, shuffle=False, num_workers=3) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() # --------- 4. inference for each image --------- from datetime import datetime a = datetime.now() total_inf = 0 for i_test, data_test in enumerate(test_salobj_dataloader): print("inferencing frame:", i_test * batch_size) # print("dl:", datetime.now()-a) a = datetime.now() inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = inputs_test.cuda() d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) print("inf:", total_inf / (i_test + 1)) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) total_inf += (datetime.now() - a).microseconds # save results to test_results folder # TODO: dynamically remember input sizes somehow, hardcoded for now for j in range(pred.shape[0]): save_output(img_name_list[batch_size * i_test + j], pred[j:j + 1], prediction_dir) del d1, d2, d3, d4, d5, d6, d7 a = datetime.now()
def main(): # --------- 1. get image path and name --------- model_name = 'u2net' # u2netp image_dir_0 = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir_0 = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir_0 = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'saved_models', model_name + '.pth') parser = argparse.ArgumentParser() parser.add_argument("-m", "--model_path", default=model_dir_0, help="pretrained model as .pth file") parser.add_argument("-i", "--input_path", default=image_dir_0, help="folder with input images") parser.add_argument("-o", "--output_path", default=prediction_dir_0, help="folder with output images") args = parser.parse_args() image_dir = args.input_path prediction_dir = args.output_path pretrained_model_path = args.model_path if not os.path.exists(pretrained_model_path): print( f"Could not find pretrained U^2 Net model at {pretrained_model_path}. " f"Please run CMake to download the default pretrained model or specify path to the pretrained model" f" as the -m argument to the u2net_test script.") return -1 img_name_list = glob.glob(image_dir + os.sep + '*') 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) # --------- 3. model define --------- if model_name == 'u2net': print("...load U2NET---173.6 MB") net = U2NET(3, 1) elif model_name == 'u2netp': print("...load U2NEP---4.7 MB") net = U2NETP(3, 1) else: raise ValueError(f"Unsuppored model name: {model_name}") if torch.cuda.is_available(): net.load_state_dict(torch.load(pretrained_model_path)) net.cuda() else: net.load_state_dict( torch.load(pretrained_model_path, 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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7 return 0
def main(): args = get_parameters() # --------- 1. get image path and name --------- model_name = 'u2net' # u2netp error_file_link = args.errorFile img_name_list = [] with open(args.input, 'r') as file: for line in file: line = line.strip() # preprocess line img_name_list.append(line) prediction_dir = args.output_dir model_dir = './saved_models/' + model_name + '.pth' #print(img_name_list) print("Num of image paths in ", str(args.input), "is: ", len(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, collate_fn=my_collate) # --------- 3. model define --------- if model_name == 'u2net': print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() # --------- 4. inference for each image --------- for i_test, data_test in enumerate(test_salobj_dataloader): try: print("\r------In processing file {} with name {}--------".format( i_test + 1, img_name_list[i_test].split("/")[-1]), end='') inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7 except Exception as error: print(error) with open(error_file_link, 'a+') as err_file: error_mess = img_name_list[i_test] + '*' + str(error) + '\n' err_file.write(error_mess) continue
def main(): # --------------------------------------------------------- # Configurations # --------------------------------------------------------- heavy_augmentation = True # False to use author's default implementation gan_training = False mixup_augmentation = False fullsize_training = False multiscale_training = False multi_gpu = True mixed_precision_training = True model_name = "u2net" # "u2net", "u2netp", "u2net_heavy" se_type = None # "csse", "sse", "cse", None; None to use author's default implementation # checkpoint = "saved_models/u2net/u2net.pth" checkpoint = None checkpoint_netD = None w_adv = 0.2 w_vgg = 0.2 train_dirs = [ "../datasets/sky_segmentation_dataset/datasets/cvprw2020_sky_seg/train/" ] train_dirs_file_limit = [ None, ] image_ext = '.jpg' label_ext = '.png' dataset_name = "cvprw2020_sky_seg" lr = 0.0003 epoch_num = 500 batch_size_train = 48 # batch_size_val = 1 workers = 16 save_frq = 1000 # save the model every 2000 iterations save_debug_samples = False debug_samples_dir = "./debug/" # --------------------------------------------------------- model_dir = './saved_models/' + model_name + '/' os.makedirs(model_dir, exist_ok=True) writer = SummaryWriter() if fullsize_training: batch_size_train = 1 multiscale_training = False # --------------------------------------------------------- # 1. Construct data input pipeline # --------------------------------------------------------- # Get dataset name dataset_name = dataset_name.replace(" ", "_") # Get training data assert len(train_dirs) == len(train_dirs_file_limit), \ "Different train dirs and train dirs file limit length!" tra_img_name_list = [] tra_lbl_name_list = [] for d, flimit in zip(train_dirs, train_dirs_file_limit): img_files = glob.glob(d + '**/*' + image_ext, recursive=True) if flimit: img_files = np.random.choice(img_files, size=flimit, replace=False) print(f"directory: {d}, files: {len(img_files)}") for img_path in img_files: lbl_path = img_path.replace("/image/", "/alpha/") \ .replace(image_ext, label_ext) if os.path.exists(img_path) and os.path.exists(lbl_path): assert os.path.splitext( os.path.basename(img_path))[0] == os.path.splitext( os.path.basename(lbl_path))[0], "Wrong filename." tra_img_name_list.append(img_path) tra_lbl_name_list.append(lbl_path) else: print( f"Warning, dropping sample {img_path} because label file {lbl_path} not found!" ) tra_img_name_list, tra_lbl_name_list = shuffle(tra_img_name_list, tra_lbl_name_list) train_num = len(tra_img_name_list) # val_num = 0 # unused print(f"dataset name : {dataset_name}") print(f"training samples : {train_num}") # Construct data input pipeline if heavy_augmentation: transform = AlbuSampleTransformer( get_heavy_transform( fullsize_training=fullsize_training, transform_size=False if (fullsize_training or multiscale_training) else True)) else: transform = transforms.Compose([ RescaleT(320), RandomCrop(288), ]) # Create dataset and dataloader dataset_kwargs = dict(img_name_list=tra_img_name_list, lbl_name_list=tra_lbl_name_list, transform=transforms.Compose([ transform, ] + ([ SaveDebugSamples(out_dir=debug_samples_dir), ] if save_debug_samples else []) + ([ ToTensorLab(flag=0), ] if not multiscale_training else []))) if mixup_augmentation: _dataset_cls = MixupAugSalObjDataset else: _dataset_cls = SalObjDataset salobj_dataset = _dataset_cls(**dataset_kwargs) salobj_dataloader = DataLoader( salobj_dataset, batch_size=batch_size_train, collate_fn=multi_scale_collater if multiscale_training else None, shuffle=True, pin_memory=True, num_workers=workers) # --------------------------------------------------------- # 2. Load model # --------------------------------------------------------- # Instantiate model if model_name == "u2net": net = U2NET(3, 1, se_type=se_type) elif model_name == "u2netp": net = U2NETP(3, 1, se_type=se_type) elif model_name == "u2net_heavy": net = u2net_heavy() elif model_name == "custom": net = CustomNet() else: raise ValueError(f"Unknown model_name: {model_name}") # Restore model weights from checkpoint if checkpoint: if not os.path.exists(checkpoint): raise FileNotFoundError(f"Checkpoint file not found: {checkpoint}") try: print(f"Restoring from checkpoint: {checkpoint}") net.load_state_dict(torch.load(checkpoint, map_location="cpu")) print(" - [x] success") except: print(" - [!] error") if torch.cuda.is_available(): net.cuda() if gan_training: netD = MultiScaleNLayerDiscriminator() if checkpoint_netD: if not os.path.exists(checkpoint_netD): raise FileNotFoundError( f"Discriminator checkpoint file not found: {checkpoint_netD}" ) try: print( f"Restoring discriminator from checkpoint: {checkpoint_netD}" ) netD.load_state_dict( torch.load(checkpoint_netD, map_location="cpu")) print(" - [x] success") except: print(" - [!] error") if torch.cuda.is_available(): netD.cuda() vgg19 = VGG19Features() vgg19.eval() if torch.cuda.is_available(): vgg19 = vgg19.cuda() # --------------------------------------------------------- # 3. Define optimizer # --------------------------------------------------------- optimizer = optim.Adam(net.parameters(), lr=lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0) # optimizer = optim.SGD(net.parameters(), lr=lr) # scheduler = optim.lr_scheduler.CyclicLR(optimizer, base_lr=lr/4, max_lr=lr, # mode="triangular2", # step_size_up=2 * len(salobj_dataloader)) if gan_training: optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(0.5, 0.9)) # --------------------------------------------------------- # 4. Initialize AMP and data parallel stuffs # --------------------------------------------------------- GOT_AMP = False if mixed_precision_training: try: print("Checking for Apex AMP support...") from apex import amp GOT_AMP = True print(" - [x] yes") except ImportError: print(" - [!] no") if GOT_AMP: amp.register_float_function(torch, 'sigmoid') net, optimizer = amp.initialize(net, optimizer, opt_level="O1") if gan_training: netD, optimizerD = amp.initialize(netD, optimizerD, opt_level="O1") vgg19 = amp.initialize(vgg19, opt_level="O1") if torch.cuda.device_count() > 1 and multi_gpu: print(f"Multi-GPU training using {torch.cuda.device_count()} GPUs.") net = nn.DataParallel(net) if gan_training: netD = nn.DataParallel(netD) vgg19 = nn.DataParallel(vgg19) else: print(f"Training using {torch.cuda.device_count()} GPUs.") # --------------------------------------------------------- # 5. Training # --------------------------------------------------------- print("Start training...") ite_num = 0 ite_num4val = 0 running_loss = 0.0 running_bce_loss = 0.0 running_tar_loss = 0.0 running_adv_loss = 0.0 running_per_loss = 0.0 running_fake_loss = 0.0 running_real_loss = 0.0 running_lossD = 0.0 for epoch in tqdm(range(0, epoch_num), desc="All epochs"): net.train() if gan_training: netD.train() for i, data in enumerate( tqdm(salobj_dataloader, desc=f"Epoch #{epoch}")): ite_num = ite_num + 1 ite_num4val = ite_num4val + 1 image_key = "image" label_key = "label" inputs, labels = data[image_key], data[label_key] # tqdm.write(f"input tensor shape: {inputs.shape}") 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) # # Zero the parameter gradients # optimizer.zero_grad() # Forward + backward + optimize d6 = 0 if model_name == "custom": d0, d1, d2, d3, d4, d5 = net(inputs_v) else: d0, d1, d2, d3, d4, d5, d6 = net(inputs_v) if gan_training: optimizerD.zero_grad() dis_fake = netD(inputs_v, d0.detach()) dis_real = netD(inputs_v, labels_v) loss_fake = bce_with_logits_loss(dis_fake, torch.zeros_like(dis_fake)) loss_real = bce_with_logits_loss(dis_real, torch.ones_like(dis_real)) lossD = loss_fake + loss_real if GOT_AMP: with amp.scale_loss(lossD, optimizerD) as scaled_loss: scaled_loss.backward() else: lossD.backward() optimizerD.step() writer.add_scalar("lossD/fake", loss_fake.item(), ite_num) writer.add_scalar("lossD/real", loss_real.item(), ite_num) writer.add_scalar("lossD/sum", lossD.item(), ite_num) running_fake_loss += loss_fake.item() running_real_loss += loss_real.item() running_lossD += lossD.item() # Zero the parameter gradients optimizer.zero_grad() if model_name == "custom": loss2, loss = multi_bce_loss_fusion5(d0, d1, d2, d3, d4, d5, labels_v) else: loss2, loss = multi_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v) writer.add_scalar("lossG/bce", loss.item(), ite_num) running_bce_loss += loss.item() if gan_training: # Adversarial loss loss_adv = 0.0 if w_adv: dis_fake = netD(inputs_v, d0) loss_adv = bce_with_logits_loss(dis_fake, torch.ones_like(dis_fake)) # Perceptual loss loss_per = 0.0 if w_vgg: vgg19_fm_pred = vgg19(inputs_v * d0) vgg19_fm_label = vgg19(inputs_v * labels_v) loss_per = mae_loss(vgg19_fm_pred, vgg19_fm_label) loss = loss + w_adv * loss_adv + w_vgg * loss_per if GOT_AMP: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() # scheduler.step() writer.add_scalar("lossG/sum", loss.item(), ite_num) writer.add_scalar("lossG/loss2", loss2.item(), ite_num) running_loss += loss.item() running_tar_loss += loss2.item() if gan_training: writer.add_scalar("lossG/adv", loss_adv.item(), ite_num) writer.add_scalar("lossG/perceptual", loss_per.item(), ite_num) running_adv_loss += loss_adv.item() running_per_loss += loss_per.item() if ite_num % 200 == 0: writer.add_images("inputs", inv_normalize(inputs_v), ite_num) writer.add_images("labels", labels_v, ite_num) writer.add_images("preds", d0, ite_num) # Delete temporary outputs and loss del d0, d1, d2, d3, d4, d5, d6, loss2, loss if gan_training: del dis_fake, dis_real, loss_fake, loss_real, lossD, loss_adv, vgg19_fm_pred, vgg19_fm_label, loss_per # Print stats tqdm.write( "[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train G/sum: %3f, G/bce: %3f, G/bce_tar: %3f, G/adv: %3f, G/percept: %3f, D/fake: %3f, D/real: %3f, D/sum: %3f" % (epoch + 1, epoch_num, (i + 1) * batch_size_train, train_num, ite_num, running_loss / ite_num4val, running_bce_loss / ite_num4val, running_tar_loss / ite_num4val, running_adv_loss / ite_num4val, running_per_loss / ite_num4val, running_fake_loss / ite_num4val, running_real_loss / ite_num4val, running_lossD / ite_num4val)) if ite_num % save_frq == 0: # Save checkpoint torch.save( net.module.state_dict() if hasattr( net, "module") else net.state_dict(), model_dir + model_name + (("_" + se_type) if se_type else "") + ("_" + dataset_name) + ("_mixup_aug" if mixup_augmentation else "") + ("_heavy_aug" if heavy_augmentation else "") + ("_fullsize" if fullsize_training else "") + ("_multiscale" if multiscale_training else "") + "_bce_itr_%d_train_%3f_tar_%3f.pth" % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val)) if gan_training: torch.save( netD.module.state_dict() if hasattr(netD, "module") else netD.state_dict(), model_dir + "netD_" + model_name + (("_" + se_type) if se_type else "") + ("_" + dataset_name) + ("_mixup_aug" if mixup_augmentation else "") + ("_heavy_aug" if heavy_augmentation else "") + ("_fullsize" if fullsize_training else "") + ("_multiscale" if multiscale_training else "") + "itr_%d.pth" % (ite_num)) # Reset stats running_loss = 0.0 running_bce_loss = 0.0 running_tar_loss = 0.0 running_adv_loss = 0.0 running_per_loss = 0.0 running_fake_loss = 0.0 running_real_loss = 0.0 running_lossD = 0.0 ite_num4val = 0 net.train() # resume train if gan_training: netD.train() writer.close() print("Training completed successfully.")
def main(): # --------- 1. get image path and name --------- model_name = 'u2net' #u2netp #模型名字和训练一致 # image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') image_dir = 'D:/wcs/U-2-Net/data/RIVER/Test/Image/' #测试图像路径 # prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) prediction_dir = 'D:/wcs/U-2-Net/data/RIVER/Test/pre/' #保存结果路径 model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') # 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) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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.load_state_dict( torch.load( './saved_models/u2net/u2net_bce_itr_36000_train_0.091362_tar_0.003286.pth' )) #加载自己的模型 if torch.cuda.is_available(): net.cuda() 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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization # print(d7.shape) pred = d1[:, 0, :, :] #注意这里,这个d1是融合了d2,d3,d4,d5,d6,d7的,如果想了解具体就翻到网络模型去自习看看 pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): full_sized_testing = False model_name = 'u2net' # model_name = 'u2netp' # image_dir = './test_images/' # image_dir = "../detectron2_mask_prediction/predictions/image/" image_dir = "../datasets/DUTS-TE/image/" # prediction_dir = './test_data/' + model_name + '_results/' # model_dir = './saved_models/' + model_name + '/' + model_name + '.pth' model_dir = "./saved_models/u2net/u2net.pth" # model_dir = "./u2net_mixed_person_n_portraits_heavy_aug_multiscale_bce_itr_8000_train_0.384033_tar_0.046964.pth" assert os.path.isfile(model_dir) prediction_dir = f"./predictions{'_fullsize' if full_sized_testing else ''}_{os.path.splitext(os.path.basename(model_dir))[0]}/" os.makedirs(prediction_dir, exist_ok=True) img_name_list = glob.glob(image_dir + '*') img_exts = [".jpg", ".jpeg", ".png", ".jfif"] img_name_list = list( filter(lambda p: os.path.splitext(p)[-1].lower() in img_exts, img_name_list)) # print(img_name_list) # --------- 2. dataloader --------- # 1. dataloader test_salobj_dataset = SalObjDataset( img_name_list=img_name_list, lbl_name_list=[], transform=transforms.Compose(([] if full_sized_testing else [ 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'): print("...load U2NET---173.6 MB") net = U2NET(3, 1, se_type=None) elif (model_name == 'u2netp'): print("...load U2NEP---4.7 MB") net = U2NETP(3, 1) elif (model_name == 'custom'): net = CustomNet() net.load_state_dict(torch.load(model_dir)) if torch.cuda.is_available(): net.cuda() 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]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d7 = 0 if model_name == "custom": d1, d2, d3, d4, d5, d6 = net(inputs_test) else: d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'u2netp' # u2netp u2net data_dir = '/data2/wangjiajie/datasets/scene_segment1023/u2data/' image_dir = os.path.join(data_dir, 'test_imgs') prediction_dir = os.path.join('./outputs/', model_name + '/') if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) # tra_label_dir = 'test_lbls/' image_ext = '.jpg' # label_ext = '.jpg' # '.png' model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') print(f'test img numbers are: {len(img_name_list)}') # --------- 2. dataloader --------- #1. dataloader test_salobj_dataset = SalObjDataset(img_name_list=img_name_list, lbl_name_list=[], transform=Compose([ SmallestMaxSize(max_size=320), ])) test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) checkpoint = torch.load(model_dir) d = collections.OrderedDict() for key, value in checkpoint.items(): tmp = key[7:] d[tmp] = value net.load_state_dict(d) if torch.cuda.is_available(): net.cuda() 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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = 1.0 - d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'u2net' #u2netp image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir = './saved_models/u2net_portrait/u2net_portrait.pth' #os.path.join(os.getcwd(), 'saved_models', model_name, 'u2net_bce_itr_8000_train_1.003468_tar_0.108501' + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') print(img_name_list) # --------- 2. dataloader --------- #1. dataloader transforms = A.Compose([ A.Resize(512, 512), A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ToTensorV2() ], p=1.0) test_salobj_dataset = SalObjDataset(img_name_list=img_name_list, lbl_name_list=[], transform=transforms) test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() 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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization #pred = 1.0 - d1[:,0,:,:]#reversed pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'u2net' #u2netp image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) # image_dir = os.path.join('/nfs/project/huxiaoliang/data/white_or_not/white_bg_image') # prediction_dir = os.path.join('/nfs/project/huxiaoliang/data/white_or_not/white_bg_image_pred'+ os.sep) model_dir = os.path.join('/nfs/private/modelfiles/u2net-saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') 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) # --------- 3. model define --------- if (model_name == 'u2net'): print("...load U2NET---173.6 MB") 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)) if torch.cuda.is_available(): net.cuda() net.eval() # --------- 4. inference for each image --------- error = [] for i_test, data_test in enumerate(test_salobj_dataloader): try: print("inferencing:", img_name_list[i_test].split(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7 except Exception as ex: traceback.print_exc() error.append(img_name_list[i_test].split(os.sep)[-1]) print('异常数据:', error)
saturation=(0.9, 1.1), hue=0.05), 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(4, 1) net.load_state_dict(torch.load(model_dir), strict=False) elif (model_name == 'u2netp'): net = U2NETP(4, 1) net.load_state_dict(torch.load(model_dir), strict=False) 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) # lr = 0.001 # ------- 5. training process -------- print("---start training...")
def main(): # --------- 1. get image path and name --------- model_name = 'u2netp' # fixed as u2netp image_dir = os.path.join( os.getcwd(), 'input' ) # changed to 'images' directory which is populated while running the script prediction_dir = os.path.join( os.getcwd(), 'output/' ) # changed to 'results' directory which is populated after the predictions model_dir = os.path.join(os.getcwd(), model_name + '.pth') # path to u2netp pretrained weights img_name_list = glob.glob(image_dir + os.sep + '*') 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) # --------- 3. model define --------- 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=torch.device('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(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir) del d1, d2, d3, d4, d5, d6, d7
def main(): # --------- 1. get image path and name --------- model_name = 'rgbd_u2net' #u2netp image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) # model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '_s6457_2021-02-18_12_36_53u2net_6457_bce_itr_38000_train_0.107382_tar_0.009159.pth') model_dir = 'saved_models/rgbd_u2net/rgbd_u2net_s30857_2021-04-06_03_52_47/rgbd_u2net_30857_bce_itr_220000_train_0.089360_tar_0.008674.pth' # path_files = Path('/pool/2021-03-31_22-11-41') path_files = Path('/pool/2021-03-31_22-32-41') img_name_list1 = sorted([str(x) for x in path_files.rglob('**/rgb/*.png')]) img_name_list1 = [ str(x) for x in img_name_list1 if '2021-03-31' in str(x) or '2021-04-01' in str(x) ] # path_files2 = Path('/dataset') # img_name_list2 = sorted([str(x) for x in path_files2.rglob('**/rgb/*.png') if not Path(str(x).replace('rgb','annotation')).exists()]) # img_name_list2 = [str(x) for x in img_name_list2 if '2021-03-10' in str(x)] img_name_list = img_name_list1 #+ img_name_list2 print('img len', len(img_name_list)) depth_name_list = [ x.replace('rgb', 'aligned_depth') for x in img_name_list ] # --------- 2. dataloader --------- #1. dataloader test_salobj_dataset = RGBD_SalObjDataset(img_name_list=img_name_list, depth_name_list=depth_name_list, lbl_name_list=[], transform=transforms.Compose([ RGBD_RescaleT(320), RGBD_ToTensorLab(flag=0) ])) test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1) # --------- 3. model define --------- if (model_name == 'rgbd_u2net'): print("...load U2NET---173.6 MB") net = U2NET(4, 1) elif (model_name == 'u2netp'): print("...load U2NEP---4.7 MB") net = U2NETP(3, 1) net.load_state_dict(torch.load(model_dir)) if torch.cuda.is_available(): net.cuda() net.eval() # --------- 4. inference for each image --------- for i_test, data_test in enumerate(tqdm(test_salobj_dataloader)): # print("inferencing:",img_name_list[i_test].split(os.sep)[-1]) inputs_test = data_test['image'] depth = data_test['depth'] inputs_test = torch.cat((inputs_test, torch.unsqueeze(depth, dim=1)), dim=1) # H x W x 4 inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = Variable(inputs_test.cuda()) else: inputs_test = Variable(inputs_test) with torch.no_grad(): d1, d2, d3, d4, d5, d6, d7 = net(inputs_test) # normalization pred = d1[:, 0, :, :] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test], pred, prediction_dir, i_test) del d1, d2, d3, d4, d5, d6, d7
def main(model_name, img_dir, retrain, weight, model_dir): model_name = 'u2net' #'u2netp' #data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep) # tra_image_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'im_aug' + os.sep) # tra_label_dir = os.path.join('DUTS', 'DUTS-TR', 'DUTS-TR', 'gt_aug' + os.sep) tra_image_dir = os.path.join(img_dir, 'origin') tra_label_dir = os.path.join(img_dir, 'mask') # train_image_dir = os.path.join('') image_ext = '.jpg' label_ext = '.png' model_dir = os.path.join(os.getcwd(), 'saved_models', model_name + os.sep) epoch_start = 0 epoch_num = 500 batch_size_train = 20 batch_size_val = 1 train_num = 4000 val_num = 500 # tra_img_name_list = glob.glob(tra_image_dir + '*' + image_ext) tra_img_name_list = os.listdir(tra_image_dir) for i,item in enumerate(tra_img_name_list): tra_img_name_list[i] = os.path.join(tra_image_dir, item) tra_lbl_name_list = os.listdir(tra_label_dir) for i,item in enumerate(tra_lbl_name_list): tra_lbl_name_list[i] = os.path.join(tra_label_dir, item) print(tra_img_name_list) # for img_path in tra_img_name_list: # img_name = img_path.split(os.sep)[-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(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=4) # ------- 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...") if retrain == True: checkpoint = torch.load(weight) net.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) epoch = checkpoint['epoch'] # loss = checkpoint['loss'] 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 # print(data) 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.item() running_tar_loss += loss2.data.item() # 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)) # # torch.save({ # # 'epoch': epoch, # # 'model_state_dict': net.state_dict(), # # 'optimizer_state_dict': optimizer.state_dict(), # # 'loss': loss, # # }, model_dir + model_name + epoch) # running_loss = 0.0 # running_tar_loss = 0.0 # net.train() # resume train # ite_num4val = 0 torch.save( { 'epoch': epoch, 'model_state_dict': net.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), }, os.path.join(model_dir, model_name + str(epoch)) )
def main(): # --------- 1. get image path and name --------- model_name='u2netp'#u2netp image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images') prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep) model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth') img_name_list = glob.glob(image_dir + os.sep + '*') # 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) # --------- 3. model define --------- if(model_name=='u2net'): print("...load U2NET---173.6 MB") net = U2NET(3,1) elif(model_name=='u2netp'): print("...load U2NEP---4.7 MB") net = U2NETP(3,1) if torch.cuda.is_available(): net.load_state_dict(torch.load("saved_models/u2netp/u2netp_bce_itr_10000_train_1.384799_tar_0.185377.pth")) net.cuda() else: net.load_state_dict(torch.load(model_dir, map_location=torch.device("cpu"))) net.eval() inference_time = 0.0 # --------- 4. inference for each image --------- for i_test, data_test in enumerate(test_salobj_dataloader): # print("inferencing:",img_name_list[i_test].split(os.sep)[-1]) inputs_test = data_test['image'] inputs_test = inputs_test.type(torch.FloatTensor) if torch.cuda.is_available(): inputs_test = inputs_test.cuda() start = time.time() with torch.no_grad(): d1,d2,d3,d4,d5,d6,d7= net(inputs_test) inference_time += (time.time() - start) # normalization pred = d1[:,0,:,:] pred = normPRED(pred) # save results to test_results folder if not os.path.exists(prediction_dir): os.makedirs(prediction_dir, exist_ok=True) save_output(img_name_list[i_test],pred,prediction_dir) del d1,d2,d3,d4,d5,d6,d7 print( f"Predicted {len(img_name_list)} images in {inference_time:.2f}s" )