def evaluate():
## network
net = DeepLabLargeFOV(3, 21)
net.eval()
net.cuda()
model_pth = './res/model.pkl'
net.load_state_dict(torch.load(model_pth))
## dataset
ds = PascalVoc('./data/VOCdevkit/', mode='val')
## inference
ious = []
for i, (im, lb) in enumerate(ds):
im_org = im
im = ds.trans(im)
im = im.cuda().unsqueeze(0)
scores = net(im)
scores = F.interpolate(scores, im.size()[2:], mode='bilinear')
scores = scores.detach().cpu().numpy()
mask = np.argmax(scores, axis=1).squeeze(0).astype(np.uint8)
lb = np.asarray(lb)
iou = compute_iou(mask, lb)
ious.append(iou)
# if iou == 0:
# from colormap import color_map
# cm = color_map()
# H, W = mask.shape
# pic = np.empty((H, W, 3))
# for i in range(21):
# pic[mask == i] = cm[i]
# im_org = cv2.cvtColor(np.asarray(im_org), cv2.COLOR_RGB2BGR)
# cv2.imshow('pic', pic)
# cv2.imshow('org', im_org)
# cv2.waitKey(0)
print('iou is: {}'.format(iou))
mIOU = sum(ious) / len(ious)
print('iou in whole is: {}'.format(mIOU))
def load_pascal(image_set, year, devkit_path, shuffle=False):
"""
wrapper function for loading pascal voc dataset
Parameters:
----------
image_set : str
train, trainval...
year : str
2007, 2012 or combinations splitted by comma
devkit_path : str
root directory of dataset
shuffle : bool
whether to shuffle initial list
Returns:
----------
Imdb
"""
image_set = [y.strip() for y in image_set.split(',')]
assert image_set, "No image_set specified"
year = [y.strip() for y in year.split(',')]
assert year, "No year specified"
# make sure (# sets == # years)
if len(image_set) > 1 and len(year) == 1:
year = year * len(image_set)
if len(image_set) == 1 and len(year) > 1:
image_set = image_set * len(year)
assert len(image_set) == len(year), "Number of sets and year mismatch"
imdbs = []
for s, y in zip(image_set, year):
imdbs.append(PascalVoc(s, y, devkit_path, shuffle, is_train=True))
if len(imdbs) > 1:
return ConcatDB(imdbs, shuffle)
else:
return imdbs[0]
def load_pascal(image_set, devkit_path, shuffle=False, class_names=None):
"""
wrapper function for loading pascal voc dataset
Parameters:
----------
image_set : str
train, trainval...
devkit_path : str
root directory of dataset
shuffle : bool
whether to shuffle initial list
Returns:
----------
Imdb
"""
assert image_set, "No image_set specified"
return PascalVoc(image_set,
devkit_path,
shuffle,
is_train=True,
class_names=class_names)
def train():
## modules and losses
logger.info('creating model and loss module')
ignore_label = 255
net = DeepLabLargeFOV(3, 21)
net.train()
net.cuda()
net = nn.DataParallel(net)
Loss = nn.CrossEntropyLoss(ignore_index=ignore_label)
Loss.cuda()
## dataset
logger.info('creating dataset and dataloader')
batchsize = 30
ds = PascalVoc('./data/VOCdevkit/', mode='train')
dl = DataLoader(ds,
batch_size=batchsize,
shuffle=True,
num_workers=6,
drop_last=True)
## optimizer
logger.info('creating optimizer')
warmup_start_lr = 1e-6
warmup_iter = 1000
start_lr = 1e-3
lr_steps = [5000, 7000]
lr_factor = 0.1
momentum = 0.9
weight_decay = 1e-4
optimizer = Optimizer(params=net.parameters(),
warmup_start_lr=warmup_start_lr,
warmup_iter=warmup_iter,
start_lr=start_lr,
lr_steps=lr_steps,
lr_factor=lr_factor,
momentum=momentum,
weight_decay=weight_decay)
## train loop
iter_num = 8000
log_iter = 20
loss_avg = []
st = time.time()
alpha = 0.1
diter = iter(dl)
logger.info('start training')
for it in range(iter_num):
try:
im, lb = next(diter)
if not im.size()[0] == batchsize: continue
except StopIteration:
diter = iter(dl)
im, lb = next(diter)
im = im.cuda()
lb = lb.cuda()
lam = np.random.beta(alpha, alpha)
idx = torch.randperm(batchsize)
mix_im = im * lam + (1. - lam) * im[idx, :]
mix_lb = lb[idx, :]
optimizer.zero_grad()
out = net(mix_im)
out = F.interpolate(out, im.size()[2:], mode='bilinear')
lb = torch.squeeze(lb)
mix_lb = torch.squeeze(mix_lb)
loss = lam * Loss(out, lb) + (1. - lam) * Loss(out, mix_lb)
loss.backward()
optimizer.step()
# optimizer.zero_grad()
# out = net(im)
# out = F.interpolate(out, im.size()[2:], mode = 'bilinear')
# lb = torch.squeeze(lb)
# # out = out.permute(0, 2, 3, 1).contiguous().view(-1, 21)
# # lb = lb.permute(0, 2, 3, 1).contiguous().view(-1,)
# loss = Loss(out, lb)
# loss.backward()
# optimizer.step()
loss = loss.detach().cpu().numpy()
loss_avg.append(loss)
## log message
if it % log_iter == 0 and not it == 0:
loss_avg = sum(loss_avg) / len(loss_avg)
ed = time.time()
t_int = ed - st
lr = optimizer.get_lr()
msg = 'iter: {}/{}, loss: {:3f}'.format(it, iter_num, loss_avg)
msg = '{}, lr: {:4f}, time: {:3f}'.format(msg, lr, t_int)
logger.info(msg)
st = ed
loss_avg = []
## dump model
model_pth = './res/model.pkl'
torch.save(net.module.state_dict(), model_pth)
logger.info('training done, model saved to: {}'.format(model_pth))
def load_pascal(image_set,
year,
devkit_path,
shuffle=False,
class_names=None,
true_negative=None):
"""
wrapper function for loading pascal voc dataset
Parameters:
----------
image_set : str
train, trainval...
year : str
2007, 2012 or combinations splitted by comma
devkit_path : str
root directory of dataset
shuffle : bool
whether to shuffle initial list
Returns:
----------
Imdb
"""
# image_set is the set, like trainval, val, test.
image_set = [y.strip() for y in image_set.split(',')]
assert image_set, "No image_set specified"
# year, only used while the dataset us PASCAL VOC.
year = [y.strip() for y in year.split(',')]
assert year, "No year specified"
# make sure (# sets == # years)
# list * integer, return a list. This operation will copy the elements of list for integer times.
# Must ensure len(image_set) == len(year).
if len(image_set) > 1 and len(year) == 1:
year = year * len(image_set)
if len(image_set) == 1 and len(year) > 1:
image_set = image_set * len(year)
assert len(image_set) == len(year), "Number of sets and year mismatch"
imdbs = []
for s, y in zip(image_set, year):
'''
For example:
i=0, s="trainval", year="2007"
i=1, s="trainval", year="2012"
'''
imdbs.append(
PascalVoc(s,
y,
devkit_path,
shuffle,
is_train=True,
class_names=class_names,
names="label_map.txt",
true_negative_images=true_negative))
'''
devkit_path is the root path, can be "/home/ly/caffe-ssd/data/VOC0712".
class_names, default is None.
default class_names is None, so there needs a class file, which is the label_map.txt
in caffe-ssd. There are 21 classes, the first class is background. label form 0 to 20.
Must appoint the dirname.
'''
if len(imdbs) > 1:
return ConcatDB(imdbs, shuffle)
# class ConcatDB, ConcatDB is used to concatenate multiple imdbs to form a larger db.
else:
return imdbs[0]