def main(args):
"""Create the model and start the evaluation process."""
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(ListDataSet(args.data_dir,
args.img_list,
args.lbl_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
split=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
with torch.no_grad():
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, _, name = batch
if args.model == 'DeeplabMulti':
output1, output2 = model(Variable(image).cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG':
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' %
(args.save, name.split('.')[0]))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=input_size, mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True)
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
image = image.to(device)
if args.model == 'DeeplabMulti':
output1, output2,_,_ = model(image)
output = interp(output2).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output = model(image)
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'Deeplab':
model = Res_Deeplab(num_classes=args.num_classes)
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
model_dict.update(saved_state_dict)
###
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
else:
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
if args.model == 'Deeplab':
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG':
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
def eval(pth, cityscapes_eval_dir, i_iter):
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if args.model == 'ResNet':
model = Res_Deeplab(num_classes=args.num_classes)
saved_state_dict = torch.load(pth)
elif args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes)
saved_state_dict = torch.load(pth)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
cityscapesloader = data.DataLoader(cityscapesDataSet(
args.cityscapes_data_dir,
args.cityscapes_data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
for index, batch in enumerate(cityscapesloader):
with torch.no_grad():
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (cityscapes_eval_dir, name))
output_col.save('%s/%s_color.png' %
(cityscapes_eval_dir, name.split('.')[0]))
if i_iter == 0:
break
def main():
args = get_arguments()
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(args.cuda_device_id)
if args.model == 'ResNet':
model = DeeplabMulti(num_classes=args.num_classes)
if args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes)
device = torch.device("cuda" if not args.cpu else "cpu")
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.eval()
model.to(device)
trainloader = data.DataLoader(GTA5DataSet(args.data_dir, args.data_list, crop_size=(1024, 512),scale=False, mirror=False, mean=IMG_MEAN),
batch_size=1, shuffle=False, pin_memory=True)
count_class = np.zeros((19, 1))
class_center_temp = np.zeros((19, 256))
for index, batch in enumerate(trainloader):
if index % 100 == 0:
print( '%d processd' % index)
images, labels, _, _= batch
images = images.to(device)
labels = labels.long().to(device)
with torch.no_grad():
feature, _ = model(images)
class_center,count_class_t = class_center_precal(feature,labels)
count_class = count_class + count_class_t.numpy()
class_center_temp += class_center.cpu().data[0].numpy()
count_class[count_class==0] = 1 #in case divide 0 error
class_center = class_center_temp/count_class
np.save('./source_center.npy',class_center)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
config_path = os.path.join(os.path.dirname(args.restore_from), 'opts.yaml')
with open(config_path, 'r') as stream:
config = yaml.load(stream)
args.model = config['model']
print('ModelType:%s' % args.model)
print('NormType:%s' % config['norm_style'])
gpu0 = args.gpu
batchsize = args.batchsize
model_name = os.path.basename(os.path.dirname(args.restore_from))
#args.save += model_name
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes,
use_se=config['use_se'],
train_bn=False,
norm_style=config['norm_style'])
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
try:
model.load_state_dict(saved_state_dict)
except:
model = torch.nn.DataParallel(model)
model.load_state_dict(saved_state_dict)
model = torch.nn.DataParallel(model)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(robotDataSet(args.data_dir,
args.data_list,
crop_size=(960, 1280),
resize_size=(1280, 960),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
scale = 1.25
testloader2 = data.DataLoader(robotDataSet(
args.data_dir,
args.data_list,
crop_size=(round(960 * scale), round(1280 * scale)),
resize_size=(round(1280 * scale), round(960 * scale)),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(960, 1280),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(960, 1280), mode='bilinear')
sm = torch.nn.Softmax(dim=1)
for index, img_data in enumerate(zip(testloader, testloader2)):
batch, batch2 = img_data
image, _, _, name = batch
image2, _, _, name2 = batch2
print(image.shape)
inputs = image.cuda()
inputs2 = image2.cuda()
print('\r>>>>Extracting feature...%04d/%04d' %
(index * batchsize, NUM_STEPS),
end='')
if args.model == 'DeepLab':
with torch.no_grad():
output1, output2 = model(inputs)
output_batch = interp(sm(0.5 * output1 + output2))
output1, output2 = model(fliplr(inputs))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
del output1, output2, inputs
output1, output2 = model(inputs2)
output_batch += interp(sm(0.5 * output1 + output2))
output1, output2 = model(fliplr(inputs2))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
del output1, output2, inputs2
output_batch = output_batch.cpu().data.numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output_batch = model(Variable(image).cuda())
output_batch = interp(output_batch).cpu().data.numpy()
output_batch = output_batch.transpose(0, 2, 3, 1)
score_batch = np.max(output_batch, axis=3)
output_batch = np.asarray(np.argmax(output_batch, axis=3),
dtype=np.uint8)
#output_batch[score_batch<3.6] = 255 #3.6 = 4*0.9
for i in range(output_batch.shape[0]):
output = output_batch[i, :, :]
output_col = colorize_mask(output)
output = Image.fromarray(output)
name_tmp = name[i].split('/')[-1]
dir_name = name[i].split('/')[-2]
save_path = args.save + '/' + dir_name
#save_path = re.replace(save_path, 'leftImg8bit', 'pseudo')
#print(save_path)
if not os.path.isdir(save_path):
os.mkdir(save_path)
output.save('%s/%s' % (save_path, name_tmp))
print('%s/%s' % (save_path, name_tmp))
output_col.save('%s/%s_color.png' %
(save_path, name_tmp.split('.')[0]))
return args.save
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
#model = Res_Deeplab(num_classes=args.num_classes)
model = DeepLab(backbone='resnet', output_stride=8)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {
k: v
for k, v in saved_state_dict.items() if k in model_dict
}
model_dict.update(saved_state_dict)
model.load_state_dict(saved_state_dict)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
num_classes = 20
tp_list = [0] * num_classes
fp_list = [0] * num_classes
fn_list = [0] * num_classes
iou_list = [0] * num_classes
hist = np.zeros((21, 21))
group = 1
scorer = SegScorer(num_classes=21)
datalayer = SSDatalayer(group)
cos_similarity_func = nn.CosineSimilarity()
for count in tqdm(range(1000)):
dat = datalayer.dequeue()
ref_img = dat['second_img'][0] # (3, 457, 500)
query_img = dat['first_img'][0] # (3, 375, 500)
query_label = dat['second_label'][0] # (1, 375, 500)
ref_label = dat['first_label'][0] # (1, 457, 500)
# query_img = dat['second_img'][0]
# ref_img = dat['first_img'][0]
# ref_label = dat['second_label'][0]
# query_label = dat['first_label'][0]
deploy_info = dat['deploy_info']
semantic_label = deploy_info['first_semantic_labels'][0][0] - 1 # 2
ref_img, ref_label = torch.Tensor(ref_img).cuda(), torch.Tensor(
ref_label).cuda()
query_img, query_label = torch.Tensor(query_img).cuda(), torch.Tensor(
query_label[0, :, :]).cuda()
#ref_img, ref_label = torch.Tensor(ref_img), torch.Tensor(ref_label)
#query_img, query_label = torch.Tensor(query_img), torch.Tensor(query_label[0, :, :])
# ref_img = ref_img*ref_label
ref_img_var, query_img_var = Variable(ref_img), Variable(query_img)
query_label_var, ref_label_var = Variable(query_label), Variable(
ref_label)
ref_img_var = torch.unsqueeze(ref_img_var, dim=0) # [1, 3, 457, 500]
ref_label_var = torch.unsqueeze(ref_label_var,
dim=1) # [1, 1, 457, 500]
query_img_var = torch.unsqueeze(query_img_var,
dim=0) # [1, 3, 375, 500]
query_label_var = torch.unsqueeze(query_label_var,
dim=0) # [1, 375, 500]
samples = torch.cat([ref_img_var, query_img_var], 0)
pred = model(samples, ref_label_var)
w, h = query_label.size()
pred = F.upsample(pred, size=(w, h), mode='bilinear') #[2, 416, 416]
pred = F.softmax(pred, dim=1).squeeze()
values, pred = torch.max(pred, dim=0)
#print(pred.shape)
pred = pred.data.cpu().numpy().astype(np.int32) # (333, 500)
#print(pred.shape)
org_img = get_org_img(
query_img.squeeze().cpu().data.numpy()) # 查询集的图片(375, 500, 3)
#print(org_img.shape)
img = mask_to_img(pred, org_img) # (375, 500, 3)mask和原图加权后的彩色图片
cv2.imwrite('save_bins/que_pred/query_set_1_%d.png' % (count), img)
query_label = query_label.cpu().numpy().astype(np.int32) # (333, 500)
class_ind = int(deploy_info['first_semantic_labels'][0][0]
) - 1 # because class indices from 1 in data layer,0
scorer.update(pred, query_label, class_ind + 1)
tp, tn, fp, fn = measure(query_label, pred)
# iou_img = tp/float(max(tn+fp+fn,1))
tp_list[class_ind] += tp
fp_list[class_ind] += fp
fn_list[class_ind] += fn
# max in case both pred and label are zero
iou_list = [
tp_list[ic] /
float(max(tp_list[ic] + fp_list[ic] + fn_list[ic], 1))
for ic in range(num_classes)
]
tmp_pred = pred
tmp_pred[tmp_pred > 0.5] = class_ind + 1
tmp_gt_label = query_label
tmp_gt_label[tmp_gt_label > 0.5] = class_ind + 1
hist += Metrics.fast_hist(tmp_pred, query_label, 21)
print("-------------GROUP %d-------------" % (group))
print(iou_list)
class_indexes = range(group * 5, (group + 1) * 5)
print('Mean:', np.mean(np.take(iou_list, class_indexes)))
'''
for group in range(2):
datalayer = SSDatalayer(group+1)
restore(args, model, group+1)
for count in tqdm(range(1000)):
dat = datalayer.dequeue()
ref_img = dat['second_img'][0]#(3, 457, 500)
query_img = dat['first_img'][0]#(3, 375, 500)
query_label = dat['second_label'][0]#(1, 375, 500)
ref_label = dat['first_label'][0]#(1, 457, 500)
# query_img = dat['second_img'][0]
# ref_img = dat['first_img'][0]
# ref_label = dat['second_label'][0]
# query_label = dat['first_label'][0]
deploy_info = dat['deploy_info']
semantic_label = deploy_info['first_semantic_labels'][0][0] - 1#2
ref_img, ref_label = torch.Tensor(ref_img).cuda(), torch.Tensor(ref_label).cuda()
query_img, query_label = torch.Tensor(query_img).cuda(), torch.Tensor(query_label[0,:,:]).cuda()
#ref_img, ref_label = torch.Tensor(ref_img), torch.Tensor(ref_label)
#query_img, query_label = torch.Tensor(query_img), torch.Tensor(query_label[0, :, :])
# ref_img = ref_img*ref_label
ref_img_var, query_img_var = Variable(ref_img), Variable(query_img)
query_label_var, ref_label_var = Variable(query_label), Variable(ref_label)
ref_img_var = torch.unsqueeze(ref_img_var,dim=0)#[1, 3, 457, 500]
ref_label_var = torch.unsqueeze(ref_label_var, dim=1)#[1, 1, 457, 500]
query_img_var = torch.unsqueeze(query_img_var, dim=0)#[1, 3, 375, 500]
query_label_var = torch.unsqueeze(query_label_var, dim=0)#[1, 375, 500]
logits = model(query_img_var, ref_img_var, ref_label_var,ref_label_var)
# w, h = query_label.size()
# outB_side = F.upsample(outB_side, size=(w, h), mode='bilinear')
# out_side = F.softmax(outB_side, dim=1).squeeze()
# values, pred = torch.max(out_side, dim=0)
values, pred = model.get_pred(logits, query_img_var)#values[2, 333, 500]
pred = pred.data.cpu().numpy().astype(np.int32)#(333, 500)
query_label = query_label.cpu().numpy().astype(np.int32)#(333, 500)
class_ind = int(deploy_info['first_semantic_labels'][0][0])-1 # because class indices from 1 in data layer,0
scorer.update(pred, query_label, class_ind+1)
tp, tn, fp, fn = measure(query_label, pred)
# iou_img = tp/float(max(tn+fp+fn,1))
tp_list[class_ind] += tp
fp_list[class_ind] += fp
fn_list[class_ind] += fn
# max in case both pred and label are zero
iou_list = [tp_list[ic] /
float(max(tp_list[ic] + fp_list[ic] + fn_list[ic],1))
for ic in range(num_classes)]
tmp_pred = pred
tmp_pred[tmp_pred>0.5] = class_ind+1
tmp_gt_label = query_label
tmp_gt_label[tmp_gt_label>0.5] = class_ind+1
hist += Metrics.fast_hist(tmp_pred, query_label, 21)
print("-------------GROUP %d-------------"%(group))
print(iou_list)
class_indexes = range(group*5, (group+1)*5)
print('Mean:', np.mean(np.take(iou_list, class_indexes)))
print('BMVC IOU', np.mean(np.take(iou_list, range(0,20))))
miou = Metrics.get_voc_iou(hist)
print('IOU:', miou, np.mean(miou))
'''
binary_hist = np.array((hist[0, 0], hist[0, 1:].sum(), hist[1:, 0].sum(),
hist[1:, 1:].sum())).reshape((2, 2))
bin_iu = np.diag(binary_hist) / (binary_hist.sum(1) + binary_hist.sum(0) -
np.diag(binary_hist))
print('Bin_iu:', bin_iu)
scores = scorer.score()
for k in scores.keys():
print(k, np.mean(scores[k]), scores[k])
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
model_dict.update(saved_state_dict)
###
model.load_state_dict(saved_state_dict)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
batch_size=1, shuffle=False, pin_memory=True)
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
image = image.to(device)
_, output2, features = model(image)
output = output2.cpu().data[0].numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
features = features.cpu().data[0].numpy()
#### tsne plot ####
gt_dir = '/project/AdaptSegNet/data/Cityscapes/data/gtFine/val/'
gt_file = name[0].split('/')[-1]
gt_file = gt_file.replace('leftImg8bit', 'gtFine_labelIds')
if 'frankfurt' in gt_file:
gt_file = gt_dir + 'frankfurt/' + gt_file
elif 'lindau' in gt_file:
gt_file = gt_dir + 'lindau/' + gt_file
elif 'munster' in gt_file:
gt_file = gt_dir + 'munster/' + gt_file
"""
json_dir = '/project/AdaptSegNet/dataset/cityscapes_list/'
with open(join(json_dir, 'info.json'), 'r') as fp:
info = json.load(fp)
num_classes = np.int(info['classes'])
print('Num classes', num_classes)
name_classes = np.array(info['label'], dtype=np.str)
mapping = np.array(info['label2train'], dtype=np.int)"""
labels = np.array(Image.open(gt_file))
labels = label_mask(labels).astype(np.float64)
labels = skimage.transform.resize(labels, [65, 129])
labels = label_mask(labels).astype(np.uint8)
F, H, W = np.squeeze(features).shape
features = features.reshape(F, H*W).transpose(1,0)
labels = labels.reshape(H*W,)
# features = features[1:1000, :]
# labels = labels[1:1000]
tsne = TSNE(random_state=RS).fit_transform(features)
pdb.set_trace()
scatter_plot(tsne, labels)
#### tsne plot ####
pdb.set_trace()
def main():
"""Create the model and start the training."""
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
gpu = args.gpu
tau = torch.ones(1) * args.tau
tau = tau.cuda(args.gpu)
# Create network
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params, False)
elif args.model == 'DeepLabVGG':
model = DeeplabVGG(pretrained=True, num_classes=args.num_classes)
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes)
model_D2 = FCDiscriminator(num_classes=args.num_classes)
model_D1.train()
model_D1.cuda(args.gpu)
model_D2.train()
model_D2.cuda(args.gpu)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
weak_transform = transforms.Compose([
# transforms.RandomCrop(32, 4),
# transforms.RandomRotation(30),
# transforms.Resize(1024),
transforms.ToTensor(),
# transforms.Normalize(mean, std),
# RandomCrop(768)
])
target_transform = transforms.Compose([
# transforms.RandomCrop(32, 4),
# transforms.RandomRotation(30),
# transforms.Normalize(mean, std)
# transforms.Resize(1024),
# transforms.ToTensor(),
# RandomCrop(768)
])
label_set = GTA5(
root=args.data_dir,
num_cls=19,
split='all',
remap_labels=True,
transform=weak_transform,
target_transform=target_transform,
scale=input_size,
# crop_transform=RandomCrop(int(768*(args.scale/1024))),
)
unlabel_set = Cityscapes(
root=args.data_dir_target,
split=args.set,
remap_labels=True,
transform=weak_transform,
target_transform=target_transform,
scale=input_size_target,
# crop_transform=RandomCrop(int(768*(args.scale/1024))),
)
test_set = Cityscapes(
root=args.data_dir_target,
split='val',
remap_labels=True,
transform=weak_transform,
target_transform=target_transform,
scale=input_size_target,
# crop_transform=RandomCrop(768)
)
label_loader = data.DataLoader(label_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=False)
unlabel_loader = data.DataLoader(unlabel_set,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=False)
test_loader = data.DataLoader(test_set,
batch_size=2,
shuffle=False,
num_workers=args.num_workers,
pin_memory=False)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
[model, model_D2,
model_D2], [optimizer, optimizer_D1, optimizer_D2
] = amp.initialize([model, model_D2, model_D2],
[optimizer, optimizer_D1, optimizer_D2],
opt_level="O1",
num_losses=7)
optimizer.zero_grad()
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
if args.gan == 'Vanilla':
bce_loss = torch.nn.BCEWithLogitsLoss()
elif args.gan == 'LS':
bce_loss = torch.nn.MSELoss()
interp = Interpolate(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
interp_target = Interpolate(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
interp_test = Interpolate(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# interp_test = Interpolate(size=(1024, 2048), mode='bilinear', align_corners=True)
normalize_transform = transforms.Compose([
torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
# labels for adversarial training
source_label = 0
target_label = 1
max_mIoU = 0
total_loss_seg_value1 = []
total_loss_adv_target_value1 = []
total_loss_D_value1 = []
total_loss_con_value1 = []
total_loss_seg_value2 = []
total_loss_adv_target_value2 = []
total_loss_D_value2 = []
total_loss_con_value2 = []
hist = np.zeros((num_cls, num_cls))
# for i_iter in range(args.num_steps):
for i_iter, (batch, batch_un) in enumerate(
zip(roundrobin_infinite(label_loader),
roundrobin_infinite(unlabel_loader))):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_con_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
loss_con_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
images, labels = batch
images_orig = images
images = transform_batch(images, normalize_transform)
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp(pred1)
pred2 = interp(pred2)
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss = loss_seg2 + args.lambda_seg * loss_seg1
# proper normalization
loss = loss / args.iter_size
with amp.scale_loss(loss, optimizer, loss_id=0) as scaled_loss:
scaled_loss.backward()
# loss.backward()
loss_seg_value1 += loss_seg1.data.cpu().numpy() / args.iter_size
loss_seg_value2 += loss_seg2.data.cpu().numpy() / args.iter_size
# train with target
images_tar, labels_tar = batch_un
images_tar_orig = images_tar
images_tar = transform_batch(images_tar, normalize_transform)
images_tar = Variable(images_tar).cuda(args.gpu)
pred_target1, pred_target2 = model(images_tar)
pred_target1 = interp_target(pred_target1)
pred_target2 = interp_target(pred_target2)
D_out1 = model_D1(F.softmax(pred_target1, dim=1))
D_out2 = model_D2(F.softmax(pred_target2, dim=1))
loss_adv_target1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(args.gpu))
loss_adv_target2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(args.gpu))
loss = args.lambda_adv_target1 * loss_adv_target1 + args.lambda_adv_target2 * loss_adv_target2
loss = loss / args.iter_size
with amp.scale_loss(loss, optimizer, loss_id=1) as scaled_loss:
scaled_loss.backward()
# loss.backward()
loss_adv_target_value1 += loss_adv_target1.data.cpu().numpy(
) / args.iter_size
loss_adv_target_value2 += loss_adv_target2.data.cpu().numpy(
) / args.iter_size
# train with consistency loss
# unsupervise phase
policies = RandAugment().get_batch_policy(args.batch_size)
rand_p1 = np.random.random(size=args.batch_size)
rand_p2 = np.random.random(size=args.batch_size)
random_dir = np.random.choice([-1, 1], size=[args.batch_size, 2])
images_aug = aug_batch_tensor(images_tar_orig, policies, rand_p1,
rand_p2, random_dir)
images_aug_orig = images_aug
images_aug = transform_batch(images_aug, normalize_transform)
images_aug = Variable(images_aug).cuda(args.gpu)
pred_target_aug1, pred_target_aug2 = model(images_aug)
pred_target_aug1 = interp_target(pred_target_aug1)
pred_target_aug2 = interp_target(pred_target_aug2)
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
max_pred1, psuedo_label1 = torch.max(F.softmax(pred_target1, dim=1), 1)
max_pred2, psuedo_label2 = torch.max(F.softmax(pred_target2, dim=1), 1)
psuedo_label1 = psuedo_label1.cpu().numpy().astype(np.float32)
psuedo_label1_thre = psuedo_label1.copy()
psuedo_label1_thre[(max_pred1 < tau).cpu().numpy().astype(
np.bool)] = 255 # threshold to don't care
psuedo_label1_thre = aug_batch_numpy(psuedo_label1_thre, policies,
rand_p1, rand_p2, random_dir)
psuedo_label2 = psuedo_label2.cpu().numpy().astype(np.float32)
psuedo_label2_thre = psuedo_label2.copy()
psuedo_label2_thre[(max_pred2 < tau).cpu().numpy().astype(
np.bool)] = 255 # threshold to don't care
psuedo_label2_thre = aug_batch_numpy(psuedo_label2_thre, policies,
rand_p1, rand_p2, random_dir)
psuedo_label1_thre = Variable(psuedo_label1_thre).cuda(args.gpu)
psuedo_label2_thre = Variable(psuedo_label2_thre).cuda(args.gpu)
if (psuedo_label1_thre != 255).sum().cpu().numpy() > 0:
# nll_loss doesn't support empty tensors
loss_con1 = loss_calc(pred_target_aug1, psuedo_label1_thre,
args.gpu)
loss_con_value1 += loss_con1.data.cpu().numpy() / args.iter_size
else:
loss_con1 = torch.tensor(0.0, requires_grad=True).cuda(args.gpu)
if (psuedo_label2_thre != 255).sum().cpu().numpy() > 0:
# nll_loss doesn't support empty tensors
loss_con2 = loss_calc(pred_target_aug2, psuedo_label2_thre,
args.gpu)
loss_con_value2 += loss_con2.data.cpu().numpy() / args.iter_size
else:
loss_con2 = torch.tensor(0.0, requires_grad=True).cuda(args.gpu)
loss = args.lambda_con * loss_con1 + args.lambda_con * loss_con2
# proper normalization
loss = loss / args.iter_size
with amp.scale_loss(loss, optimizer, loss_id=2) as scaled_loss:
scaled_loss.backward()
# loss.backward()
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
pred2 = pred2.detach()
D_out1 = model_D1(F.softmax(pred1, dim=1))
D_out2 = model_D2(F.softmax(pred2, dim=1))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(source_label)).cuda(args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
with amp.scale_loss(loss_D1, optimizer_D1, loss_id=3) as scaled_loss:
scaled_loss.backward()
# loss_D1.backward()
with amp.scale_loss(loss_D2, optimizer_D2, loss_id=4) as scaled_loss:
scaled_loss.backward()
# loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
# train with target
pred_target1 = pred_target1.detach()
pred_target2 = pred_target2.detach()
D_out1 = model_D1(F.softmax(pred_target1, dim=1))
D_out2 = model_D2(F.softmax(pred_target2, dim=1))
loss_D1 = bce_loss(
D_out1,
Variable(
torch.FloatTensor(
D_out1.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D2 = bce_loss(
D_out2,
Variable(
torch.FloatTensor(
D_out2.data.size()).fill_(target_label)).cuda(args.gpu))
loss_D1 = loss_D1 / args.iter_size / 2
loss_D2 = loss_D2 / args.iter_size / 2
with amp.scale_loss(loss_D1, optimizer_D1, loss_id=5) as scaled_loss:
scaled_loss.backward()
# loss_D1.backward()
with amp.scale_loss(loss_D2, optimizer_D2, loss_id=6) as scaled_loss:
scaled_loss.backward()
# loss_D2.backward()
loss_D_value1 += loss_D1.data.cpu().numpy()
loss_D_value2 += loss_D2.data.cpu().numpy()
optimizer.step()
optimizer_D1.step()
optimizer_D2.step()
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}, loss_con1 = {8:.3f}, loss_con2 = {9:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1, loss_seg_value2,
loss_adv_target_value1, loss_adv_target_value2,
loss_D_value1, loss_D_value2, loss_con_value1,
loss_con_value2))
total_loss_seg_value1.append(loss_seg_value1)
total_loss_adv_target_value1.append(loss_adv_target_value1)
total_loss_D_value1.append(loss_D_value1)
total_loss_con_value1.append(loss_con_value1)
total_loss_seg_value2.append(loss_seg_value2)
total_loss_adv_target_value2.append(loss_adv_target_value2)
total_loss_D_value2.append(loss_D_value2)
total_loss_con_value2.append(loss_con_value2)
hist += fast_hist(
labels.cpu().numpy().flatten().astype(int),
torch.argmax(pred2, dim=1).cpu().numpy().flatten().astype(int),
num_cls)
if i_iter % 10 == 0:
print('({}/{})'.format(i_iter + 1, int(args.num_steps)))
acc_overall, acc_percls, iu, fwIU = result_stats(hist)
mIoU = np.mean(iu)
per_class = [[classes[i], acc] for i, acc in list(enumerate(iu))]
per_class = np.array(per_class).flatten()
print(
('per cls IoU :' + ('\n{:>14s} : {}') * 19).format(*per_class))
print('mIoU : {:0.2f}'.format(np.mean(iu)))
print('fwIoU : {:0.2f}'.format(fwIU))
print('pixel acc : {:0.2f}'.format(acc_overall))
per_class = [[classes[i], acc]
for i, acc in list(enumerate(acc_percls))]
per_class = np.array(per_class).flatten()
print(
('per cls acc :' + ('\n{:>14s} : {}') * 19).format(*per_class))
avg_train_acc = acc_overall
avg_train_loss_seg1 = np.mean(total_loss_seg_value1)
avg_train_loss_adv1 = np.mean(total_loss_adv_target_value1)
avg_train_loss_dis1 = np.mean(total_loss_D_value1)
avg_train_loss_con1 = np.mean(total_loss_con_value1)
avg_train_loss_seg2 = np.mean(total_loss_seg_value2)
avg_train_loss_adv2 = np.mean(total_loss_adv_target_value2)
avg_train_loss_dis2 = np.mean(total_loss_D_value2)
avg_train_loss_con2 = np.mean(total_loss_con_value2)
print('avg_train_acc :', avg_train_acc)
print('avg_train_loss_seg1 :', avg_train_loss_seg1)
print('avg_train_loss_adv1 :', avg_train_loss_adv1)
print('avg_train_loss_dis1 :', avg_train_loss_dis1)
print('avg_train_loss_con1 :', avg_train_loss_con1)
print('avg_train_loss_seg2 :', avg_train_loss_seg2)
print('avg_train_loss_adv2 :', avg_train_loss_adv2)
print('avg_train_loss_dis2 :', avg_train_loss_dis2)
print('avg_train_loss_con2 :', avg_train_loss_con2)
writer['train'].add_scalar('log/mIoU', mIoU, i_iter)
writer['train'].add_scalar('log/acc', avg_train_acc, i_iter)
writer['train'].add_scalar('log1/loss_seg', avg_train_loss_seg1,
i_iter)
writer['train'].add_scalar('log1/loss_adv', avg_train_loss_adv1,
i_iter)
writer['train'].add_scalar('log1/loss_dis', avg_train_loss_dis1,
i_iter)
writer['train'].add_scalar('log1/loss_con', avg_train_loss_con1,
i_iter)
writer['train'].add_scalar('log2/loss_seg', avg_train_loss_seg2,
i_iter)
writer['train'].add_scalar('log2/loss_adv', avg_train_loss_adv2,
i_iter)
writer['train'].add_scalar('log2/loss_dis', avg_train_loss_dis2,
i_iter)
writer['train'].add_scalar('log2/loss_con', avg_train_loss_con2,
i_iter)
hist = np.zeros((num_cls, num_cls))
total_loss_seg_value1 = []
total_loss_adv_target_value1 = []
total_loss_D_value1 = []
total_loss_con_value1 = []
total_loss_seg_value2 = []
total_loss_adv_target_value2 = []
total_loss_D_value2 = []
total_loss_con_value2 = []
fig = plt.figure(figsize=(15, 15))
labels = labels[0].cpu().numpy().astype(np.float32)
ax = fig.add_subplot(331)
ax.imshow(print_palette(Image.fromarray(labels).convert('L')))
ax.axis("off")
ax.set_title('labels')
ax = fig.add_subplot(337)
images = images_orig[0].cpu().numpy().transpose((1, 2, 0))
# images += IMG_MEAN
ax.imshow(images)
ax.axis("off")
ax.set_title('datas')
_, pred2 = torch.max(pred2, dim=1)
pred2 = pred2[0].cpu().numpy().astype(np.float32)
ax = fig.add_subplot(334)
ax.imshow(print_palette(Image.fromarray(pred2).convert('L')))
ax.axis("off")
ax.set_title('predicts')
labels_tar = labels_tar[0].cpu().numpy().astype(np.float32)
ax = fig.add_subplot(332)
ax.imshow(print_palette(Image.fromarray(labels_tar).convert('L')))
ax.axis("off")
ax.set_title('tar_labels')
ax = fig.add_subplot(338)
ax.imshow(images_tar_orig[0].cpu().numpy().transpose((1, 2, 0)))
ax.axis("off")
ax.set_title('tar_datas')
_, pred_target2 = torch.max(pred_target2, dim=1)
pred_target2 = pred_target2[0].cpu().numpy().astype(np.float32)
ax = fig.add_subplot(335)
ax.imshow(print_palette(
Image.fromarray(pred_target2).convert('L')))
ax.axis("off")
ax.set_title('tar_predicts')
print(policies[0], 'p1', rand_p1[0], 'p2', rand_p2[0],
'random_dir', random_dir[0])
psuedo_label2_thre = psuedo_label2_thre[0].cpu().numpy().astype(
np.float32)
ax = fig.add_subplot(333)
ax.imshow(
print_palette(
Image.fromarray(psuedo_label2_thre).convert('L')))
ax.axis("off")
ax.set_title('psuedo_labels')
ax = fig.add_subplot(339)
ax.imshow(images_aug_orig[0].cpu().numpy().transpose((1, 2, 0)))
ax.axis("off")
ax.set_title('aug_datas')
_, pred_target_aug2 = torch.max(pred_target_aug2, dim=1)
pred_target_aug2 = pred_target_aug2[0].cpu().numpy().astype(
np.float32)
ax = fig.add_subplot(336)
ax.imshow(
print_palette(Image.fromarray(pred_target_aug2).convert('L')))
ax.axis("off")
ax.set_title('aug_predicts')
# plt.show()
writer['train'].add_figure('image/',
fig,
global_step=i_iter,
close=True)
if i_iter % 500 == 0:
loss1 = []
loss2 = []
for test_i, batch in enumerate(test_loader):
images, labels = batch
images_orig = images
images = transform_batch(images, normalize_transform)
images = Variable(images).cuda(args.gpu)
pred1, pred2 = model(images)
pred1 = interp_test(pred1)
pred1 = pred1.detach()
pred2 = interp_test(pred2)
pred2 = pred2.detach()
loss_seg1 = loss_calc(pred1, labels, args.gpu)
loss_seg2 = loss_calc(pred2, labels, args.gpu)
loss1.append(loss_seg1.item())
loss2.append(loss_seg2.item())
hist += fast_hist(
labels.cpu().numpy().flatten().astype(int),
torch.argmax(pred2,
dim=1).cpu().numpy().flatten().astype(int),
num_cls)
print('test')
fig = plt.figure(figsize=(15, 15))
labels = labels[-1].cpu().numpy().astype(np.float32)
ax = fig.add_subplot(311)
ax.imshow(print_palette(Image.fromarray(labels).convert('L')))
ax.axis("off")
ax.set_title('labels')
ax = fig.add_subplot(313)
ax.imshow(images_orig[-1].cpu().numpy().transpose((1, 2, 0)))
ax.axis("off")
ax.set_title('datas')
_, pred2 = torch.max(pred2, dim=1)
pred2 = pred2[-1].cpu().numpy().astype(np.float32)
ax = fig.add_subplot(312)
ax.imshow(print_palette(Image.fromarray(pred2).convert('L')))
ax.axis("off")
ax.set_title('predicts')
# plt.show()
writer['test'].add_figure('test_image/',
fig,
global_step=i_iter,
close=True)
acc_overall, acc_percls, iu, fwIU = result_stats(hist)
mIoU = np.mean(iu)
per_class = [[classes[i], acc] for i, acc in list(enumerate(iu))]
per_class = np.array(per_class).flatten()
print(
('per cls IoU :' + ('\n{:>14s} : {}') * 19).format(*per_class))
print('mIoU : {:0.2f}'.format(mIoU))
print('fwIoU : {:0.2f}'.format(fwIU))
print('pixel acc : {:0.2f}'.format(acc_overall))
per_class = [[classes[i], acc]
for i, acc in list(enumerate(acc_percls))]
per_class = np.array(per_class).flatten()
print(
('per cls acc :' + ('\n{:>14s} : {}') * 19).format(*per_class))
avg_test_loss1 = np.mean(loss1)
avg_test_loss2 = np.mean(loss2)
avg_test_acc = acc_overall
print('avg_test_loss2 :', avg_test_loss1)
print('avg_test_loss1 :', avg_test_loss2)
print('avg_test_acc :', avg_test_acc)
writer['test'].add_scalar('log1/loss_seg', avg_test_loss1, i_iter)
writer['test'].add_scalar('log2/loss_seg', avg_test_loss2, i_iter)
writer['test'].add_scalar('log/acc', avg_test_acc, i_iter)
writer['test'].add_scalar('log/mIoU', mIoU, i_iter)
hist = np.zeros((num_cls, num_cls))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D2.pth'))
break
if max_mIoU < mIoU:
max_mIoU = mIoU
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + 'best_iter' + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + 'best_iter' + '_D1.pth'))
torch.save(
model_D2.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + 'best_iter' + '_D2.pth'))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {
k: v
for k, v in saved_state_dict.items() if k in model_dict
}
model_dict.update(saved_state_dict)
###
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
log_dir = args.save
if not os.path.isdir(log_dir):
os.mkdir(log_dir)
exp_name = datetime.datetime.now().strftime("%H%M%S-%Y%m%d")
log_dir = os.path.join(log_dir, exp_name)
writer = SummaryWriter(log_dir)
# testloader = data.DataLoader(SyntheticSmokeTrain(args={}, dataset_limit=-1, #args.num_steps * args.iter_size * args.batch_size,
# image_shape=(360,640), dataset_mean=IMG_MEAN),
# batch_size=1, shuffle=True, pin_memory=True)
testloader = data.DataLoader(SmokeDataset(image_size=(640, 360),
dataset_mean=IMG_MEAN),
batch_size=1,
shuffle=True,
pin_memory=True)
# testloader = data.DataLoader(SimpleSmokeTrain(args = {}, image_size=(640,360), dataset_mean=IMG_MEAN),
# batch_size=1, shuffle=True, pin_memory=True)
# testloader = data.DataLoader(cityscapesDataSet(args.data_dir, args.data_list, crop_size=(1024, 512), mean=IMG_MEAN, scale=False, mirror=False, set=args.set),
# batch_size=1, shuffle=False, pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(640, 360),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(640, 360),
mode='bilinear',
align_corners=True)
count = 0
iou_sum_fg = 0
iou_count_fg = 0
iou_sum_bg = 0
iou_count_bg = 0
for index, batch in enumerate(testloader):
if (index + 1) % 100 == 0:
print('%d processd' % index)
# print("Processed {}/{}".format(index, len(testloader)))
# if count > 5:
# break
image, label, name = batch
if args.model == 'DeeplabMulti':
with torch.no_grad():
output1, output2 = model(Variable(image).cuda(gpu0))
# print(output1.shape)
# print(output2.shape)
output = interp(output2).cpu()
orig_output = output.detach().clone()
output = output.data[0].numpy()
# output = (output > 0.5).astype(np.uint8)*255
# print(np.all(output==0), np.all(output==255))
# print(np.min(output), np.max(output))
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
with torch.no_grad():
output = model(Variable(image).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
classes_seen = set(output.ravel().tolist())
# print(classes_seen)
# print(output.shape, name[0])
output_col = colorize_mask(output)
output = Image.fromarray(output)
# print("name", name)
name = name[0]
# name = name[0].split('/')[-1]
if len(classes_seen) > 1:
count += 1
print(classes_seen)
print(Counter(np.asarray(output).ravel()))
image = image.squeeze()
for c in range(3):
image[c, :, :] += IMG_MEAN[c]
# image2[c,:,:] += IMG_MEAN[2-c]
image = (image - image.min()) / (image.max() - image.min())
image = image[[2, 1, 0], :, :]
print(image.shape, image.min(), image.max())
output.save(os.path.join(args.save, name + '.png'))
output_col.save(os.path.join(args.save, name + '_color.png'))
# output.save('%s/%s.png' % (args.save, name))
# output_col.save('%s/%s_color.png' % (args.save, name))#.split('.')[0]))
output_argmaxs = torch.argmax(orig_output.squeeze(), dim=0)
mask1 = (output_argmaxs == 0).float() * 255
label = label.squeeze()
iou_fg = iou_pytorch(mask1, label)
print("foreground IoU", iou_fg)
iou_sum_fg += iou_fg
iou_count_fg += 1
mask2 = (output_argmaxs > 0).float() * 255
label2 = label.max() - label
iou_bg = iou_pytorch(mask2, label2)
print("IoU for background: ", iou_bg)
iou_sum_bg += iou_bg
iou_count_bg += 1
writer.add_images(f'input_images',
tf.resize(image[[2, 1, 0]], [1080, 1920]),
index,
dataformats='CHW')
print("shape of label", label.shape)
label_reshaped = tf.resize(label.unsqueeze(0),
[1080, 1920]).squeeze()
print("label reshaped: ", label_reshaped.shape)
writer.add_images(f'labels',
label_reshaped,
index,
dataformats='HW')
writer.add_images(
f'output/1',
255 - np.asarray(tf.resize(output, [1080, 1920])) * 255,
index,
dataformats='HW')
# writer.add_images(f'output/1',np.asarray(output)*255, index,dataformats='HW')
# writer.add_images(f'output/2',np.asarray(output_col), index, dataformats='HW')
writer.add_scalar(f'iou/smoke', iou_fg, index)
writer.add_scalar(f'iou/background', iou_bg, index)
writer.add_scalar(f'iou/mean', (iou_bg + iou_fg) / 2, index)
writer.flush()
if iou_count_fg > 0:
print("Mean IoU, foreground: {}".format(iou_sum_fg / iou_count_fg))
print("Mean IoU, background: {}".format(iou_sum_bg / iou_count_bg))
print("Mean IoU, averaged over classes: {}".format(
(iou_sum_fg + iou_sum_bg) / (iou_count_fg + iou_count_bg)))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
if not os.path.exists(args.save):
os.makedirs(args.save)
nyu_nyu_dict = {11:255, 13:255, 15:255, 17:255, 19:255, 20:255, 21: 255, 23: 255,
24:255, 25:255, 26:255, 27:255, 28:255, 29:255, 31:255, 32:255, 33:255}
nyu_nyu_map = lambda x: nyu_nyu_dict.get(x+1,x)
nyu_nyu_map = np.vectorize(nyu_nyu_map)
args.nyu_nyu_map = nyu_nyu_map
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http' :
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {k: v for k, v in saved_state_dict.items() if k in model_dict}
model_dict.update(saved_state_dict)
###
model.load_state_dict(saved_state_dict)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
metrics = StreamSegMetrics(args.num_classes)
metrics_remap = StreamSegMetrics(args.num_classes)
ignore_label = 255
value_scale = 255
mean = [0.485, 0.456, 0.406]
mean = [item * value_scale for item in mean]
std = [0.229, 0.224, 0.225]
std = [item * value_scale for item in std]
val_transform = transforms.Compose([
# et.ExtResize( 512 ),
transforms.Crop([args.height+1, args.width+1], crop_type='center', padding=IMG_MEAN, ignore_label=ignore_label),
transforms.ToTensor(),
transforms.Normalize(mean=IMG_MEAN,
std=[1, 1, 1]),
])
val_dst = NYU(root=args.data_dir, opt=args,
split='val', transform=val_transform,
imWidth = args.width, imHeight = args.height, phase="TEST",
randomize = False)
print("Dset Length {}".format(len(val_dst)))
testloader = data.DataLoader(val_dst,
batch_size=1, shuffle=False, pin_memory=True)
interp = nn.Upsample(size=(args.height+1, args.width+1), mode='bilinear', align_corners=True)
metrics.reset()
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, targets, name = batch
image = image.to(device)
print(index)
if args.model == 'DeeplabMulti':
output1, output2 = model(image)
output = interp(output2).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output = model(image)
output = interp(output).cpu().data[0].numpy()
targets = targets.cpu().numpy()
output = output.transpose(1,2,0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
preds = output[None,:,:]
#input_ = image.cpu().numpy()[0].transpose(1,2,0) + np.array(IMG_MEAN)
metrics.update(targets, preds)
targets = args.nyu_nyu_map(targets)
preds = args.nyu_nyu_map(preds)
metrics_remap.update(targets,preds)
#input_ = Image.fromarray(input_.astype(np.uint8))
#output_col = colorize_mask(output)
#output = Image.fromarray(output)
#name = name[0].split('/')[-1]
#input_.save('%s/%s' % (args.save, name))
#output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
print(metrics.get_results())
print(metrics_remap.get_results())
def main():
"""Create the model and start the evaluation process."""
for i in range(26, 126): #1 126
model_path = './snapshots/GTA2Cityscapes/GTA5_{0:d}.pth'.format(i *
2000)
save_path = './result/GTA2Cityscapes_{0:d}'.format(i * 2000)
args = get_arguments()
gpu0 = args.gpu
if not os.path.exists(save_path):
os.makedirs(save_path)
if args.model == 'ResNet':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes)
saved_state_dict = torch.load(model_path)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
with torch.no_grad():
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
_, output2 = model(Variable(image).cuda(gpu0))
output = interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (save_path, name))
output_col.save('%s/%s_color.png' %
(save_path, name.split('.')[0]))
print(save_path)
def main():
"""Create the model and start the training."""
w, h = map(int, args.input_size_source.split(','))
input_size_source = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
# Create network
if args.model == 'ResNet':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
# model.load_state_dict(saved_state_dict)
elif args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes,
pretrained=True,
vgg16_caffe_path=args.restore_from)
# saved_state_dict = torch.load(args.restore_from)
# model.load_state_dict(saved_state_dict)
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
model.train()
model.cuda(args.gpu)
cudnn.benchmark = True
#Discrimintator setting
model_D = FCDiscriminator(num_classes=args.num_classes)
model_D.train()
model_D.cuda(args.gpu)
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
# labels for adversarial training
source_adv_label = 0
target_adv_label = 1
#Dataloader
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(GTA5DataSet(args.translated_data_dir,
args.data_list,
max_iters=args.num_steps *
args.iter_size * args.batch_size,
crop_size=input_size_source,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
style_trainloader = data.DataLoader(GTA5DataSet(
args.stylized_data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_source,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
style_trainloader_iter = enumerate(style_trainloader)
if STAGE == 1:
targetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
else:
#Dataloader for self-training
targetloader = data.DataLoader(cityscapesDataSetLabel(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
mean=IMG_MEAN,
set=args.set,
label_folder='Path to generated pseudo labels'),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
interp = nn.Upsample(size=(input_size_source[1], input_size_source[0]),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# load checkpoint
model, model_D, optimizer, start_iter = load_checkpoint(
model,
model_D,
optimizer,
filename=args.snapshot_dir + 'checkpoint_' + CHECKPOINT + '.pth.tar')
for i_iter in range(start_iter, args.num_steps):
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
#train segementation network
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
if STAGE == 1:
if i_iter % 2 == 0:
_, batch = next(trainloader_iter)
else:
_, batch = next(style_trainloader_iter)
else:
_, batch = next(trainloader_iter)
image_source, label, _, _ = batch
image_source = Variable(image_source).cuda(args.gpu)
pred_source = model(image_source)
pred_source = interp(pred_source)
loss_seg_source = loss_calc(pred_source, label, args.gpu)
loss_seg_source_value = loss_seg_source.item()
loss_seg_source.backward()
if STAGE == 2:
# train with target
_, batch = next(targetloader_iter)
image_target, target_label, _, _ = batch
image_target = Variable(image_target).cuda(args.gpu)
pred_target = model(image_target)
pred_target = interp_target(pred_target)
#target segmentation loss
loss_seg_target = loss_calc(pred_target,
target_label,
gpu=args.gpu)
loss_seg_target.backward()
# optimize
optimizer.step()
if STAGE == 1:
# train with target
_, batch = next(targetloader_iter)
image_target, _, _ = batch
image_target = Variable(image_target).cuda(args.gpu)
pred_target = model(image_target)
pred_target = interp_target(pred_target)
#output-level adversarial training
D_output_target = model_D(F.softmax(pred_target))
loss_adv = bce_loss(
D_output_target,
Variable(
torch.FloatTensor(D_output_target.data.size()).fill_(
source_adv_label)).cuda(args.gpu))
loss_adv = loss_adv * args.lambda_adv
loss_adv.backward()
#train discriminator
for param in model_D.parameters():
param.requires_grad = True
pred_source = pred_source.detach()
pred_target = pred_target.detach()
D_output_source = model_D(F.softmax(pred_source))
D_output_target = model_D(F.softmax(pred_target))
loss_D_source = bce_loss(
D_output_source,
Variable(
torch.FloatTensor(D_output_source.data.size()).fill_(
source_adv_label)).cuda(args.gpu))
loss_D_target = bce_loss(
D_output_target,
Variable(
torch.FloatTensor(D_output_target.data.size()).fill_(
target_adv_label)).cuda(args.gpu))
loss_D_source = loss_D_source / 2
loss_D_target = loss_D_target / 2
loss_D_source.backward()
loss_D_target.backward()
#optimize
optimizer_D.step()
print('exp = {}'.format(args.snapshot_dir))
print('iter = {0:8d}/{1:8d}, loss_seg_source = {2:.5f}'.format(
i_iter, args.num_steps, loss_seg_source_value))
if i_iter % args.save_pred_every == 0:
print('taking snapshot ...')
state = {
'iter': i_iter,
'model': model.state_dict(),
'model_D': model_D.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(
state,
osp.join(args.snapshot_dir,
'checkpoint_' + str(i_iter) + '.pth.tar'))
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_D_' + str(i_iter) + '.pth'))
cityscapes_eval_dir = osp.join(args.cityscapes_eval_dir,
str(i_iter))
if not os.path.exists(cityscapes_eval_dir):
os.makedirs(cityscapes_eval_dir)
eval(osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'),
cityscapes_eval_dir, i_iter)
iou19, iou13, iou = compute_mIoU(cityscapes_eval_dir, i_iter)
outputfile = open(args.output_file, 'a')
outputfile.write(
str(i_iter) + '\t' + str(iou19) + '\t' +
str(iou.replace('\n', ' ')) + '\n')
outputfile.close()
def main():
seed = 1338
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
"""Create the model and start the evaluation process."""
args = get_arguments()
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
# if args.restore_from[:4] == 'http' :
# saved_state_dict = model_zoo.load_url(args.restore_from)
# else:
# saved_state_dict = torch.load(args.restore_from)
for files in range(int(args.num_steps_stop / args.save_pred_every)):
print('Step: ', (files + 1) * args.save_pred_every)
if SOURCE_ONLY:
saved_state_dict = torch.load('./snapshots/source_only/GTA5_' +
str((files + 1) *
args.save_pred_every) + '.pth')
else:
if args.level == 'single-level':
saved_state_dict = torch.load(
'./snapshots/single_level/GTA5_' +
str((files + 1) * args.save_pred_every) + '.pth')
elif args.level == 'multi-level':
saved_state_dict = torch.load('./snapshots/multi_level/GTA5_' +
str((files + 1) *
args.save_pred_every) +
'.pth')
else:
raise NotImplementedError(
'level choice {} is not implemented'.format(args.level))
### for running different versions of pytorch
model_dict = model.state_dict()
saved_state_dict = {
k: v
for k, v in saved_state_dict.items() if k in model_dict
}
model_dict.update(saved_state_dict)
###
model.load_state_dict(saved_state_dict)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
if args.multi_gpu:
model = nn.DataParallel(model)
model.eval()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
image = image.to(device)
if args.model == 'DeeplabMulti':
output1, output2 = model(image)
output = interp(output2).cpu().data[0].numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output = model(image)
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
if SOURCE_ONLY:
if not os.path.exists(
os.path.join(
args.save, 'source_only', 'step' + str(
(files + 1) * args.save_pred_every))):
os.makedirs(
os.path.join(
args.save, 'source_only', 'step' + str(
(files + 1) * args.save_pred_every)))
output.save(
os.path.join(
args.save, 'source_only', 'step' + str(
(files + 1) * args.save_pred_every), name))
output_col.save(
os.path.join(
args.save, 'source_only', 'step' + str(
(files + 1) * args.save_pred_every),
name.split('.')[0] + '_color.png'))
else:
if args.level == 'single-level':
if not os.path.exists(
os.path.join(
args.save, 'single_level', 'step' + str(
(files + 1) * args.save_pred_every))):
os.makedirs(
os.path.join(
args.save, 'single_level', 'step' + str(
(files + 1) * args.save_pred_every)))
output.save(
os.path.join(
args.save, 'single_level', 'step' + str(
(files + 1) * args.save_pred_every), name))
output_col.save(
os.path.join(
args.save, 'single_level', 'step' + str(
(files + 1) * args.save_pred_every),
name.split('.')[0] + '_color.png'))
elif args.level == 'multi-level':
if not os.path.exists(
os.path.join(
args.save, 'multi_level', 'step' + str(
(files + 1) * args.save_pred_every))):
os.makedirs(
os.path.join(
args.save, 'multi_level', 'step' + str(
(files + 1) * args.save_pred_every)))
output.save(
os.path.join(
args.save, 'multi_level', 'step' + str(
(files + 1) * args.save_pred_every), name))
output_col.save(
os.path.join(
args.save, 'multi_level', 'step' + str(
(files + 1) * args.save_pred_every),
name.split('.')[0] + '_color.png'))
else:
raise NotImplementedError(
'level choice {} is not implemented'.format(
args.level))
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
w, h = map(int, args.input_size.split(','))
config_path = os.path.join(os.path.dirname(args.restore_from), 'opts.yaml')
with open(config_path, 'r') as stream:
config = yaml.load(stream)
args.model = config['model']
print('ModelType:%s' % args.model)
print('NormType:%s' % config['norm_style'])
gpu0 = args.gpu
batchsize = args.batchsize
model_name = os.path.basename(os.path.dirname(args.restore_from))
#args.save += model_name
if not os.path.exists(args.save):
os.makedirs(args.save)
confidence_path = os.path.join(args.save, 'submit/confidence')
label_path = os.path.join(args.save, 'submit/labelTrainIds')
label_invalid_path = os.path.join(args.save,
'submit/labelTrainIds_invalid')
for path in [confidence_path, label_path, label_invalid_path]:
if not os.path.exists(path):
os.makedirs(path)
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes,
use_se=config['use_se'],
train_bn=False,
norm_style=config['norm_style'])
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
try:
model.load_state_dict(saved_state_dict)
except:
model = torch.nn.DataParallel(model)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(DarkZurichDataSet(args.data_dir,
args.data_list,
crop_size=(h, w),
resize_size=(w, h),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
scale = 1.25
testloader2 = data.DataLoader(DarkZurichDataSet(
args.data_dir,
args.data_list,
crop_size=(round(h * scale), round(w * scale)),
resize_size=(round(w * scale), round(h * scale)),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(1080, 1920),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(1080, 1920), mode='bilinear')
sm = torch.nn.Softmax(dim=1)
log_sm = torch.nn.LogSoftmax(dim=1)
kl_distance = nn.KLDivLoss(reduction='none')
prior = np.load('./utils/prior_all.npy').transpose(
(2, 0, 1))[np.newaxis, :, :, :]
prior = torch.from_numpy(prior)
for index, img_data in enumerate(zip(testloader, testloader2)):
batch, batch2 = img_data
image, _, name = batch
image2, _, name2 = batch2
inputs = image.cuda()
inputs2 = image2.cuda()
print('\r>>>>Extracting feature...%04d/%04d' %
(index * batchsize, args.batchsize * len(testloader)),
end='')
if args.model == 'DeepLab':
with torch.no_grad():
output1, output2 = model(inputs)
output_batch = interp(sm(0.5 * output1 + output2))
heatmap_batch = torch.sum(kl_distance(log_sm(output1),
sm(output2)),
dim=1)
output1, output2 = model(fliplr(inputs))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
del output1, output2, inputs
output1, output2 = model(inputs2)
output_batch += interp(sm(0.5 * output1 + output2))
output1, output2 = model(fliplr(inputs2))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
del output1, output2, inputs2
ratio = 0.95
output_batch = output_batch.cpu() / 4
# output_batch = output_batch *(ratio + (1 - ratio) * prior)
output_batch = output_batch.data.numpy()
heatmap_batch = heatmap_batch.cpu().data.numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output_batch = model(Variable(image).cuda())
output_batch = interp(output_batch).cpu().data.numpy()
output_batch = output_batch.transpose(0, 2, 3, 1)
score_batch = np.max(output_batch, axis=3)
output_batch = np.asarray(np.argmax(output_batch, axis=3),
dtype=np.uint8)
threshold = 0.3274
for i in range(output_batch.shape[0]):
output_single = output_batch[i, :, :]
output_col = colorize_mask(output_single)
output = Image.fromarray(output_single)
name_tmp = name[i].split('/')[-1]
dir_name = name[i].split('/')[-2]
save_path = args.save + '/' + dir_name
if not os.path.isdir(save_path):
os.mkdir(save_path)
output.save('%s/%s' % (save_path, name_tmp))
print('%s/%s' % (save_path, name_tmp))
output_col.save('%s/%s_color.png' %
(save_path, name_tmp.split('.')[0]))
# heatmap_tmp = heatmap_batch[i,:,:]/np.max(heatmap_batch[i,:,:])
# fig = plt.figure()
# plt.axis('off')
# heatmap = plt.imshow(heatmap_tmp, cmap='viridis')
# fig.colorbar(heatmap)
# fig.savefig('%s/%s_heatmap.png' % (save_path, name_tmp.split('.')[0]))
if args.set == 'test' or args.set == 'val':
# label
output.save('%s/%s' % (label_path, name_tmp))
# label invalid
output_single[score_batch[i, :, :] < threshold] = 255
output = Image.fromarray(output_single)
output.save('%s/%s' % (label_invalid_path, name_tmp))
# conficence
confidence = score_batch[i, :, :] * 65535
confidence = np.asarray(confidence, dtype=np.uint16)
print(confidence.min(), confidence.max())
iio.imwrite('%s/%s' % (confidence_path, name_tmp), confidence)
return args.save
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
elif args.model == 'DeeplabVGGBN':
deeplab_vggbn.BatchNorm = SyncBatchNorm2d
model = deeplab_vggbn.DeeplabVGGBN(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict, strict=False)
print(model)
device = torch.device("cuda" if not args.cpu else "cpu")
model = model.to(device)
model.eval()
testloader = data.DataLoader(BDDDataSet(args.data_dir,
args.data_list,
crop_size=(960, 540),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True) # 960 540
interp = nn.Upsample(size=(720, 1280), mode='bilinear', align_corners=True)
if args.save_confidence:
select = open('list.txt', 'w')
c_list = []
for index, batch in enumerate(testloader):
if index % 10 == 0:
print('%d processd' % index)
image, _, name = batch
image = image.to(device)
output = model(image)
if args.save_confidence:
confidence = get_confidence(output)
confidence = confidence.cpu().item()
c_list.append([confidence, name])
name = name[0].split('/')[-1]
save_path = '%s/%s_c.txt' % (args.save, name.split('.')[0])
record = open(save_path, 'w')
record.write('%.5f' % confidence)
record.close()
else:
name = name[0].split('/')[-1]
output = interp(output).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
output.save('%s/%s' % (args.save, name[:-4] + '.png'))
output_col.save('%s/%s_color.png' % (args.save, name.split('.')[0]))
def takeFirst(elem):
return elem[0]
if args.save_confidence:
c_list.sort(key=takeFirst, reverse=True)
length = len(c_list)
for i in range(length // 3):
print(c_list[i][0])
print(c_list[i][1])
select.write(c_list[i][1][0])
select.write('\n')
select.close()
print(args.save)
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = args.gpu
batchsize = args.batchsize
model_name = os.path.basename(os.path.dirname(args.restore_from))
args.save += model_name
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes,
train_bn=False,
norm_style='in')
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
try:
model.load_state_dict(saved_state_dict)
except:
model = torch.nn.DataParallel(model)
model.load_state_dict(saved_state_dict)
model.eval()
model.cuda()
testloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
crop_size=(640, 1280),
resize_size=(1280, 640),
mean=IMG_MEAN,
scale=False,
mirror=False),
batch_size=batchsize,
shuffle=False,
pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(640, 1280),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(640, 1280), mode='bilinear')
sm = torch.nn.Softmax(dim=1)
for index, batch in enumerate(testloader):
if (index * batchsize) % 100 == 0:
print('%d processd' % (index * batchsize))
image, _, _, name = batch
print(image.shape)
inputs = Variable(image).cuda()
if args.model == 'DeeplabMulti':
output1, output2 = model(inputs)
output_batch = interp(sm(0.5 * output1 +
output2)).cpu().data.numpy()
#output1, output2 = model(fliplr(inputs))
#output2 = fliplr(output2)
#output_batch += interp(output2).cpu().data.numpy()
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output_batch = model(Variable(image).cuda())
output_batch = interp(output_batch).cpu().data.numpy()
output_batch = output_batch.transpose(0, 2, 3, 1)
output_batch = np.asarray(np.argmax(output_batch, axis=3),
dtype=np.uint8)
for i in range(output_batch.shape[0]):
output = output_batch[i, :, :]
output_col = colorize_mask(output)
output = Image.fromarray(output)
name_tmp = name[i].split('/')[-1]
output.save('%s/%s' % (args.save, name_tmp))
output_col.save('%s/%s_color.png' %
(args.save, name_tmp.split('.')[0]))
return args.save
def main():
"""Create the model and start the training."""
global args
args = get_arguments()
if args.dist:
init_dist(args.launcher, backend=args.backend)
world_size = 1
rank = 0
if args.dist:
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device("cuda" if not args.cpu else "cpu")
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
# Create network
if args.model == 'Deeplab':
model = DeeplabMulti(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from, strict=False)
new_params = model.state_dict().copy()
for i in saved_state_dict:
i_parts = i.split('.')
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
model.load_state_dict(new_params)
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict, strict=False)
elif args.model == 'DeeplabVGGBN':
deeplab_vggbn.BatchNorm = SyncBatchNorm2d
model = deeplab_vggbn.DeeplabVGGBN(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict, strict=False)
del saved_state_dict
model.train()
model.to(device)
if args.dist:
broadcast_params(model)
if rank == 0:
print(model)
cudnn.benchmark = True
# init D
model_D1 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D2 = FCDiscriminator(num_classes=args.num_classes).to(device)
model_D1.train()
model_D1.to(device)
if args.dist:
broadcast_params(model_D1)
if args.restore_D is not None:
D_dict = torch.load(args.restore_D)
model_D1.load_state_dict(D_dict, strict=False)
del D_dict
model_D2.train()
model_D2.to(device)
if args.dist:
broadcast_params(model_D2)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
train_data = GTA5BDDDataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps * args.iter_size *
args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN)
train_sampler = None
if args.dist:
train_sampler = DistributedSampler(train_data)
trainloader = data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=False if train_sampler else True,
num_workers=args.num_workers,
pin_memory=False,
sampler=train_sampler)
trainloader_iter = enumerate(cycle(trainloader))
target_data = BDDDataSet(args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size *
args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set)
target_sampler = None
if args.dist:
target_sampler = DistributedSampler(target_data)
targetloader = data.DataLoader(target_data,
batch_size=args.batch_size,
shuffle=False if target_sampler else True,
num_workers=args.num_workers,
pin_memory=False,
sampler=target_sampler)
targetloader_iter = enumerate(cycle(targetloader))
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D1 = optim.Adam(model_D1.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D1.zero_grad()
optimizer_D2 = optim.Adam(model_D2.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D2.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
#interp = nn.Upsample(size=(input_size[1], input_size[0]), mode='bilinear', align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
# set up tensor board
if args.tensorboard and rank == 0:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
torch.cuda.empty_cache()
for i_iter in range(args.num_steps):
loss_seg_value1 = 0
loss_adv_target_value1 = 0
loss_D_value1 = 0
loss_seg_value2 = 0
loss_adv_target_value2 = 0
loss_D_value2 = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D1.zero_grad()
optimizer_D2.zero_grad()
adjust_learning_rate_D(optimizer_D1, i_iter)
adjust_learning_rate_D(optimizer_D2, i_iter)
for sub_i in range(args.iter_size):
# train G
# don't accumulate grads in D
for param in model_D1.parameters():
param.requires_grad = False
for param in model_D2.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.__next__()
images, labels, size, _ = batch
images = images.to(device)
labels = labels.long().to(device)
interp = nn.Upsample(size=(size[1], size[0]),
mode='bilinear',
align_corners=True)
pred1 = model(images)
pred1 = interp(pred1)
loss_seg1 = seg_loss(pred1, labels)
loss = loss_seg1
# proper normalization
loss = loss / args.iter_size / world_size
loss.backward()
loss_seg_value1 += loss_seg1.item() / args.iter_size
_, batch = targetloader_iter.__next__()
# train with target
images, _, _ = batch
images = images.to(device)
pred_target1 = model(images)
pred_target1 = interp_target(pred_target1)
D_out1 = model_D1(F.softmax(pred_target1))
loss_adv_target1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss = args.lambda_adv_target1 * loss_adv_target1
loss = loss / args.iter_size / world_size
loss.backward()
loss_adv_target_value1 += loss_adv_target1.item() / args.iter_size
# train D
# bring back requires_grad
for param in model_D1.parameters():
param.requires_grad = True
for param in model_D2.parameters():
param.requires_grad = True
# train with source
pred1 = pred1.detach()
D_out1 = model_D1(F.softmax(pred1))
loss_D1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(source_label).to(device))
loss_D1 = loss_D1 / args.iter_size / 2 / world_size
loss_D1.backward()
loss_D_value1 += loss_D1.item()
# train with target
pred_target1 = pred_target1.detach()
D_out1 = model_D1(F.softmax(pred_target1))
loss_D1 = bce_loss(
D_out1,
torch.FloatTensor(
D_out1.data.size()).fill_(target_label).to(device))
loss_D1 = loss_D1 / args.iter_size / 2 / world_size
loss_D1.backward()
if args.dist:
average_gradients(model)
average_gradients(model_D1)
average_gradients(model_D2)
loss_D_value1 += loss_D1.item()
optimizer.step()
optimizer_D1.step()
if rank == 0:
if args.tensorboard:
scalar_info = {
'loss_seg1': loss_seg_value1,
'loss_seg2': loss_seg_value2,
'loss_adv_target1': loss_adv_target_value1,
'loss_adv_target2': loss_adv_target_value2,
'loss_D1': loss_D_value1 * world_size,
'loss_D2': loss_D_value2 * world_size,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg1 = {2:.3f} loss_seg2 = {3:.3f} loss_adv1 = {4:.3f}, loss_adv2 = {5:.3f} loss_D1 = {6:.3f} loss_D2 = {7:.3f}'
.format(i_iter, args.num_steps, loss_seg_value1,
loss_seg_value2, loss_adv_target_value1,
loss_adv_target_value2, loss_D_value1, loss_D_value2))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D1.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D1.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(i_iter) + '_D1.pth'))
print(args.snapshot_dir)
if args.tensorboard and rank == 0:
writer.close()
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
config_path = os.path.join(os.path.dirname(args.restore_from), 'opts.yaml')
with open(config_path, 'r') as stream:
config = yaml.load(stream)
args.model = config['model']
print('ModelType:%s' % args.model)
print('NormType:%s' % config['norm_style'])
gpu0 = args.gpu
batchsize = args.batchsize
model_name = os.path.basename(os.path.dirname(args.restore_from))
args.save += model_name
if not os.path.exists(args.save):
os.makedirs(args.save)
if args.model == 'DeepLab':
model = DeeplabMulti(num_classes=args.num_classes,
use_se=config['use_se'],
train_bn=False,
norm_style=config['norm_style'])
elif args.model == 'Oracle':
model = Res_Deeplab(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_ORC
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
try:
model.load_state_dict(saved_state_dict)
except:
model = torch.nn.DataParallel(model)
model.load_state_dict(saved_state_dict)
#model = torch.nn.DataParallel(model)
model.eval()
model.cuda(gpu0)
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(512, 1024),
resize_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
scale = 1.25
testloader2 = data.DataLoader(cityscapesDataSet(
args.data_dir,
args.data_list,
crop_size=(round(512 * scale), round(1024 * scale)),
resize_size=(round(1024 * scale), round(512 * scale)),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
scale = 0.9
testloader3 = data.DataLoader(cityscapesDataSet(
args.data_dir,
args.data_list,
crop_size=(round(512 * scale), round(1024 * scale)),
resize_size=(round(1024 * scale), round(512 * scale)),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=batchsize,
shuffle=False,
pin_memory=True,
num_workers=4)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
else:
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
sm = torch.nn.Softmax(dim=1)
log_sm = torch.nn.LogSoftmax(dim=1)
kl_distance = nn.KLDivLoss(reduction='none')
for index, img_data in enumerate(zip(testloader, testloader2,
testloader3)):
batch, batch2, batch3 = img_data
image, _, _, name = batch
image2, _, _, name2 = batch2
#image3, _, _, name3 = batch3
inputs = image.cuda()
inputs2 = image2.cuda()
#inputs3 = Variable(image3).cuda()
print('\r>>>>Extracting feature...%03d/%03d' %
(index * batchsize, NUM_STEPS),
end='')
if args.model == 'DeepLab':
with torch.no_grad():
output1, output2 = model(inputs)
output_batch = interp(sm(0.5 * output1 + output2))
heatmap_output1, heatmap_output2 = output1, output2
#output_batch = interp(sm(output1))
#output_batch = interp(sm(output2))
output1, output2 = model(fliplr(inputs))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
heatmap_output1, heatmap_output2 = heatmap_output1 + output1, heatmap_output2 + output2
#output_batch += interp(sm(output1))
#output_batch += interp(sm(output2))
del output1, output2, inputs
output1, output2 = model(inputs2)
output_batch += interp(sm(0.5 * output1 + output2))
#output_batch += interp(sm(output1))
#output_batch += interp(sm(output2))
output1, output2 = model(fliplr(inputs2))
output1, output2 = fliplr(output1), fliplr(output2)
output_batch += interp(sm(0.5 * output1 + output2))
#output_batch += interp(sm(output1))
#output_batch += interp(sm(output2))
del output1, output2, inputs2
output_batch = output_batch.cpu().data.numpy()
heatmap_batch = torch.sum(kl_distance(log_sm(heatmap_output1),
sm(heatmap_output2)),
dim=1)
heatmap_batch = torch.log(
1 + 10 * heatmap_batch) # for visualization
heatmap_batch = heatmap_batch.cpu().data.numpy()
#output1, output2 = model(inputs3)
#output_batch += interp(sm(0.5* output1 + output2)).cpu().data.numpy()
#output1, output2 = model(fliplr(inputs3))
#output1, output2 = fliplr(output1), fliplr(output2)
#output_batch += interp(sm(0.5 * output1 + output2)).cpu().data.numpy()
#del output1, output2, inputs3
elif args.model == 'DeeplabVGG' or args.model == 'Oracle':
output_batch = model(Variable(image).cuda())
output_batch = interp(output_batch).cpu().data.numpy()
output_batch = output_batch.transpose(0, 2, 3, 1)
scoremap_batch = np.asarray(np.max(output_batch, axis=3))
output_batch = np.asarray(np.argmax(output_batch, axis=3),
dtype=np.uint8)
output_iterator = []
heatmap_iterator = []
scoremap_iterator = []
for i in range(output_batch.shape[0]):
output_iterator.append(output_batch[i, :, :])
heatmap_iterator.append(heatmap_batch[i, :, :] /
np.max(heatmap_batch[i, :, :]))
scoremap_iterator.append(1 - scoremap_batch[i, :, :] /
np.max(scoremap_batch[i, :, :]))
name_tmp = name[i].split('/')[-1]
name[i] = '%s/%s' % (args.save, name_tmp)
with Pool(4) as p:
p.map(save, zip(output_iterator, name))
p.map(save_heatmap, zip(heatmap_iterator, name))
p.map(save_scoremap, zip(scoremap_iterator, name))
del output_batch
return args.save
def main():
"""Create the model and start the evaluation process."""
args = get_arguments()
gpu0 = 2
torch.cuda.manual_seed(1337)
torch.cuda.set_device(2)
if not os.path.exists(args.save):
os.makedirs(args.save)
for i in range(5):
if not os.path.exists(args.save + '/' + str(i)):
os.makedirs(args.save + '/' + str(i))
if args.model == 'DeeplabMulti':
model = DeeplabMulti(num_classes=args.num_classes)
elif args.model == 'DeeplabVGG':
model = DeeplabVGG(num_classes=args.num_classes)
if args.restore_from == RESTORE_FROM:
args.restore_from = RESTORE_FROM_VGG
print("begin")
if args.restore_from[:4] == 'http':
print("1112222")
#saved_state_dict = model_zoo.load_url(args.restore_from)
else:
print("2222222", gpu0)
# saved_state_dict = torch.load(args.restore_from)
print(args.restore_from)
model.load_state_dict(torch.load(args.restore_from))
model.cuda(gpu0)
# print(sys.getsizeof(model))
# model.eval()
# exit()
testloader = data.DataLoader(cityscapesDataSet(args.data_dir,
args.data_list,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set=args.set),
batch_size=1,
shuffle=False,
pin_memory=True)
if version.parse(torch.__version__) >= version.parse('0.4.0'):
# interp = nn.Upsample(size=(1024, 2048), mode='bilinear', align_corners=True)
interp = Upsample_function
else:
interp = nn.Upsample(size=(1024, 2048), mode='bilinear')
with torch.no_grad():
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, labels, _, name = batch
image = Variable(image).cuda(gpu0)
final = []
if args.model == 'DeeplabMulti':
output1, output2 = model(image)
output1 = F.softmax(output1, 1)
output2 = F.softmax(output2, 1)
for i in [0, 3, 7, 10]:
final_output = i / 10.0 * output1 + (10.0 -
i) / 10.0 * output2
output = interp(final_output).cpu().data[0].numpy()
final.append(output)
break
labels = labels.cpu().data[0].numpy()
elif args.model == 'DeeplabVGG':
output = model(Variable(image, volatile=True).cuda(gpu0))
output = interp(output).cpu().data[0].numpy()
name = name[0].split('/')[-1]
# labels_col = colorize_mask(labels)
# labels_col.save('%s/%s_real.png' % (args.save, name.split('.')[0]))
for i in range(4):
output = final[i]
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_col = colorize_mask(output)
output = Image.fromarray(output)
output.save('%s/%s/%s' % (args.save, str(i), name))
output_col.save('%s/%s/%s_color.png' %
(args.save, str(i), name.split('.')[0]))
break
def main():
"""Create the model and start the training."""
device = torch.device("cuda" if not args.cpu else "cpu")
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
# Create network
if args.model == 'ResNet':
model = DeeplabMulti(num_classes=args.num_classes)
if args.restore_from[:4] == 'http':
saved_state_dict = model_zoo.load_url(args.restore_from)
else:
saved_state_dict = torch.load(args.restore_from)
new_params = model.state_dict().copy()
for i in saved_state_dict:
# Scale.layer5.conv2d_list.3.weight
i_parts = i.split('.')
# print i_parts
if not args.num_classes == 19 or not i_parts[1] == 'layer5':
new_params['.'.join(i_parts[1:])] = saved_state_dict[i]
# print i_parts
model.load_state_dict(new_params)
if args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes,
vgg16_caffe_path='./model/vgg16_init.pth',
pretrained=True)
model.train()
model.to(device)
cudnn.benchmark = True
# init D
if args.model == 'ResNet':
model_D = FCDiscriminator(num_classes=2048).to(device)
if args.model == 'VGG':
model_D = FCDiscriminator(num_classes=1024).to(device)
model_D.train()
model_D.to(device)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps *
args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
# implement model.optim_parameters(args) to handle different models' lr setting
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
# set up tensor board
if args.tensorboard:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
for i_iter in range(args.num_steps):
loss_seg = 0
loss_adv_target_value = 0
loss_D_value = 0
loss_cla_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.to(device)
labels = labels.long().to(device)
feature, prediction = model(images)
prediction = interp(prediction)
loss = seg_loss(prediction, labels)
loss.backward()
loss_seg = loss.item()
# train with target
_, batch = targetloader_iter.__next__()
images, _, _ = batch
images = images.to(device)
feature_target, _ = model(images)
_, D_out = model_D(feature_target)
loss_adv_target = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
#print(args.lambda_adv_target)
loss = args.lambda_adv_target * loss_adv_target
loss.backward()
loss_adv_target_value = loss_adv_target.item()
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with source
feature = feature.detach()
cla, D_out = model_D(feature)
cla = interp(cla)
loss_cla = seg_loss(cla, labels)
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
loss_D = loss_D / 2
#print(args.lambda_s)
loss_Disc = args.lambda_s * loss_cla + loss_D
loss_Disc.backward()
loss_cla_value = loss_cla.item()
loss_D_value = loss_D.item()
# train with target
feature_target = feature_target.detach()
_, D_out = model_D(feature_target)
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(target_label).to(device))
loss_D = loss_D / 2
loss_D.backward()
loss_D_value += loss_D.item()
optimizer.step()
optimizer_D.step()
if args.tensorboard:
scalar_info = {
'loss_seg': loss_seg,
'loss_cla': loss_cla_value,
'loss_adv_target': loss_adv_target_value,
'loss_D': loss_D_value,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
#print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f} loss_adv = {3:.3f} loss_D = {4:.3f} loss_cla = {5:.3f}'
.format(i_iter, args.num_steps, loss_seg, loss_adv_target_value,
loss_D_value, loss_cla_value))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D.pth'))
if args.tensorboard:
writer.close()
def main():
"""Create the model and start the training."""
device = torch.device("cuda" if not args.cpu else "cpu")
w, h = map(int, args.input_size.split(','))
input_size = (w, h)
w, h = map(int, args.input_size_target.split(','))
input_size_target = (w, h)
cudnn.enabled = True
bestIoU = 0
bestIter = 0
# Create network
if args.model == 'ResNet':
model = DeeplabMulti(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
if args.model == 'VGG':
model = DeeplabVGG(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(saved_state_dict)
model.train()
model.to(device)
cudnn.benchmark = True
# init D
if args.model == 'ResNet':
model_D = FCDiscriminator(num_classes=256).to(device)
saved_state_dict = torch.load('./snapshots/BestGTA5_D.pth')
model_D.load_state_dict(saved_state_dict)
if args.model == 'VGG':
model_D = FCDiscriminator(num_classes=256).to(device)
model_D.train()
model_D.to(device)
if not os.path.exists(args.snapshot_dir):
os.makedirs(args.snapshot_dir)
trainloader = data.DataLoader(GTA5DataSet(args.data_dir,
args.data_list,
max_iters=args.num_steps *
args.iter_size * args.batch_size,
crop_size=input_size,
scale=args.random_scale,
mirror=args.random_mirror,
mean=IMG_MEAN),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
trainloader_iter = enumerate(trainloader)
targetloader = data.DataLoader(cityscapesDataSetLabel(
args.data_dir_target,
args.data_list_target,
max_iters=args.num_steps * args.iter_size * args.batch_size,
crop_size=input_size_target,
scale=False,
mirror=args.random_mirror,
mean=IMG_MEAN,
set=args.set),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True)
targetloader_iter = enumerate(targetloader)
optimizer = optim.SGD(model.optim_parameters(args),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.zero_grad()
optimizer_D = optim.Adam(model_D.parameters(),
lr=args.learning_rate_D,
betas=(0.9, 0.99))
optimizer_D.zero_grad()
bce_loss = torch.nn.BCEWithLogitsLoss()
seg_loss = torch.nn.CrossEntropyLoss(ignore_index=255)
interp = nn.Upsample(size=(input_size[1], input_size[0]),
mode='bilinear',
align_corners=True)
interp_target = nn.Upsample(size=(input_size_target[1],
input_size_target[0]),
mode='bilinear',
align_corners=True)
test_interp = nn.Upsample(size=(1024, 2048),
mode='bilinear',
align_corners=True)
# labels for adversarial training
source_label = 0
target_label = 1
# load calculated class center for initilization
class_center_source_ori = np.load('./source_center.npy')
class_center_source_ori = torch.from_numpy(class_center_source_ori)
class_center_target_ori = np.load('./target_center.npy')
class_center_target_ori = torch.from_numpy(class_center_target_ori)
# set up tensor board
if args.tensorboard:
if not os.path.exists(args.log_dir):
os.makedirs(args.log_dir)
writer = SummaryWriter(args.log_dir)
for i_iter in range(args.num_steps):
loss_seg = 0
loss_adv_target_value = 0
loss_D_value = 0
loss_cla_value = 0
loss_square_value = 0
loss_st_value = 0
optimizer.zero_grad()
adjust_learning_rate(optimizer, i_iter)
optimizer_D.zero_grad()
adjust_learning_rate_D(optimizer_D, i_iter)
# train G
# don't accumulate grads in D
for param in model_D.parameters():
param.requires_grad = False
# train with source
_, batch = trainloader_iter.__next__()
images, labels, _, _ = batch
images = images.to(device)
labels_s = labels # copy for center calculation
labels = labels.long().to(device)
feature, prediction = model(images)
feature_s = feature # copy for center calculation
prediction = interp(prediction)
loss = seg_loss(prediction, labels)
loss.backward(retain_graph=True)
loss_seg = loss.item()
# train with target
_, batch = targetloader_iter.__next__()
images, labels_pseudo, _, _ = batch
labels_t = labels_pseudo # copy for center calculation
images = images.to(device)
labels_pseudo = labels_pseudo.long().to(device)
feature_target, pred_target = model(images)
feature_t = feature_target # copy for center calculation
_, D_out = model_D(feature_target)
loss_adv_target = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
#print(args.lambda_adv_target)
loss = args.lambda_adv_target * loss_adv_target
loss.backward(retain_graph=True)
loss_adv_target_value = loss_adv_target.item()
pred_target = interp_target(pred_target)
loss_st = seg_loss(pred_target, labels_pseudo)
loss_st.backward(retain_graph=True)
loss_st_value = loss_st.item()
# class center alignment begin
if i_iter > 10000:
class_center_source = class_center_cal(feature_s, labels_s)
class_center_target = class_center_cal(feature_t, labels_t)
class_center_source_ori = class_center_update(
class_center_source, class_center_source_ori,
args.lambda_center_update)
class_center_target_ori = class_center_update(
class_center_target, class_center_target_ori,
args.lambda_center_update)
class_center_source_ori = class_center_source_ori.detach(
) #align target center to source
center_diff = class_center_source_ori - class_center_target_ori
loss_square = torch.pow(center_diff, 2).sum()
loss = args.lambda_center * loss_square
loss.backward()
loss_square_value = loss_square.item()
# class center alignment end
# train D
# bring back requires_grad
for param in model_D.parameters():
param.requires_grad = True
# train with source
feature = feature.detach()
cla, D_out = model_D(feature)
cla = interp(cla)
loss_cla = seg_loss(cla, labels)
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(source_label).to(device))
loss_D = loss_D / 2
#print(args.lambda_s)
loss_Disc = args.lambda_s * loss_cla + loss_D
loss_Disc.backward()
loss_cla_value = loss_cla.item()
loss_D_value = loss_D.item()
# train with target
feature_target = feature_target.detach()
_, D_out = model_D(feature_target)
loss_D = bce_loss(
D_out,
torch.FloatTensor(
D_out.data.size()).fill_(target_label).to(device))
loss_D = loss_D / 2
loss_D.backward()
loss_D_value += loss_D.item()
optimizer.step()
optimizer_D.step()
class_center_target_ori = class_center_target_ori.detach()
if args.tensorboard:
scalar_info = {
'loss_seg': loss_seg,
'loss_cla': loss_cla_value,
'loss_adv_target': loss_adv_target_value,
'loss_st_value': loss_st_value,
'loss_D': loss_D_value,
}
if i_iter % 10 == 0:
for key, val in scalar_info.items():
writer.add_scalar(key, val, i_iter)
#print('exp = {}'.format(args.snapshot_dir))
print(
'iter = {0:8d}/{1:8d}, loss_seg = {2:.3f} loss_adv = {3:.3f} loss_D = {4:.3f} loss_cla = {5:.3f} loss_st = {6:.5f} loss_square = {7:.5f}'
.format(i_iter, args.num_steps, loss_seg, loss_adv_target_value,
loss_D_value, loss_cla_value, loss_st_value,
loss_square_value))
if i_iter >= args.num_steps_stop - 1:
print('save model ...')
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir,
'GTA5_' + str(args.num_steps_stop) + '_D.pth'))
break
if i_iter % args.save_pred_every == 0 and i_iter != 0:
print('taking snapshot ...')
if not os.path.exists(args.save):
os.makedirs(args.save)
testloader = data.DataLoader(cityscapesDataSet(
args.data_dir_target,
args.data_list_target_test,
crop_size=(1024, 512),
mean=IMG_MEAN,
scale=False,
mirror=False,
set='val'),
batch_size=1,
shuffle=False,
pin_memory=True)
model.eval()
for index, batch in enumerate(testloader):
if index % 100 == 0:
print('%d processd' % index)
image, _, name = batch
with torch.no_grad():
output1, output2 = model(Variable(image).to(device))
output = test_interp(output2).cpu().data[0].numpy()
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output = Image.fromarray(output)
name = name[0].split('/')[-1]
output.save('%s/%s' % (args.save, name))
mIoUs = compute_mIoU(osp.join(args.data_dir_target, 'gtFine/val'),
args.save, 'dataset/cityscapes_list')
mIoU = round(np.nanmean(mIoUs) * 100, 2)
print('===> current mIoU: ' + str(mIoU))
print('===> last best mIoU: ' + str(bestIoU))
print('===> last best iter: ' + str(bestIter))
if mIoU > bestIoU:
bestIoU = mIoU
bestIter = i_iter
torch.save(model.state_dict(),
osp.join(args.snapshot_dir, 'BestGTA5.pth'))
torch.save(model_D.state_dict(),
osp.join(args.snapshot_dir, 'BestGTA5_D.pth'))
torch.save(
model.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '.pth'))
torch.save(
model_D.state_dict(),
osp.join(args.snapshot_dir, 'GTA5_' + str(i_iter) + '_D.pth'))
model.train()
if args.tensorboard:
writer.close()