def main():
check_dir = './' + name
if not os.path.exists(check_dir):
os.mkdir(check_dir)
# data
train_loader = torch.utils.data.DataLoader(Folder(
opt.train_dir,
transform=transforms.Compose([transforms.Resize(
(img_size, img_size))]),
mean=mean,
std=std),
batch_size=opt.b,
shuffle=True,
num_workers=4,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(Folder(
opt.val_dir,
transform=transforms.Compose([transforms.Resize(
(img_size, img_size))]),
mean=mean,
std=std),
batch_size=opt.b / 2,
shuffle=True,
num_workers=4,
pin_memory=True)
# models
criterion = nn.BCEWithLogitsLoss().cuda()
if opt.model == 'FCN':
net = models.FCN(pretrained=True, c_output=n_classes,
base=opt.base).cuda()
else:
# net = getattr(models, opt.model)(pretrain=True, c_output=n_classes).cuda()
net = getattr(models, opt.model)(c_output=n_classes).cuda()
optimizer = torch.optim.Adam([{'params': net.parameters(), 'lr': 1e-4}])
lr = 5e-3
lr_decay = 0.9
# optimizer = torch.optim.SGD([
# {'params': [param for _name, param in net.named_parameters() if _name[-4:] == 'bias'],
# 'lr': 2 * lr},
# {'params': [param for _name, param in net.named_parameters() if _name[-4:] != 'bias'],
# 'lr': lr, 'weight_decay': 1e-4}
# ], momentum=0.9, nesterov=True)
logs = {'best_ep': 0, 'best': 0}
for epoch in range(opt.e):
optimizer.param_groups[0]['lr'] = lr * (
1 - float(epoch) / opt.e)**lr_decay # weight
train(epoch, train_loader, optimizer, criterion, net, logs)
miou = validate(val_loader, net, os.path.join(check_dir, 'results'))
logs[epoch] = {'mIOU': miou}
if miou > logs['best']:
logs['best'] = miou
logs['best_ep'] = epoch
torch.save(net.state_dict(), '%s/net.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
import json
import os
import random
from options.train_options import TrainOptions
opt = TrainOptions() # set CUDA_VISIBLE_DEVICES before import torch
opt = opt.parse()
home = os.path.expanduser("~")
val_loader = torch.utils.data.DataLoader(Folder(opt.val_img_dir,
opt.val_gt_dir,
crop=None,
flip=False,
rotate=None,
size=opt.imageSize,
mean=opt.mean,
std=opt.std,
training=False),
batch_size=opt.batchSize,
shuffle=True,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_img_dir,
'./zhang-dut-train-crf_bin',
crop=0.9,
flip=True,
rotate=None,
size=opt.imageSize,
mean=opt.mean,
def main():
img_size = 256
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
make_dir(opt.result_dir)
# data
# Load vocabulary wrapper
with open('vocab.pkl', 'rb') as f:
vocab = pickle.load(f)
loader = torch.utils.data.DataLoader(
Folder(img_dir=opt.img_dir, gt_dir=opt.gt_dir,
source_transform=transforms.Compose([transforms.Resize((img_size, img_size))]),
target_transform=transforms.Compose([transforms.Resize((img_size, img_size))]),
mean=mean, std=std),
batch_size=opt.batchSize, shuffle=False, num_workers=4, pin_memory=True)
# caption and classification networks
cls_net = FCN(base='densenet169')
cls_net = cls_net.cuda()
cap_net = EncDec(len(vocab))
cap_net = cap_net.cuda()
# saliency network
sal_net = DeepLab(base='densenet169', c_output=1)
sal_net = nn.DataParallel(sal_net).cuda()
# the 1st, 2nd and 3rd rows of Table 1
cls_net.load_state_dict(torch.load('net-cls-init.pth'))
cap_net.load_state_dict(torch.load('net-cap-init.pth'))
output_dir = '/'.join([opt.result_dir, 'init', 'cls'])
make_dir(output_dir)
validate_one(loader, cls_net, output_dir)
fm, mae, _,_ = fm_and_mae(output_dir, opt.gt_dir)
print('cls fm %.3f'%fm)
# the 2nd row of Table 1
output_dir = '/'.join([opt.result_dir, 'init', 'cap'])
make_dir(output_dir)
validate_one(loader, cap_net, output_dir)
fm, mae, _,_ = fm_and_mae(output_dir, opt.gt_dir)
print('cap fm %.3f'%fm)
# the 3rd row of Table 1
output_dir = '/'.join([opt.result_dir, 'init', 'avg'])
make_dir(output_dir)
validate_two(loader, cls_net, cap_net, output_dir)
fm, mae, _,_ = fm_and_mae(output_dir, opt.gt_dir)
print('cls cap fm %.3f'%fm)
# the 4th row of Table 1
cls_net.load_state_dict(torch.load('cls-two-woun.pth'))
cap_net.load_state_dict(torch.load('cap-two-woun.pth'))
output_dir = '/'.join([opt.result_dir, 'at', 'avg'])
make_dir(output_dir)
validate_two(loader, cls_net, cap_net, output_dir)
fm, mae, _,_ = fm_and_mae(output_dir, opt.gt_dir)
print('cls cap at fm %.3f'%fm)
# the 5th row of Table 1
cls_net.load_state_dict(torch.load('cls-two-mr.pth'))
cap_net.load_state_dict(torch.load('cap-two-mr.pth'))
output_dir = '/'.join([opt.result_dir, 'ac', 'avg'])
make_dir(output_dir)
validate_two(loader, cls_net, cap_net, output_dir)
fm, mae, _,_ = fm_and_mae(output_dir, opt.gt_dir)
print('cls cap at ac fm %.3f'%fm)
# the 6th row of Table 1
sal_net.load_state_dict(torch.load('sal.pth'))
output_dir = '/'.join([opt.result_dir, 'sal'])
make_dir(output_dir)
validate_one(loader, sal_net, output_dir)
fm, mae, _,_ = fm_and_mae(output_dir, opt.gt_dir)
print('sal fm %.3f'%fm)
def main():
check_dir = '../LPSfiles/' + name
if not os.path.exists(check_dir):
os.mkdir(check_dir)
# data
val_loader = torch.utils.data.DataLoader(PriorFolder(opt.val_dir,
opt.prior_dir,
size=256,
mean=mean,
std=std),
batch_size=opt.b * 3,
shuffle=False,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_dir,
scales=[64] * 3 +
[128, 256],
crop=0.9,
flip=True,
rotate=None,
mean=mean,
std=std),
collate_fn=collate_more,
batch_size=opt.b * 6,
shuffle=True,
num_workers=4,
pin_memory=True)
# models
p = 5
net = Net(base=opt.base)
fcn = FCN(net)
net = nn.DataParallel(net).cuda()
net.train()
"""
# fcn = nn.DataParallel(fcn).cuda()
# sdict =torch.load('/home/crow/LPSfiles/Train2_vgg16/fcn-iter13800.pth')
# fcn.load_state_dict(sdict)
fcn = nn.DataParallel(fcn).cuda()
fcn.train()
optimizer = torch.optim.Adam([
{'params': fcn.parameters(), 'lr': 1e-4},
])
logs = {'best_it':0, 'best': 0}
sal_data_iter = iter(train_loader)
i_sal_data = 0
for it in tqdm(range(opt.max)):
# for it in tqdm(range(1)):
# if it > 1000 and it % 100 == 0:
# optimizer.param_groups[0]['lr'] *= 0.5
if i_sal_data >= len(train_loader):
sal_data_iter = iter(train_loader)
i_sal_data = 0
data, lbls, _ = sal_data_iter.next()
i_sal_data += 1
data = data.cuda()
lbls = [lbl.unsqueeze(1).cuda() for lbl in lbls]
msks = fcn(data)
loss = sum([F.binary_cross_entropy_with_logits(msk, lbl) for msk, lbl in zip(msks, lbls)])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if it % 10 == 0:
writer.add_scalar('loss', loss.item(), it)
image = make_image_grid(data[:6], mean, std)
writer.add_image('Image', torchvision.utils.make_grid(image), it)
big_msk = F.sigmoid(msks[-1]).expand(-1, 3, -1, -1)
writer.add_image('msk', torchvision.utils.make_grid(big_msk.data[:6]), it)
big_msk = lbls[-1].expand(-1, 3, -1, -1)
writer.add_image('gt', torchvision.utils.make_grid(big_msk.data[:6]), it)
# if it % 100 == 0:
if it != 0 and it % 100 == 0:
fm, mae = validate(val_loader, fcn, os.path.join(check_dir, 'results'),
os.path.join(opt.val_dir, 'masks'))
print(u'损失: %.4f'%(loss.item()))
print(u'最大FM: iteration %d的%.4f, 这次FM: %.4f'%(logs['best_it'], logs['best'], fm))
logs[it] = {'FM': fm}
if fm > logs['best']:
logs['best'] = fm
logs['best_it'] = it
torch.save(fcn.state_dict(), '%s/fcn-best.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
torch.save(fcn.state_dict(), '%s/fcn-iter%d.pth' % (check_dir, it))
"""
###################################################################################################
val_loader = torch.utils.data.DataLoader(PriorFolder(opt.val_dir,
opt.prior_dir,
size=256,
mean=mean,
std=std),
batch_size=opt.b,
shuffle=False,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_dir,
scales=[256],
crop=0.9,
flip=True,
rotate=None,
mean=mean,
std=std),
collate_fn=collate_more,
batch_size=opt.b,
shuffle=True,
num_workers=4,
pin_memory=True)
optimizer = torch.optim.Adam([
{
'params': net.parameters(),
'lr': 1e-4
},
])
logs = {'best_it': 0, 'best': 0}
sal_data_iter = iter(train_loader)
i_sal_data = 0
for it in tqdm(range(opt.max)):
# if it > 1000 and it % 100 == 0:
# optimizer.param_groups[0]['lr'] *= 0.5
if i_sal_data >= len(train_loader):
sal_data_iter = iter(train_loader)
i_sal_data = 0
data, lbl, _ = sal_data_iter.next()
i_sal_data += 1
data = data.cuda()
lbl = lbl[0].unsqueeze(1)
noisy_label = (lbl.numpy() +
np.random.binomial(1,
float(p) / 100.0, (256, 256))) % 2
noisy_label = torch.Tensor(noisy_label).cuda()
lbl = lbl.cuda()
msk = net(data, noisy_label)
loss = F.binary_cross_entropy_with_logits(msk, lbl)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if it % 10 == 0:
writer.add_scalar('loss', loss.item(), it)
image = make_image_grid(data[:6], mean, std)
writer.add_image('Image', torchvision.utils.make_grid(image), it)
big_msk = F.sigmoid(msk).expand(-1, 3, -1, -1)
writer.add_image('msk',
torchvision.utils.make_grid(big_msk.data[:6]), it)
big_msk = lbl.expand(-1, 3, -1, -1)
writer.add_image('gt',
torchvision.utils.make_grid(big_msk.data[:6]), it)
# if it % 200 == 0:
if it != 0 and it % 100 == 0:
fm, mae = validate(val_loader, net,
os.path.join(check_dir, 'results'),
os.path.join(opt.val_dir, 'masks'))
print(u'损失: %.4f' % (loss.item()))
print(u'最大FM: iteration %d的%.4f, 这次FM: %.4f' %
(logs['best_it'], logs['best'], fm))
logs[it] = {'FM': fm}
if fm > logs['best']:
logs['best'] = fm
logs['best_it'] = it
torch.save(net.state_dict(), '%s/net-best.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
torch.save(net.state_dict(), '%s/net-iter%d.pth' % (check_dir, it))
writer = Logger(pathlog)
with open("{}_capsal.txt".format(args.pathsave), "w") as f:
for k, v in vars(args).items():
line = "{}: {}\n".format(k, v)
print(line)
f.write(line)
with open('vocab_trainval.pkl', 'rb') as f:
vocab = pickle.load(f)
sal_val_loader = torch.utils.data.DataLoader(Folder(args.pathimg_val_sal,
args.pathann_val_sal,
crop=None,
flip=False,
rotate=None,
size=256,
mean=None,
std=None,
training=False),
batch_size=args.batchsize,
shuffle=False,
num_workers=args.numworkers,
pin_memory=True)
cap_train_loader = torch.utils.data.DataLoader(
CocoCaption(args.pathimg_train_cap,
args.pathann_train_cap,
vocab,
crop=None,
flip=True,
rotate=None,
voc_val_gt_dir = '%s/data/datasets/segmentation_Dataset/VOCdevkit/VOC2012/SegmentationClass' % home
voc_train_split = '%s/data/datasets/segmentation_Dataset/VOCdevkit/VOC2012/ImageSets/Segmentation/argtrain.txt' % home
voc_val_split = '%s/data/datasets/segmentation_Dataset/VOCdevkit/VOC2012/ImageSets/Segmentation/val.txt' % home
sal_train_img_dir = '%s/data/datasets/saliency_Dataset/DUT-train/images' % home
sal_train_gt_dir = '%s/data/datasets/saliency_Dataset/DUT-train/masks' % home
sal_val_img_dir = '%s/data/datasets/saliency_Dataset/ECSSD/images' % home
sal_val_gt_dir = '%s/data/datasets/saliency_Dataset/ECSSD/masks' % home
sal_train_loader = torch.utils.data.DataLoader(Folder(sal_train_img_dir,
sal_train_gt_dir,
crop=0.9,
flip=True,
rotate=None,
size=opt.imageSize,
mean=opt.mean,
std=opt.std,
training=True),
batch_size=opt.batchSize,
shuffle=True,
num_workers=4,
pin_memory=True)
sal_val_loader = torch.utils.data.DataLoader(Folder(sal_val_img_dir,
sal_val_gt_dir,
crop=None,
flip=False,
rotate=None,
size=opt.imageSize,
from options.test_options import TestOptions
opt = TestOptions() # set CUDA_VISIBLE_DEVICES before import torch
opt = opt.parse()
home = os.path.expanduser("~")
with open('vocab.pkl', 'rb') as f:
vocab = pickle.load(f)
val_loader = torch.utils.data.DataLoader(Folder(opt.val_img_dir,
opt.val_gt_dir,
crop=None,
flip=False,
rotate=None,
size=opt.imageSize,
mean=opt.mean,
std=opt.std,
training=False),
batch_size=opt.batchSize,
shuffle=True,
num_workers=4,
pin_memory=True)
model = getattr(models, opt.model)(opt, vocab_size=len(vocab), n_class=200)
def test(model):
print("============================= TEST ============================")
model.switch_to_eval()
voc_test_loader = torch.utils.data.DataLoader(
VOC(voc_val_img_dir, voc_val_gt_dir, voc_val_split,
crop=None, flip=False, rotate=None, size=opt.imageSize,
mean=opt.mean, std=opt.std, training=False),
batch_size=opt.batchSize, shuffle=True, num_workers=4, pin_memory=True)
voc_train_loader = torch.utils.data.DataLoader(
VOC(voc_train_img_dir, voc_train_gt_dir, voc_train_split,
crop=None, flip=False, rotate=None, size=opt.imageSize,
mean=opt.mean, std=opt.std, training=-1),
batch_size=opt.batchSize, shuffle=False, num_workers=4, pin_memory=True)
sal_val_loader = torch.utils.data.DataLoader(
Folder(sal_val_img_dir, sal_val_gt_dir,
crop=None, flip=False, rotate=None, size=opt.imageSize,
mean=opt.mean, std=opt.std, training=False),
batch_size=opt.batchSize, shuffle=True, num_workers=4, pin_memory=True)
def test(model):
print("============================= TEST ============================")
model.switch_to_eval()
for i, (img, name, WW, HH) in tqdm(enumerate(voc_test_loader), desc='testing'):
model.test(img, name, WW, HH)
model.switch_to_train()
miou = evaluate_iou(opt.results_dir, voc_val_gt_dir, c_output)
model.performance = {'miou': miou}
print(miou)
with open('val_voc.json', 'w') as f:
json.dump(model.performance, f)