def main():
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
random.seed(args.seed)
torch.manual_seed(args.seed)
writer = SummaryWriter(args.snapshot_dir)
os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
cudnn.enabled = True
deeplab = get_segmentation_model("_".join([args.network, args.method]), num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
new_params = deeplab.state_dict().copy()
if 'wide' in args.network:
saved_state_dict = saved_state_dict['state_dict']
if 'vistas' in args.method:
saved_state_dict = saved_state_dict['body']
for i in saved_state_dict:
new_params[i] = saved_state_dict[i]
else:
for i in saved_state_dict:
i_parts = i.split('.')
if not 'classifier' in i_parts:
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
elif 'mobilenet' in args.network:
for i in saved_state_dict:
i_parts = i.split('.')
if not (i_parts[0]=='features' and i_parts[1]=='18') and not i_parts[0]=='classifier':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
else:
for i in saved_state_dict:
i_parts = i.split('.')
if not i_parts[0]=='fc' and not i_parts[0]=='last_linear' and not i_parts[0]=='classifier':
new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
if args.start_iters > 0:
deeplab.load_state_dict(saved_state_dict)
else:
deeplab.load_state_dict(new_params)
model = DataParallelModel(deeplab)
# model = nn.DataParallel(deeplab)
model.train()
model.float()
model.cuda()
criterion = CriterionCrossEntropy()
if "dsn" in args.method:
if args.ohem:
if args.ohem_single:
print('use ohem only for the second prediction map.')
criterion = CriterionOhemDSN_single(thres=args.ohem_thres, min_kept=args.ohem_keep, dsn_weight=float(args.dsn_weight))
else:
criterion = CriterionOhemDSN(thres=args.ohem_thres, min_kept=args.ohem_keep, dsn_weight=float(args.dsn_weight), use_weight=True)
else:
criterion = CriterionDSN(dsn_weight=float(args.dsn_weight), use_weight=True)
criterion = DataParallelCriterion(criterion)
criterion.cuda()
cudnn.benchmark = True
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(get_segmentation_dataset(args.dataset, root=args.data_dir, list_path=args.data_list,
max_iters=args.num_steps*args.batch_size, crop_size=input_size,
scale=args.random_scale, mirror=args.random_mirror, network=args.network),
batch_size=args.batch_size, shuffle=True, num_workers=1, pin_memory=True)
optimizer = optim.SGD([{'params': filter(lambda p: p.requires_grad, deeplab.parameters()), 'lr': args.learning_rate }],
lr=args.learning_rate, momentum=args.momentum, weight_decay=args.weight_decay)
optimizer.zero_grad()
for i_iter, batch in enumerate(trainloader):
sys.stdout.flush()
i_iter += args.start_iters
images, labels, _, _ = batch
images = Variable(images.cuda())
labels = Variable(labels.long().cuda())
optimizer.zero_grad()
lr = adjust_learning_rate(optimizer, i_iter)
if args.fix_lr:
lr = args.learning_rate
print('learning_rate: {}'.format(lr))
if 'gt' in args.method:
preds = model(images, labels)
else:
preds = model(images)
loss = criterion(preds, labels)
loss.backward()
optimizer.step()
if i_iter % 100 == 0:
writer.add_scalar('learning_rate', lr, i_iter)
writer.add_scalar('loss', loss.data.cpu().numpy(), i_iter)
print('iter = {} of {} completed, loss = {}'.format(i_iter, args.num_steps, loss.data.cpu().numpy()))
if i_iter >= args.num_steps-1:
print('save model ...')
torch.save(deeplab.state_dict(),osp.join(args.snapshot_dir, 'CS_scenes_'+str(args.num_steps)+'.pth'))
break
if i_iter % args.save_pred_every == 0:
print('taking snapshot ...')
torch.save(deeplab.state_dict(),osp.join(args.snapshot_dir, 'CS_scenes_'+str(i_iter)+'.pth'))
end = timeit.default_timer()
print(end-start,'seconds')
def main():
"""Create the model and start the evaluation process."""
args = Parameters().parse()
# #
# args.method = 'student_res18_pre'
args.method = 'student_esp_d'
args.dataset = 'camvid_light'
args.data_list = "/ssd/yifan/SegNet/CamVid/test.txt"
args.data_dir = "/ssd/yifan/"
args.num_classes = 11
# args.method='psp_dsn_floor'
args.restore_from = "./checkpoint/Camvid/ESP/base_57.8.pth"
# args.restore_from="/teamscratch/msravcshare/v-yifan/ESPNet/train/0.4results_enc_01_enc_2_8/model_298.pth"
# args.restore_from = "/teamscratch/msravcshare/v-yifacd n/sd_pytorch0.5/checkpoint/snapshots_psp_dsn_floor_1e-2_40000_TEACHER864/CS_scenes_40000.pth"
# args.restore_from = "/teamscratch/msravcshare/v-yifan/sd_pytorch0.5/checkpoint/snapshots_psp_dsn_floor_1e-2_40000_TEACHER5121024_esp/CS_scenes_40000.pth"
# args.data_list = '/teamscratch/msravcshare/v-yifan/deeplab_v3/dataset/list/cityscapes/train.lst'
args.batch_size = 1
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
print(args)
output_path = args.output_path
if not os.path.exists(output_path):
os.makedirs(output_path)
# args.method='psp_dsn'
deeplab = get_segmentation_model(args.method, num_classes=args.num_classes)
ignore_label = 255
id_to_trainid = {
-1: ignore_label,
0: ignore_label,
1: ignore_label,
2: ignore_label,
3: ignore_label,
4: ignore_label,
5: ignore_label,
6: ignore_label,
7: 0,
8: 1,
9: ignore_label,
10: ignore_label,
11: 2,
12: 3,
13: 4,
14: ignore_label,
15: ignore_label,
16: ignore_label,
17: 5,
18: ignore_label,
19: 6,
20: 7,
21: 8,
22: 9,
23: 10,
24: 11,
25: 12,
26: 13,
27: 14,
28: 15,
29: ignore_label,
30: ignore_label,
31: 16,
32: 17,
33: 18
}
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
# args.restore_from="/teamscratch/msravcshare/v-yifan/sd_pytorch0.3/checkpoint/snapshots_resnet_psp_dsn_1e-4_5e-4_8_20000_DSN_0.4_769light/CS_scenes_20000.pth"
# if 'dense' in args.method:
#
if args.restore_from is not None:
saved_state_dict = torch.load(args.restore_from)
c_keys = saved_state_dict.keys()
for i in c_keys:
flag = i.split('.')[0]
if 'module' in flag:
deeplab = nn.DataParallel(deeplab)
deeplab.load_state_dict(saved_state_dict)
if 'module' not in flag:
deeplab = nn.DataParallel(deeplab)
# if 'dense' not in args.method:
# deeplab = nn.DataParallel(deeplab)
model = deeplab
model.eval()
model.cuda()
# args.dataset='cityscapes_light'
testloader = data.DataLoader(get_segmentation_dataset(
args.dataset,
root=args.data_dir,
list_path=args.data_list,
crop_size=(360, 480),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=args.batch_size,
shuffle=False,
pin_memory=True)
data_list = []
confusion_matrix = np.zeros((args.num_classes, args.num_classes))
palette = get_palette(20)
image_id = 0
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
if args.side:
image, label, _, size, name = batch
elif 'sd' in args.dataset:
_, image, label, size, name = batch
else:
image, label, size, name = batch
# print('image name: {}'.format(name))
size = size[0].numpy()
output = predict_esp(model, image)
# seg_pred = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)
result = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)
# result=cv2.resize(result, (1024, 1024), interpolation=cv2.INTER_NEAREST)
m_seg_pred = ma.masked_array(result, mask=torch.eq(label, 255))
ma.set_fill_value(m_seg_pred, 20)
seg_pred = m_seg_pred
for i in range(image.size(0)):
image_id += 1
print('%d th segmentation map generated ...' % (image_id))
args.store_output = 'True'
output_path = './esp_camvid_base/'
if not os.path.exists(output_path):
os.mkdir(output_path)
if args.store_output == 'True':
# print('a')
output_im = PILImage.fromarray(seg_pred[i])
output_im.putpalette(palette)
output_im.save(output_path + '/' + name[i] + '.png')
seg_gt = np.asarray(label.numpy()[:, :size[0], :size[1]], dtype=np.int)
ignore_index = seg_gt != 255
seg_gt = seg_gt[ignore_index]
seg_pred = seg_pred[ignore_index]
confusion_matrix += get_confusion_matrix(seg_gt, seg_pred,
args.num_classes)
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
IU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IU = IU_array.mean()
print({'meanIU': mean_IU, 'IU_array': IU_array})
print("confusion matrix\n")
print(confusion_matrix)
def main():
"""Create the model and start the evaluation process."""
args = Parameters().parse()
# file_log = open(args.log_file, "w")
# sys.stdout = sys.stderr = file_log
print("Input arguments:")
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
sys.stdout.flush()
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
ignore_label = args.ignore_label
output_path = args.output_path
if not os.path.exists(output_path):
os.makedirs(output_path)
deeplab = get_segmentation_model("_".join([args.network, args.method]),
num_classes=args.num_classes)
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
saved_state_dict = torch.load(args.restore_from)
deeplab.load_state_dict(saved_state_dict)
model = nn.DataParallel(deeplab)
model.eval()
model.cuda()
testloader = data.DataLoader(get_segmentation_dataset(
args.dataset,
root=args.data_dir,
list_path=args.data_list,
crop_size=(1024, 2048),
scale=False,
mirror=False,
network=args.network),
batch_size=args.batch_size,
shuffle=False,
pin_memory=True)
data_list = []
confusion_matrix = np.zeros((args.num_classes, args.num_classes))
palette = get_palette(256)
id_to_trainid = {
-1: ignore_label,
0: ignore_label,
1: ignore_label,
2: ignore_label,
3: ignore_label,
4: ignore_label,
5: ignore_label,
6: ignore_label,
7: 0,
8: 1,
9: ignore_label,
10: ignore_label,
11: 2,
12: 3,
13: 4,
14: ignore_label,
15: ignore_label,
16: ignore_label,
17: 5,
18: ignore_label,
19: 6,
20: 7,
21: 8,
22: 9,
23: 10,
24: 11,
25: 12,
26: 13,
27: 14,
28: 15,
29: ignore_label,
30: ignore_label,
31: 16,
32: 17,
33: 18
}
image_id = 0
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
image, size, name = batch
size = size[0].numpy()
if torch_ver == '0.3':
if args.use_ms == 'True':
output = predict_multi_scale(model, image.numpy(),
([0.75, 1, 1.25]),
args.num_classes, args.use_flip,
args.method)
else:
output = predict_whole_img(model,
image.numpy(),
args.num_classes,
args.method,
scale=float(args.whole_scale))
else:
with torch.no_grad():
if args.use_ms == 'True':
output = predict_multi_scale(model, image.numpy(),
([0.75, 1, 1.25]),
args.num_classes,
args.use_flip, args.method)
else:
output = predict_whole_img(model,
image.numpy(),
args.num_classes,
args.method,
scale=float(args.whole_scale))
seg_pred = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)
seg_pred = id2trainId(seg_pred, id_to_trainid, reverse=True)
for i in range(image.size(0)):
image_id += 1
print('%d th segmentation map generated ...' % (image_id))
sys.stdout.flush()
if args.store_output == 'True':
output_im = PILImage.fromarray(seg_pred[i])
output_im.putpalette(palette)
output_im.save(output_path + '/' + name[i] + '.png')
def main():
"""Create the model and start the evaluation process."""
args = Parameters().parse()
# file_log = open(args.log_file, "w")
# sys.stdout = sys.stderr = file_log
print("Input arguments:")
sys.stdout.flush()
for key, val in vars(args).items():
print("{:16} {}".format(key, val))
h, w = map(int, args.input_size.split(','))
input_size = (h, w)
output_path = args.output_path
if not os.path.exists(output_path):
os.makedirs(output_path)
deeplab = get_segmentation_model("_".join([args.network, args.method]),
num_classes=args.num_classes)
ignore_label = 255
id_to_trainid = {
-1: ignore_label,
0: ignore_label,
1: ignore_label,
2: ignore_label,
3: ignore_label,
4: ignore_label,
5: ignore_label,
6: ignore_label,
7: 0,
8: 1,
9: ignore_label,
10: ignore_label,
11: 2,
12: 3,
13: 4,
14: ignore_label,
15: ignore_label,
16: ignore_label,
17: 5,
18: ignore_label,
19: 6,
20: 7,
21: 8,
22: 9,
23: 10,
24: 11,
25: 12,
26: 13,
27: 14,
28: 15,
29: ignore_label,
30: ignore_label,
31: 16,
32: 17,
33: 18
}
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
import urllib.request
local_checkpoint, _ = urllib.request.urlretrieve(
'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnet101-imagenet.pth',
'resnet101-imagenet.pth')
saved_state_dict = torch.load(local_checkpoint)
deeplab.load_state_dict(saved_state_dict)
model = nn.DataParallel(deeplab)
model.eval()
model.cuda()
testloader = data.DataLoader(get_segmentation_dataset(
args.dataset,
root=args.data_dir,
list_path=args.data_list,
crop_size=(1024, 2048),
scale=False,
mirror=False,
network=args.network),
batch_size=args.batch_size,
shuffle=False,
pin_memory=True)
data_list = []
confusion_matrix = np.zeros((args.num_classes, args.num_classes))
palette = get_palette(20)
image_id = 0
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % (index))
sys.stdout.flush()
image, label, size, name = batch
size = size[0].numpy()
if torch_ver == '0.3':
if args.use_ms == 'True':
output = predict_multi_scale(model, image.numpy(),
([0.75, 1, 1.25]), input_size,
args.num_classes, args.use_flip,
args.method)
else:
if args.use_flip == 'True':
output = predict_multi_scale(model, image.numpy(),
([args.whole_scale]),
input_size, args.num_classes,
args.use_flip, args.method)
else:
if 'gt' in args.method:
label = Variable(label.long().cuda())
output = predict_whole_img_w_label(
model,
image.numpy(),
args.num_classes,
args.method,
scale=float(args.whole_scale),
label=label)
else:
output = predict_whole_img(model,
image.numpy(),
args.num_classes,
args.method,
scale=float(
args.whole_scale))
else:
with torch.no_grad():
if args.use_ms == 'True':
output = predict_multi_scale(model, image.numpy(),
([0.75, 1, 1.25]), input_size,
args.num_classes,
args.use_flip, args.method)
else:
if args.use_flip == 'True':
output = predict_multi_scale(model, image.numpy(),
([args.whole_scale]),
input_size,
args.num_classes,
args.use_flip,
args.method)
else:
if 'gt' in args.method:
output = predict_whole_img_w_label(
model,
image.numpy(),
args.num_classes,
args.method,
scale=float(args.whole_scale),
label=Variable(label.long().cuda()))
else:
output = predict_whole_img(model,
image.numpy(),
args.num_classes,
args.method,
scale=float(
args.whole_scale))
seg_pred = np.asarray(np.argmax(output, axis=3), dtype=np.uint8)
m_seg_pred = ma.masked_array(seg_pred, mask=torch.eq(label, 255))
ma.set_fill_value(m_seg_pred, 20)
seg_pred = m_seg_pred
for i in range(image.size(0)):
image_id += 1
print('%d th segmentation map generated ...' % (image_id))
sys.stdout.flush()
if args.store_output == 'True':
output_im = PILImage.fromarray(seg_pred[i])
output_im.putpalette(palette)
output_im.save(output_path + '/' + name[i] + '.png')
seg_gt = np.asarray(label.numpy()[:, :size[0], :size[1]], dtype=np.int)
ignore_index = seg_gt != 255
seg_gt = seg_gt[ignore_index]
seg_pred = seg_pred[ignore_index]
confusion_matrix += get_confusion_matrix(seg_gt, seg_pred,
args.num_classes)
pos = confusion_matrix.sum(1)
res = confusion_matrix.sum(0)
tp = np.diag(confusion_matrix)
IU_array = (tp / np.maximum(1.0, pos + res - tp))
mean_IU = IU_array.mean()
print({'meanIU': mean_IU, 'IU_array': IU_array})
print("confusion matrix\n")
print(confusion_matrix)
sys.stdout.flush()
import numpy as np
from torch.autograd import Variable
import timeit
args = Parameters().parse()
args.batch_size = 1
args.dataset = 'cityscapes_light'
methods = ['student_res18_pre']
args.data_list = '/teamscratch/msravcshare/v-yifan/deeplab_v3/dataset/list/cityscapes/val.lst'
IMG_MEAN = np.array((104.00698793, 116.66876762, 122.67891434),
dtype=np.float32)
testloader = data.DataLoader(get_segmentation_dataset(args.dataset,
root=args.data_dir,
list_path=args.data_list,
crop_size=(1024, 2048),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=args.batch_size,
shuffle=False,
pin_memory=True)
for method in methods:
args.method = method
student = get_segmentation_model(args.method, num_classes=args.num_classes)
# from network.md import MobileNet
# student=MobileNet()
student = add_flops_counting_methods(student)
student = student.cuda()
student = student.eval()