def getNames(self):
names = list(map(lambda x: sorted(os.listdir(x)), self.path))
paths = [
list(map(lambda x: utils.genPath(path, x), name))
for name, path in zip(names, self.path)
]
self.names = {'X': names[0], 'Y': names[1]}
self.paths = paths
self.size = len(paths[0])
def eval_Saliency(self, Model, epoch=0, supervised=True):
savedict = {}
Outputs = {
val: Metrics.getOutPuts(Model,
valdata['X'],
self.Loader,
supervised=supervised)
for val, valdata in self.Loader.valdatas.items()
}
for val in self.Loader.valdatas.keys():
pred, valdata = Outputs[val]['final'], self.Loader.valdatas[val][
'Y']
F = Metrics.maxF(pred, valdata, self.Loader.ids[-1])
M = Metrics.mae(pred, valdata)
saves = self.scores[val].update([F, M], epoch)
savedict[val] = saves
for val, score in self.scores.items():
score.print_present()
print('-----------------------------------------')
if self.Loader.MODE == 'train':
torch.save(utils.makeDict(Model.state_dict()),
utils.genPath(self.Loader.spath, 'present.pkl'))
for val, saves in savedict.items():
for idx, save in enumerate(saves):
if save:
torch.save(
utils.makeDict(Model.state_dict()),
utils.genPath(self.Loader.spath, val + '_' +
['F', 'M'][idx] + '.pkl'))
for val, score in self.scores.items():
score.print_best()
else:
for val in self.Loader.valdatas.keys():
Outputs[val]['Name'] = self.Loader.valdatas[val]['Name']
Outputs[val]['Shape'] = self.Loader.valdatas[val]['Shape']
return Outputs if self.Loader.save else None
def __init__(self, name, mode='train', shape=224, debug=False):
super(_DataSet, self).__init__()
self.name = name
self.tic = time.time()
self.mean = np.array([0.485, 0.458, 0.407]).reshape([1, 3, 1, 1])
self.std = np.array([0.229, 0.224, 0.225]).reshape([1, 3, 1, 1])
self.dataset = utils.loadJson(utils.genPath('setup'))['DATASET']
assert name in self.dataset
self.path = self.dataset[self.name]
assert mode in ['train', 'test']
self.mode = mode
self.shape = shape
self.flip_code = [1, 0, -1]
self.getNames()
def save_preds(self):
for valname, output in self.Outputs.items():
rpath = utils.genPath(self.Loader.rpath, self.Loader.sub, valname)
if not os.path.exists(rpath):
os.makedirs(rpath)
names, shapes, finals, time = output['Name']['Y'], output[
'Shape'], output['final'] * 255., output['time']
for name, shape, final in zip(names, shapes, finals):
ppath = utils.genPath(rpath, 'final')
if not os.path.exists(ppath):
os.makedirs(ppath)
Image.fromarray(np.uint8(final)).resize(
(shape), Image.BICUBIC).save(utils.genPath(ppath, name))
if self.supervised:
preds, conts = output['preds'] * 255., output['contour'] * 255.
for name, shape, pred, cont in zip(names, shapes, preds,
conts):
for idx, pre in enumerate(pred):
pred_path = utils.genPath(rpath,
'pred_' + str(idx + 1))
if not os.path.exists(pred_path):
os.makedirs(pred_path)
Image.fromarray(np.uint8(pre)).convert('L').save(
utils.genPath(pred_path,
name.split('.')[0] + '.png'))
for idx, pre in enumerate(cont):
pred_path = utils.genPath(rpath,
'cont_' + str(idx + 1))
if not os.path.exists(pred_path):
os.makedirs(pred_path)
Image.fromarray(np.uint8(pre)).convert('L').save(
utils.genPath(pred_path,
name.split('.')[0] + '.png'))
print('Save predictions of datasets: {}.'.format(valname))
def loading(self, opt):
self.sub = opt.sub
self.debug = opt.debug
self.model = opt.model
if self.MODE == 'train':
self.weights = [1, 1, 1, 1, 1, 1]
self.optim = opt.optim
self.batch = opt.batch
self.scheduler = opt.scheduler
self.lr = opt.lr
self.plist = [['encoder', self.lr * 0.1], ['decoder', self.lr]]
self.trset = 'SOD' if self.debug else opt.trset
self.trSet = src.DataSet(self.trset,
mode='train',
shape=opt.size,
debug=self.debug)
self.epoch = int(math.ceil(
self.trSet.size / self.batch)) if self.MODE == 'train' else 10
self.iter = self.epoch * 25
else:
self.batch = 1
os.environ["CUDA_VISIBLE_DEVICES"] = opt.ids
num_gpu = len(opt.ids.split(','))
self.ids = list(range(num_gpu))
print('Backbone: {}, Using Gpu: {}'.format(self.model, opt.ids))
self.cpu = opt.cpu
self.save = not opt.no_save
self.supervised = opt.supervised
self.rpath = opt.rpath
if self.save and not os.path.exists(self.rpath):
os.makedirs(self.rpath)
npath = utils.genPath(self.rpath, self.sub)
if not os.path.exists(npath):
os.makedirs(npath)
self.spath = utils.genPath(opt.spath, self.model, self.sub)
if not os.path.exists(self.spath):
os.makedirs(self.spath)
self.mpath = self.spath + '/present.pkl'
if self.debug:
self.vals = ['SOD']
elif opt.vals == '':
self.vals = [
'SOD', 'PASCAL-S', 'ECSSD', 'DUTS-TE', 'HKU-IS', 'DUT-OMRON'
]
else:
self.vals = [
opt.vals,
]
self.mpath = self.spath + '/present.pkl'.format(opt.vals)
self.valdatas = collections.OrderedDict()
for val in self.vals:
self.valdatas[val] = src.DataSet(val, mode='test',
shape=opt.size).getFull()
self.channel = 16 if self.model.startswith('mobile') else 64
self.Model = baseline(self.model, self.channel)
print_model_parm_flops(self.Model)
print_model_parm_nums(self.Model)
self.mode = self.MODE