def main():
# --------- 1. get image path and name ---------
model_name = 'u2net_portrait' #u2netp
image_dir = './test_data/test_portrait_images/portrait_im'
prediction_dir = './test_data/test_portrait_images/portrait_results'
if (not os.path.exists(prediction_dir)):
os.mkdir(prediction_dir)
model_dir = './saved_models/u2net_portrait/u2net_portrait.pth'
img_name_list = glob.glob(image_dir + '/*')
print("Number of images: ", 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(512),
ToTensorLab(flag=0)]))
test_salobj_dataloader = DataLoader(test_salobj_dataset,
batch_size=1,
shuffle=False,
num_workers=1)
# --------- 3. model define ---------
print("...load U2NET---173.6 MB")
net = U2NET(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, :, :]
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 = './save_images/images/'
prediction_dir = './save_images/' + model_name + '_result/'
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():
image_dir = '/media/markytools/New Volume/Courses/EE298CompVis/finalproject/datasets/DUTS/DUTS-TE/DUTS-TE-Image/'
prediction_dir = './predictionout/'
model_dir = './saved_models/basnet_bsi_dataaugandarchi/basnet_bsi_epoch_207_itr_272000_train_9.345837_tar_1.082428.pth'
img_name_list = glob.glob(image_dir + '*.jpg')
# --------- 2. dataloader ---------
#1. dataload
test_salobj_dataset = SalObjDataset(
img_name_list=img_name_list,
lbl_name_list=[],
transform=transforms.Compose([RescaleT(256),
ToTensorLab(flag=0)]),
category="test")
test_salobj_dataloader = DataLoader(test_salobj_dataset,
batch_size=1,
shuffle=False,
num_workers=1)
# --------- 3. model define ---------
print("...load BASNet...")
net = BASNet(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,d8 = net(inputs_test)
d1, d2, d3, d4, d5, d6, d7, d8, d1_struct, d2_struct, d3_struct, d4_struct, d5_struct, d6_struct, d7_struct = 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, d8
def main(colored=False, imagepath=''):
# --------- 1. get image path and name ---------
model_name='u2net'#u2netp
prediction_dir = './test_data/' + model_name + '_results/'
model_dir = './saved_models/'+ model_name + '/' + model_name + '.pth'
# --------- 2. dataloader ---------
#1. dataloader
test_salobj_dataset = SalObjDataset(img_name_list = [imagepath],
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, map_location=torch.device('cpu')))
if torch.cuda.is_available():
net.cuda()
net.eval()
# --------- 4. inference for each image ---------
for _, data_test in enumerate(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)
d1,d2,d3,d4,d5,d6,d7= net(inputs_test)
# normalization
pred = d1[:,0,:,:]
pred = normPRED(pred)
del d1,d2,d3,d4,d5,d6,d7
# save results to test_results folder
return save_output(imagepath, pred, prediction_dir, colored=colored)
def main():
# --------- 1. get image path and name ---------
model_name = 'u2net'
cwd = Path(os.getcwd())
image_dir = cwd / 'test_data' / 'test_human_images'
prediction_dir = cwd / 'test_data' / 'test_human_images_results'
prediction_dir.mkdir(exist_ok=True)
model_dir = cwd / 'saved_models' / (model_name + '_human_seg') / (model_name + '_human_seg.pth')
img_name_list = list(image_dir.glob('*'))
print("Images in test:", 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)
# --------- 3. model define ---------
print("...load U2NET---173.6 MB")
net = U2NET(3, 1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.load_state_dict(torch.load(model_dir, map_location=device)).to(device)
net.eval()
# --------- 4. inference for each image ---------
for i_test, data_test in enumerate(test_salobj_dataloader):
image_path = img_name_list[i_test]
print("inferencing:", image_path.name)
inputs_test = data_test['image']
inputs_test = inputs_test.to(next(net.parameters()))
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
save_output(image_path, pred, prediction_dir)
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 infer(
net,
image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images'),
prediction_dir = os.path.join(os.getcwd(), 'test_data', 'u2net' + '_results')
):
img_name_list = glob.glob(image_dir + os.sep + '*')
prediction_dir = prediction_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)
# --------- 4. inference for each image ---------
for i_test, data_test in enumerate(test_salobj_dataloader):
print("Generating mask for:",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_images(img_name_list[i_test],pred,prediction_dir)
del d1,d2,d3,d4,d5,d6,d7
def run_prediction(files):
img_name_list = files
# --------- 2. dataloader ---------
#1. dataload
test_salobj_dataset = SalObjDataset(
img_name_list=img_name_list,
lbl_name_list=[],
transform=transforms.Compose([RescaleT(256),
ToTensorLab(flag=0)]))
test_salobj_dataloader = DataLoader(test_salobj_dataset,
batch_size=1,
shuffle=False,
num_workers=1)
# --------- 3. model define ---------
print("...load BASNet...")
net = BASNet(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, d8 = 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, d8
def remove_bg(images):
dataset = SalObjDataset(img_name_list=images,
lbl_name_list=[],
transform=transforms.Compose(
[RescaleT(320),
ToTensorLab(flag=0)]))
dataloader = DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=1)
net = U2NET(3, 1)
net.load_state_dict(torch.load(model_dir, map_location='cpu'))
net.eval()
outputs = []
for i, data in enumerate(dataloader):
inputs = data['image']
inputs = inputs.type(torch.FloatTensor)
inputs = Variable(inputs)
d1, d2, d3, d4, d5, d6, d7 = net(inputs)
pred = d1[:, 0, :, :]
pred = normPRED(pred)
filename = save_output(images[i], pred, output_dir)
outputs.append(filename)
img = cv2.imread(images[i])
mask = cv2.imread(filename, 0)
rgba = cv2.cvtColor(img, cv2.COLOR_RGB2RGBA)
rgba[:, :, 3] = mask
cv2.imwrite(filename, rgba)
del d1, d2, d3, d4, d5, d6, d7
return outputs
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 = '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 predict():
# ensure an image was properly uploaded to our endpoint
if flask.request.method == "POST":
if flask.request.files.get("image"):
file_image = flask.request.files["image"]
#image = Image.open(image)
# preprocess the image and prepare it for classification
salobjdataset = SalObjDataset(img_name_list=[file_image],
lbl_name_list=[],
transform=transforms.Compose([
RescaleT(256),
ToTensorLab(flag=0)
]))
saldataloader = DataLoader(salobjdataset,
batch_size=1,
shuffle=False,
num_workers=1)
for image in saldataloader:
input_image = image['image']
input_image = input_image.type(torch.FloatTensor)
# classify the input image and then initialize the list
# of predictions to return to the client
#img = img.type(torch.FloatTensor)
if torch.cuda.is_available():
input_image = Variable(input_image.cuda())
else:
input_image = Variable(input_image)
load_model()
d1, d2, d3, d4, d5, d6, d7, d8 = net(input_image)
pred = d1[:, 0, :, :]
pred = normPRED(pred)
predict = pred
predict = predict.squeeze()
predict_np = predict.cpu().data.numpy()
im = Image.fromarray(predict_np * 255).convert('RGB')
image = io.imread(file_image)
imo = im.resize((image.shape[1], image.shape[0]),
resample=Image.BILINEAR)
pb_np = np.array(imo)
in_img = cv2.imread(file_image)
in_img = cv2.cvtColor(in_img, cv2.COLOR_BGR2RGB)
result_img = cv2.imread(imo)
result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB)
new_img = np.where(result_img < [250, 250, 250], [254, 254, 254],
in_img)
new_img = new_img.astype('uint8')
new_img = cv2.cvtColor(new_img, cv2.COLOR_BGR2RGB)
byte_io = io.BytesIO()
new_img.save(byte_io, 'PNG')
byte_io.seek(0)
return flask.send_file(byte_io, mimetype='image/png')
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='u2net_portrait'#u2netp
image_dir = './test_data/test_portrait_images/your_portrait_im'
prediction_dir = './test_data/test_portrait_images/your_portrait_results'
if(not os.path.exists(prediction_dir)):
os.mkdir(prediction_dir)
model_dir = './saved_models/u2net_portrait/u2net_portrait.pth'
#model_dir = './saved_models/u2net/u2net_bce_itr_8000_train_1.003468_tar_0.108501.pth'
img_name_list = glob.glob(image_dir+'/*')
print("Number of images: ", len(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 ---------
print("...load U2NET---")
net = U2NET(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,:,:]#reverse black-white
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='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 = '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 train():
data_dir = '/media/markytools/New Volume/Courses/EE298CompVis/finalproject/datasets/'
test_image_dir = 'DUTS/DUTS-TE/DUTS-TE-Image/'
test_label_dir = 'DUTS/DUTS-TE/DUTS-TE-Mask/'
image_ext = '.jpg'
label_ext = '.png'
model_dir = "../saved_models/"
resume_train = True
resume_model_path = model_dir + "basnet-original.pth"
last_epoch = 1
epoch_num = 100000
batch_size_train = 8
batch_size_val = 1
train_num = 0
val_num = 0
enableInpaintAug = False
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") #set CPU to 0
# ------- 5. training process --------
print("---start training...")
test_increments = 15000
ite_num = 0
running_loss = 0.0
running_tar_loss = 0.0
ite_num4val = 1
next_test = ite_num + 0
############
############
############
############
test_img_name_list = glob.glob(data_dir + test_image_dir + '*' + image_ext)
print("data_dir + test_image_dir + '*' + image_ext: ", data_dir + test_image_dir + '*' + image_ext)
test_lbl_name_list = []
for img_path in test_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]
test_lbl_name_list.append(data_dir + test_label_dir + imidx + label_ext)
print("---")
print("test images: ", len(test_img_name_list))
print("test labels: ", len(test_lbl_name_list))
print("---")
test_num = len(test_img_name_list)
for test_lbl in test_lbl_name_list:
test_jpg = test_lbl.replace("png", "jpg")
test_jpg = test_jpg.replace("Mask", "Image")
if test_jpg not in test_img_name_list: print("test_lbl not in label: ", test_lbl)
salobj_dataset_test = SalObjDataset(
img_name_list=test_img_name_list,
lbl_name_list=test_lbl_name_list,
transform=transforms.Compose([
RescaleT(256),
RandomCrop(224),
ToTensorLab(flag=0)]),
category="test",
enableInpaintAug=enableInpaintAug)
salobj_dataloader_test = DataLoader(salobj_dataset_test, batch_size=batch_size_val, shuffle=True, num_workers=1)
# ------- 3. define model --------
# define the net
net = BASNet(3, 1)
if resume_train:
# print("resume_model_path:", resume_model_path)
checkpoint = torch.load(resume_model_path)
net.load_state_dict(checkpoint)
if torch.cuda.is_available():
net.to(device)
# ------- 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)
plotter = VisdomLinePlotter(env_name='NewlyAddedRelaxedMeasureEnv1')
best_ave_mae = 100000
best_max_fmeasure = 0
best_relaxed_fmeasure = 0
best_ave_maxf = 0
### Train network
train_loss0 = AverageMeter()
train_loss1 = AverageMeter()
train_loss2 = AverageMeter()
train_loss3 = AverageMeter()
train_loss4 = AverageMeter()
train_loss5 = AverageMeter()
train_loss6 = AverageMeter()
train_loss7 = AverageMeter()
test_loss0 = AverageMeter()
test_loss1 = AverageMeter()
test_loss2 = AverageMeter()
test_loss3 = AverageMeter()
test_loss4 = AverageMeter()
test_loss5 = AverageMeter()
test_loss6 = AverageMeter()
test_loss7 = AverageMeter()
average_mae = AverageMeter()
average_maxf = AverageMeter()
average_relaxedf = AverageMeter()
### Validate model
print("---Evaluate model---")
next_test = ite_num + test_increments
net.eval()
max_epoch_fmeasure = 0
for i, data in enumerate(salobj_dataloader_test):
inputs, labels = data['image'], data['label']
inputs = inputs.type(torch.FloatTensor)
labels = labels.type(torch.FloatTensor)
if torch.cuda.is_available():
inputs_v, labels_v = Variable(inputs.to(device), requires_grad=False), Variable(labels.to(device), requires_grad=False)
else:
inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)
d0, d1, d2, d3, d4, d5, d6, d7 = net(inputs_v)
pred = d0[:,0,:,:]
pred = normPRED(pred)
pred = pred.squeeze()
predict_np = pred.cpu().data.numpy()
im = Image.fromarray(predict_np*255).convert('RGB')
img_name = test_img_name_list[i]
image = cv2.imread(img_name)
imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR)
imo = imo.convert("L") ### Convert to grayscale 1-channel
resizedImg_np = np.array(imo) ### Result is 2D numpy array predicted salient map
img__lbl_name = test_lbl_name_list[i]
gt_img = np.array(Image.open(img__lbl_name).convert("L")) ### Ground truth salient map
### Compute metrics
img_name_png =
result_mae = getMAE(gt_img, resizedImg_np)
average_mae.update(result_mae, 1)
precision, recall = getPRCurve(gt_img, resizedImg_np)
result_maxfmeasure = getMaxFMeasure(precision, recall)
result_maxfmeasure = result_maxfmeasure.mean()
average_maxf.update(result_maxfmeasure, 1)
if (result_maxfmeasure > max_epoch_fmeasure):
max_epoch_fmeasure = result_maxfmeasure
result_relaxedfmeasure = getRelaxedFMeasure(gt_img, resizedImg_np)
average_relaxedf.update(result_relaxedfmeasure, 1)
loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, d7,labels_v,
test_loss0,
test_loss1,
test_loss2,
test_loss3,
test_loss4,
test_loss5,
test_loss6,
test_loss7)
del d0, d1, d2, d3, d4, d5, d6, d7,loss2, loss
print("Average Epoch MAE: ", average_mae.avg)
print("Max Max Epoch F-Measure: ", average_maxf.avg)
print("Average Max F-Measure: ", max_epoch_fmeasure)
print("Average Relaxed F-Measure: ", average_relaxedf.avg)
print('-------------Congratulations! Training Done!!!-------------')
imo.save(d_dir+imidx+'.png')
# --------- 1. get image path and name ---------
#image_dir = './test_data/test_images/'
image_dir = '../SOD_datasets/DUTS-TE/img/'
prediction_dir = './test_data/test_results/'
model_dir = './saved_models/basnet_bsi_original/basnet_bsi_itr_574000_train_1.013724_tar_0.054404.pth'
img_name_list = glob.glob(image_dir + '*.jpg')
# --------- 2. dataloader ---------
# 1. dataload
test_salobj_dataset = SalObjDataset(img_name_list=img_name_list, lbl_name_list=[
], transform=transforms.Compose([RescaleT(256), ToTensorLab(flag=0)]))
test_salobj_dataloader = DataLoader(
test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1)
# --------- 3. model define ---------
print("...load BASNet...")
net = BASNet(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])
running_loss = 0.0
scheduler.step()
t2 = time.time()
print(t2-t1)
print('-------------Congratulations! Training Done!!!-------------')
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--in_ch", type=int, default=3,help='rgb is 3,gray is 1')
parser.add_argument("--out_ch", type=int, default=64)
parser.add_argument('--device', type=str, default='cuda:0')
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--n_epochs', type=int, default=600, help='number of epochs to train')
parser.add_argument('--n_steps', type=int, default=200, help='number of epochs to update learning rate')
parser.add_argument('--gamma', type=float, default=0.1, help='')
parser.add_argument('--dataset_dir', type=str, default="DataSet/ori_data/DUTS-TR/trainData")
parser.add_argument('--saveModel_dir', type=str, default='experment_UFA/')
return parser.parse_args()
if __name__ == '__main__':
training_setup_seed(1)
args = parse_args()
transforms_ = [transforms.Resize((256,256), Image.BICUBIC),
# transforms.RandomHorizontalFlip(p=0.6),
transforms.ToTensor()]
train_set = SalObjDataset(dataset_dir=args.dataset_dir,transforms_=transforms_,rgb=True)
salobj_dataloader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4)
train(salobj_dataloader, args)
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 = '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)
def train():
data_dir = './train_data/'
tra_image_dir = 'DUTS-TR/DUTS-TR-Image/'
tra_label_dir = 'DUTS-TR/DUTS-TR-Mask/'
image_ext = '.jpg'
label_ext = '.png'
model_dir = "./saved_models/basnet_bsi/"
epoch_num = 100
batch_size_train = 1
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 = glob.glob(data_dir + tra_label_dir + '*' + label_ext)
# tra_lbl_name_list = []
# for img_path in tra_img_name_list:
# img_name = img_path.split("/")[-1]
# aaa = img_name.split(".")
# print(aaa)
# bbb = aaa[0:-1]
# imidx = bbb[0]
# for i in range(1,len(bbb)):
# imidx = imidx + "." + bbb[i]
# print(imidx)
# tra_lbl_name_list.append(data_dir + tra_label_dir + '*' + 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(256),
RandomCrop(224),
ToTensorLab(flag=0)
]))
salobj_dataloader = DataLoader(salobj_dataset,
batch_size=batch_size_train,
shuffle=True,
num_workers=1)
# ------- 3. define model --------
# define the net
net = BASNet(3, 1)
#if torch.cuda.is_available():
# print(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
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, d7 = net(inputs_v)
loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, d7,
labels_v)
loss.backward()
optimizer.step()
# # print statistics
running_loss += loss.data
running_tar_loss += loss2.data
# del temporary outputs and loss
del d0, d1, d2, d3, d4, d5, d6, d7, 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 % 2000 == 0: # save model every 2000 iterations
torch.save(
net.state_dict(),
model_dir + "basnet_bsi_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('-------------Congratulations! Training Done!!!-------------')
def setup(self, stage=None):
#TRAIN
if self.is_training:
print("TRAIN STAGE")
# Train Dataset Load
self.image_list_train = glob.glob(self.image_folder_train + '*' + self.image_ext)
self.label_list_train = []
for img_path in self.image_list_train:
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]
self.label_list_train.append(self.label_folder_train + imidx + self.label_ext)
# Val Dataset
self.image_list_val = glob.glob(self.image_folder_val + '*' + self.image_ext)
self.label_list_val = []
for img_path in self.image_list_val:
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]
self.label_list_val.append(self.label_folder_val + imidx + self.label_ext)
print("---")
print("train images: ", len(self.image_list_train))
print("train labels: ", len(self.label_list_train))
print("val images: ", len(self.image_list_val))
print("val labels: ", len(self.label_list_val))
print("train images: ", self.image_list_train[0])
print("train labels: ", self.label_list_train[0])
print("val images: ", self.image_list_val[0])
print("val labels: ", self.label_list_val[0])
print("---")
self.trainDataset = SalObjDataset(
img_name_list=self.image_list_train,
lbl_name_list=self.label_list_train,
transform=self.train_transform)
self.valDataset = SalObjDataset(
img_name_list=self.image_list_val,
lbl_name_list=self.label_list_val,
transform=self.val_transform)
#TEST
if not self.is_training:
print("TEST STAGE")
self.img_list_test = glob.glob(self.image_folder_test + '*')
# use same trasformation as in validation
self.testDataset = SalObjDataset(
img_name_list=self.img_list_test,
lbl_name_list=[],
transform=self.val_transform)
def train():
if os.name == 'nt':
data_dir = 'C:/Users/marky/Documents/Courses/saliency/datasets/DUTS/'
else:
data_dir = os.getenv(
"HOME") + '/Documents/Courses/EE298-CV/finalproj/datasets/DUTS/'
tra_image_dir = 'DUTS-TR/DUTS-TR-Image/'
tra_label_dir = 'DUTS-TR/DUTS-TR-Mask/'
test_image_dir = 'DUTS-TE/DUTS-TE-Image/'
test_label_dir = 'DUTS-TE/DUTS-TE-Mask/'
image_ext = '.jpg'
label_ext = '.png'
model_dir = "./saved_models/basnet_bsi_aug/"
resume_train = False
resume_model_path = model_dir + "basnet_bsi_epoch_81_itr_106839_train_1.511335_tar_0.098392.pth"
last_epoch = 1
epoch_num = 100000
batch_size_train = 8
batch_size_val = 1
train_num = 0
val_num = 0
enableInpaintAug = False
device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
# ------- 5. training process --------
print("---start training...")
test_increments = 6250
ite_num = 0
running_loss = 0.0
running_tar_loss = 0.0
ite_num4val = 1
next_test = ite_num + 0
visdom_tab_title = "StructArchWithoutStructImgs(WithHFlip)"
############
############
############
############
tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext)
print("data_dir + tra_image_dir + '*' + image_ext: ",
data_dir + tra_image_dir + '*' + image_ext)
test_img_name_list = glob.glob(data_dir + test_image_dir + '*' + image_ext)
print("data_dir + test_image_dir + '*' + image_ext: ",
data_dir + test_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)
test_lbl_name_list = []
for img_path in test_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]
test_lbl_name_list.append(data_dir + test_label_dir + imidx +
label_ext)
print("---")
print("train images: ", len(tra_img_name_list))
print("train labels: ", len(tra_lbl_name_list))
print("---")
print("---")
print("test images: ", len(test_img_name_list))
print("test labels: ", len(test_lbl_name_list))
print("---")
train_num = len(tra_img_name_list)
test_num = len(test_img_name_list)
salobj_dataset = SalObjDataset(img_name_list=tra_img_name_list,
lbl_name_list=tra_lbl_name_list,
transform=transforms.Compose([
RescaleT(256),
RandomCrop(224),
ToTensorLab(flag=0)
]),
category="train",
enableInpaintAug=enableInpaintAug)
salobj_dataset_test = SalObjDataset(img_name_list=test_img_name_list,
lbl_name_list=test_lbl_name_list,
transform=transforms.Compose([
RescaleT(256),
RandomCrop(224),
ToTensorLab(flag=0)
]),
category="test",
enableInpaintAug=enableInpaintAug)
salobj_dataloader = DataLoader(salobj_dataset,
batch_size=batch_size_train,
shuffle=True,
num_workers=1)
salobj_dataloader_test = DataLoader(salobj_dataset_test,
batch_size=batch_size_val,
shuffle=True,
num_workers=1)
# ------- 3. define model --------
# define the net
net = BASNet(3, 1)
if resume_train:
# print("resume_model_path:", resume_model_path)
checkpoint = torch.load(resume_model_path)
net.load_state_dict(checkpoint)
if torch.cuda.is_available():
net.to(device)
# ------- 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)
plotter = VisdomLinePlotter(env_name=visdom_tab_title)
best_ave_mae = 100000
best_max_fmeasure = 0
best_relaxed_fmeasure = 0
best_ave_maxf = 0
best_own_RelaxedFmeasure = 0
for epoch in range(last_epoch - 1, epoch_num):
### Train network
train_loss0 = AverageMeter()
train_loss1 = AverageMeter()
train_loss2 = AverageMeter()
train_loss3 = AverageMeter()
train_loss4 = AverageMeter()
train_loss5 = AverageMeter()
train_loss6 = AverageMeter()
train_loss7 = AverageMeter()
train_struct_loss1 = AverageMeter()
train_struct_loss2 = AverageMeter()
train_struct_loss3 = AverageMeter()
train_struct_loss4 = AverageMeter()
train_struct_loss5 = AverageMeter()
train_struct_loss6 = AverageMeter()
train_struct_loss7 = AverageMeter()
test_loss0 = AverageMeter()
test_loss1 = AverageMeter()
test_loss2 = AverageMeter()
test_loss3 = AverageMeter()
test_loss4 = AverageMeter()
test_loss5 = AverageMeter()
test_loss6 = AverageMeter()
test_loss7 = AverageMeter()
test_struct_loss1 = AverageMeter()
test_struct_loss2 = AverageMeter()
test_struct_loss3 = AverageMeter()
test_struct_loss4 = AverageMeter()
test_struct_loss5 = AverageMeter()
test_struct_loss6 = AverageMeter()
test_struct_loss7 = AverageMeter()
average_mae = AverageMeter()
average_maxf = AverageMeter()
average_relaxedf = AverageMeter()
average_own_RelaxedFMeasure = AverageMeter()
net.train()
for i, data in enumerate(salobj_dataloader):
ite_num = ite_num + 1
ite_num4val = ite_num4val + 1
inputs, labels, labels_struct = data['image'], data['label'], data[
'label2']
inputs = inputs.type(torch.FloatTensor)
labels = labels.type(torch.FloatTensor)
labels_struct = labels_struct.type(torch.FloatTensor)
# wrap them in Variable
if torch.cuda.is_available():
inputs_v, labels_v, labels_struct_v = Variable(
inputs.to(device), requires_grad=False), Variable(
labels.to(device), requires_grad=False), Variable(
labels_struct.to(device), requires_grad=False)
else:
inputs_v, labels_v, labels_struct_v = Variable(
inputs, requires_grad=False), Variable(
labels,
requires_grad=False), Variable(labels_struct,
requires_grad=False)
# y zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
d0, d1, d2, d3, d4, d5, d6, d7, d1_struct, d2_struct, d3_struct, d4_struct, d5_struct, d6_struct, d7_struct = net(
inputs_v)
loss2, loss = muti_bce_loss_fusion(
d0, d1, d2, d3, d4, d5, d6, d7, d1_struct, d2_struct,
d3_struct, d4_struct, d5_struct, d6_struct, d7_struct,
labels_v, train_loss0, train_loss1, train_loss2, train_loss3,
train_loss4, train_loss5, train_loss6, train_loss7,
train_struct_loss1, train_struct_loss2, train_struct_loss3,
train_struct_loss4, train_struct_loss5, train_struct_loss6,
train_struct_loss7)
loss.backward()
optimizer.step()
# # print statistics
running_loss += loss.data
running_tar_loss += loss2.data
# del temporary outputs and loss
del d0, d1, d2, d3, d4, d5, d6, d7, d1_struct, d2_struct, d3_struct, d4_struct, d5_struct, d6_struct, d7_struct, loss2, loss
print(
"[train 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))
plotter.plot('loss0', 'train', 'Main Loss 0', epoch + 1,
float(train_loss0.avg))
plotter.plot('loss1', 'train', 'Main Loss 1', epoch + 1,
float(train_loss1.avg))
plotter.plot('loss2', 'train', 'Main Loss 2', epoch + 1,
float(train_loss2.avg))
plotter.plot('loss3', 'train', 'Main Loss 3', epoch + 1,
float(train_loss3.avg))
plotter.plot('loss4', 'train', 'Main Loss 4', epoch + 1,
float(train_loss4.avg))
plotter.plot('loss5', 'train', 'Main Loss 5', epoch + 1,
float(train_loss5.avg))
plotter.plot('loss6', 'train', 'Main Loss 6', epoch + 1,
float(train_loss6.avg))
plotter.plot('loss7', 'train', 'Main Loss 7', epoch + 1,
float(train_loss7.avg))
plotter.plot('structloss1', 'train', 'Struct Loss 1', epoch + 1,
float(train_struct_loss1.avg))
plotter.plot('structloss2', 'train', 'Struct Loss 2', epoch + 1,
float(train_struct_loss2.avg))
plotter.plot('structloss3', 'train', 'Struct Loss 3', epoch + 1,
float(train_struct_loss3.avg))
plotter.plot('structloss4', 'train', 'Struct Loss 4', epoch + 1,
float(train_struct_loss4.avg))
plotter.plot('structloss5', 'train', 'Struct Loss 5', epoch + 1,
float(train_struct_loss5.avg))
plotter.plot('structloss6', 'train', 'Struct Loss 6', epoch + 1,
float(train_struct_loss6.avg))
plotter.plot('structloss7', 'train', 'Struct Loss 7', epoch + 1,
float(train_struct_loss7.avg))
### Validate model
print("---Evaluate model---")
if ite_num >= next_test: # test and save model 10000 iterations, due to very large DUTS-TE dataset
next_test = ite_num + test_increments
net.eval()
max_epoch_fmeasure = 0
for i, data in enumerate(salobj_dataloader_test):
inputs, labels = data['image'], data['label']
inputs = inputs.type(torch.FloatTensor)
labels = labels.type(torch.FloatTensor)
if torch.cuda.is_available():
inputs_v, labels_v = Variable(
inputs.to(device),
requires_grad=False), Variable(labels.to(device),
requires_grad=False)
else:
inputs_v, labels_v = Variable(
inputs,
requires_grad=False), Variable(labels,
requires_grad=False)
d0, d1, d2, d3, d4, d5, d6, d7, d1_struct, d2_struct, d3_struct, d4_struct, d5_struct, d6_struct, d7_struct = net(
inputs_v)
pred = d0[:, 0, :, :]
pred = normPRED(pred)
pred = pred.squeeze()
predict_np = pred.cpu().data.numpy()
im = Image.fromarray(predict_np * 255).convert('RGB')
img_name = test_img_name_list[i]
image = cv2.imread(img_name)
imo = im.resize((image.shape[1], image.shape[0]),
resample=Image.BILINEAR)
imo = imo.convert("L") ### Convert to grayscale 1-channel
resizedImg_np = np.array(
imo) ### Result is 2D numpy array predicted salient map
img__lbl_name = test_lbl_name_list[i]
gt_img = np.array(Image.open(img__lbl_name).convert(
"L")) ### Ground truth salient map
### Compute metrics
result_mae = getMAE(gt_img, resizedImg_np)
average_mae.update(result_mae, 1)
precision, recall = getPRCurve(gt_img, resizedImg_np)
result_maxfmeasure = getMaxFMeasure(precision, recall)
result_maxfmeasure = result_maxfmeasure.mean()
average_maxf.update(result_maxfmeasure, 1)
if (result_maxfmeasure > max_epoch_fmeasure):
max_epoch_fmeasure = result_maxfmeasure
result_relaxedfmeasure = getRelaxedFMeasure(
gt_img, resizedImg_np)
result_ownrelaxedfmeasure = own_RelaxedFMeasure(
gt_img, resizedImg_np)
average_relaxedf.update(result_relaxedfmeasure, 1)
average_own_RelaxedFMeasure.update(result_ownrelaxedfmeasure,
1)
loss2, loss = muti_bce_loss_fusion(
d0, d1, d2, d3, d4, d5, d6, d7, d1_struct, d2_struct,
d3_struct, d4_struct, d5_struct, d6_struct, d7_struct,
labels_v, test_loss0, test_loss1, test_loss2, test_loss3,
test_loss4, test_loss5, test_loss6, test_loss7,
test_struct_loss1, test_struct_loss2, test_struct_loss3,
test_struct_loss4, test_struct_loss5, test_struct_loss6,
test_struct_loss7)
del d0, d1, d2, d3, d4, d5, d6, d7, d1_struct, d2_struct, d3_struct, d4_struct, d5_struct, d6_struct, d7_struct, loss2, loss
print(
"[test epoch: %3d/%3d, batch: %5d/%5d, ite: %d] test loss: %3f, tar: %3f "
% (epoch + 1, epoch_num, (i + 1) * batch_size_val,
test_num, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val))
model_name = model_dir + "basnet_bsi_epoch_%d_itr_%d_train_%3f_tar_%3f.pth" % (
epoch + 1, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val)
torch.save(net.state_dict(), model_name)
running_loss = 0.0
running_tar_loss = 0.0
net.train() # resume train
ite_num4val = 1
if (average_mae.avg < best_ave_mae):
best_ave_mae = average_mae.avg
newname = model_dir + "bestMAE/basnet_bsi_epoch_%d_itr_%d_train_%3f_tar_%3f_mae_%3f.pth" % (
epoch + 1, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val, best_ave_mae)
fold_dir = newname.rsplit("/", 1)
if not os.path.isdir(fold_dir[0]): os.mkdir(fold_dir[0])
copyfile(model_name, newname)
if (max_epoch_fmeasure > best_max_fmeasure):
best_max_fmeasure = max_epoch_fmeasure
newname = model_dir + "bestEpochMaxF/basnet_bsi_epoch_%d_itr_%d_train_%3f_tar_%3f_maxfmeas_%3f.pth" % (
epoch + 1, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val, best_max_fmeasure)
fold_dir = newname.rsplit("/", 1)
if not os.path.isdir(fold_dir[0]): os.mkdir(fold_dir[0])
copyfile(model_name, newname)
if (average_maxf.avg > best_ave_maxf):
best_ave_maxf = average_maxf.avg
newname = model_dir + "bestAveMaxF/basnet_bsi_epoch_%d_itr_%d_train_%3f_tar_%3f_avemfmeas_%3f.pth" % (
epoch + 1, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val, best_ave_maxf)
fold_dir = newname.rsplit("/", 1)
if not os.path.isdir(fold_dir[0]): os.mkdir(fold_dir[0])
copyfile(model_name, newname)
if (average_relaxedf.avg > best_relaxed_fmeasure):
best_relaxed_fmeasure = average_relaxedf.avg
newname = model_dir + "bestAveRelaxF/basnet_bsi_epoch_%d_itr_%d_train_%3f_tar_%3f_averelaxfmeas_%3f.pth" % (
epoch + 1, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val, best_relaxed_fmeasure)
fold_dir = newname.rsplit("/", 1)
if not os.path.isdir(fold_dir[0]): os.mkdir(fold_dir[0])
copyfile(model_name, newname)
if (average_own_RelaxedFMeasure.avg > best_own_RelaxedFmeasure):
best_own_RelaxedFmeasure = average_own_RelaxedFMeasure.avg
newname = model_dir + "bestOwnRelaxedF/basnet_bsi_epoch_%d_itr_%d_train_%3f_tar_%3f_averelaxfmeas_%3f.pth" % (
epoch + 1, ite_num, running_loss / ite_num4val,
running_tar_loss / ite_num4val, best_own_RelaxedFmeasure)
fold_dir = newname.rsplit("/", 1)
if not os.path.isdir(fold_dir[0]): os.mkdir(fold_dir[0])
copyfile(model_name, newname)
plotter.plot('loss0', 'test', 'Main Loss 0', epoch + 1,
float(test_loss0.avg))
plotter.plot('loss1', 'test', 'Main Loss 1', epoch + 1,
float(test_loss1.avg))
plotter.plot('loss2', 'test', 'Main Loss 2', epoch + 1,
float(test_loss2.avg))
plotter.plot('loss3', 'test', 'Main Loss 3', epoch + 1,
float(test_loss3.avg))
plotter.plot('loss4', 'test', 'Main Loss 4', epoch + 1,
float(test_loss4.avg))
plotter.plot('loss5', 'test', 'Main Loss 5', epoch + 1,
float(test_loss5.avg))
plotter.plot('loss6', 'test', 'Main Loss 6', epoch + 1,
float(test_loss6.avg))
plotter.plot('loss7', 'test', 'Main Loss 7', epoch + 1,
float(test_loss7.avg))
plotter.plot('structloss1', 'test', 'Struct Loss 1', epoch + 1,
float(test_struct_loss1.avg))
plotter.plot('structloss2', 'test', 'Struct Loss 2', epoch + 1,
float(test_struct_loss2.avg))
plotter.plot('structloss3', 'test', 'Struct Loss 3', epoch + 1,
float(test_struct_loss3.avg))
plotter.plot('structloss4', 'test', 'Struct Loss 4', epoch + 1,
float(test_struct_loss4.avg))
plotter.plot('structloss5', 'test', 'Struct Loss 5', epoch + 1,
float(test_struct_loss5.avg))
plotter.plot('structloss6', 'test', 'Struct Loss 6', epoch + 1,
float(test_struct_loss6.avg))
plotter.plot('structloss7', 'test', 'Struct Loss 7', epoch + 1,
float(test_struct_loss7.avg))
plotter.plot('mae', 'test', 'Average Epoch MAE', epoch + 1,
float(average_mae.avg))
plotter.plot('max_maxf', 'test', 'Max Max Epoch F-Measure',
epoch + 1, float(max_epoch_fmeasure))
plotter.plot('ave_maxf', 'test', 'Average Max F-Measure',
epoch + 1, float(average_maxf.avg))
plotter.plot('ave_relaxedf', 'test', 'Average Relaxed F-Measure',
epoch + 1, float(average_relaxedf.avg))
plotter.plot('own_RelaxedFMeasure', 'test',
'Own Average Relaxed F-Measure', epoch + 1,
float(average_own_RelaxedFMeasure.avg))
print('-------------Congratulations! Training Done!!!-------------')
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 = '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()
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
#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)
#valid_size = round(len(tra_img_name_list) * 0.2)
valid_size = len(valid_img_name_list)
#train_size = len(tra_img_name_list) - valid_size
train_size = len(tra_img_name_list)
#print([train_size, valid_size])
#train_num = len(tra_img_name_list)
train_data = SalObjDataset(
img_name_list=tra_img_name_list,
lbl_name_list=tra_lbl_name_list,
transform=transforms.Compose([
RescaleT(256),
RandomCrop(224),
ToTensorLab(flag=0)]))
val_data = SalObjDataset(
img_name_list=valid_img_name_list,
lbl_name_list=valid_lbl_name_list,
transform=transforms.Compose([
RescaleT(256),
RandomCrop(224),
ToTensorLab(flag=0)]))
