def save_predict(output,
gt,
img_name,
dataset,
save_path,
output_grey=False,
output_color=True,
gt_color=False):
if output_grey:
output_grey = Image.fromarray(output)
output_grey.save(os.path.join(save_path, img_name + '.png'))
if output_color:
if dataset == 'cityscapes':
output_color = cityscapes_colorize_mask(output)
elif dataset == 'camvid':
output_color = camvid_colorize_mask(output)
output_color.save(os.path.join(save_path, img_name + '_color.png'))
if gt_color:
if dataset == 'cityscapes':
gt_color = cityscapes_colorize_mask(gt)
elif dataset == 'camvid':
gt_color = camvid_colorize_mask(gt)
gt_color.save(os.path.join(save_path, img_name + '_gt.png'))
def save_predict(output,
gt,
img_name,
dataset,
save_path,
output_grey=False,
output_color=True,
gt_color=False):
if output_grey:
row, col = output.shape
dst = np.ones((row, col), dtype=np.uint8) * 255
for i in range(19):
dst[output == i] = color_list[i]
output_grey = Image.fromarray(dst)
output_grey.save(os.path.join(save_path, img_name[:-1] + '.png'))
if output_color:
if dataset == 'cityscapes':
output_color = cityscapes_colorize_mask(output)
elif dataset == 'camvid':
output_color = camvid_colorize_mask(output)
output_color.save(os.path.join(save_path,
img_name[:-1] + '_color.png'))
if gt_color:
if dataset == 'cityscapes':
gt_color = cityscapes_colorize_mask(gt)
elif dataset == 'camvid':
gt_color = camvid_colorize_mask(gt)
gt_color.save(os.path.join(save_path, img_name + '_gt.png'))
def save_predict(output,
gt,
img_name,
dataset,
save_path,
output_grey=False,
output_color=True,
gt_color=False):
if output_grey:
if dataset == 'cityscapes':
output[np.where(output == 18)] = 33
output[np.where(output == 17)] = 32
output[np.where(output == 16)] = 31
output[np.where(output == 15)] = 28
output[np.where(output == 14)] = 27
output[np.where(output == 13)] = 26
output[np.where(output == 12)] = 25
output[np.where(output == 11)] = 24
output[np.where(output == 10)] = 23
output[np.where(output == 9)] = 22
output[np.where(output == 8)] = 21
output[np.where(output == 7)] = 20
output[np.where(output == 6)] = 19
output[np.where(output == 5)] = 17
output[np.where(output == 4)] = 13
output[np.where(output == 3)] = 12
output[np.where(output == 2)] = 11
output[np.where(output == 1)] = 8
output[np.where(output == 0)] = 7
output_grey = Image.fromarray(output)
output_grey.save(os.path.join(save_path, img_name + '.png'))
if output_color:
if dataset == 'cityscapes':
output_color = cityscapes_colorize_mask(output)
elif dataset == 'paris':
output_color = paris_colorize_mask(output)
elif dataset == 'road':
output_color = road_colorize_mask(output)
elif dataset == 'austin':
output_color = austin_colorize_mask(output)
elif dataset == 'postdam' or dataset == 'vaihingen':
output_color = isprs_colorize_mask(output)
output_color.save(os.path.join(save_path, img_name + '_color.png'))
if gt_color:
if dataset == 'cityscapes':
gt_color = cityscapes_colorize_mask(gt)
elif dataset == 'paris':
gt_color = paris_colorize_mask(gt)
elif dataset == 'road':
gt_color = road_colorize_mask(gt)
elif dataset == 'austin':
gt_color = austin_colorize_mask(gt)
elif dataset == 'postdam' or dataset == 'vaihingen':
gt_color = isprs_colorize_mask(gt)
gt_color.save(os.path.join(save_path, img_name + '_gt.png'))
def val(args, val_loader, model, criterion):
"""
args:
val_loader: loaded for validation dataset
model: model
criterion: loss function
return: IoU class, and mean IoU
"""
#evaluation mode
model.eval()
total_batches = len(val_loader)
data_list = []
for i, (input, label, size, name) in enumerate(val_loader):
start_time = time.time()
input_var = Variable(input, volatile=True).cuda()
output = model(input_var)
time_taken = time.time() - start_time
print('[%d/%d] time: %.2f' % (i, total_batches, time_taken))
# save seg image
output = output.cpu().data[0].numpy() # 1xCxHxW ---> CxHxW
gt = np.asarray(label[0].numpy(), dtype=np.uint8)
output = output.transpose(1, 2, 0) # CxHxW --> HxWxC
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
data_list.append([gt.flatten(), output.flatten()])
output_color = cityscapes_colorize_mask(output)
output = Image.fromarray(output)
output.save('%s/%s.png' % (args.save_seg_dir, name[0]))
output_color.save('%s/%s_color.png' % (args.save_seg_dir, name[0]))
meanIoU, IoUs = get_iou(data_list, args.classes)
print("mean IoU:", meanIoU)
print(IoUs)
def test(args, test_loader, model):
"""
args:
test_loader: loaded for test set
model: model
criterion: loss function
return: IoU class, and mean IoU
"""
#evaluation mode
model.eval()
total_batches = len(test_loader)
for i, (input, size, name) in enumerate(test_loader):
start_time = time.time()
input_var = Variable(input, volatile=True).cuda()
output = model(input_var)
time_taken = time.time() - start_time
print('[%d/%d] time: %.2f' % (i, total_batches, time_taken))
# save seg image
output = output.cpu().data[0].numpy() # 1xCxHxW ---> CxHxW
output = output.transpose(1, 2, 0) # CxHxW --> HxWxC
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_color = cityscapes_colorize_mask(output)
output = Image.fromarray(output)
#output.save( "%s/%s.png " % (args.save_seg_dir, name[0]) )
output_color.save("%s/%s_color.png" % (args.save_seg_dir, name[0]))
def test(args, test_loader, model):
'''
Args:
test_loader: loaded for test dataset
model: model
'''
#switch to evaluation mode
model.eval()
total_batches = len(test_loader)
data_list = []
for i, (input, size, name) in enumerate(test_loader):
start_time = time.time()
#input_var = Variable(input, volatile=True).cuda() torch version==0.3
with torch.no_grad():
input_var = Variable(input).cuda()
# run the mdoel
output = model(input_var)
time_taken = time.time() - start_time
print('[%d/%d] time: %.2f' % (i, total_batches, time_taken))
# save seg image
output = output.cpu().data[0].numpy() # 1xCxHxW ---> CxHxW
output = output.transpose(1, 2, 0) # CxHxW --> HxWxC
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
output_color = cityscapes_colorize_mask(output)
output = Image.fromarray(output)
output_color.save('%s/%s.png' % (args.save_seg_dir, name[0]))
def main():
args = get_arguments()
print("=====> Configure dataset and model")
configure_dataset_model(args)
print(args)
print("=====> Set GPU for training")
if args.cuda:
print("====> Use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
model = CoattentionNet(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from,
map_location=lambda storage, loc: storage)
#print(saved_state_dict.keys())
#model.load_state_dict({k.replace('pspmodule.',''):v for k,v in torch.load(args.restore_from)['state_dict'].items()})
model.load_state_dict(convert_state_dict(saved_state_dict["model"])
) #convert_state_dict(saved_state_dict["model"])
model.eval()
model.cuda()
if args.dataset == 'voc12':
testloader = data.DataLoader(VOCDataTestSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=args.img_mean),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
voc_colorize = VOCColorize()
elif args.dataset == 'cityscapes':
testloader = data.DataLoader(
CityscapesTestDataSet(args.data_dir,
args.data_list,
f_scale=args.f_scale,
mean=args.img_mean),
batch_size=1,
shuffle=False,
pin_memory=True
) # f_sale, meaning resize image at f_scale as input
interp = nn.Upsample(size=(1024, 2048), mode='bilinear') #size = (h,w)
voc_colorize = VOCColorize()
elif args.dataset == 'davis': #for davis 2016
db_test = db.PairwiseImg(
train=False,
inputRes=(473, 473),
db_root_dir=args.data_dir,
transform=None,
seq_name=None,
sample_range=args.sample_range
) #db_root_dir() --> '/path/to/DAVIS-2016' train path
testloader = data.DataLoader(db_test,
batch_size=1,
shuffle=False,
num_workers=0)
voc_colorize = VOCColorize()
else:
print("dataset error")
data_list = []
if args.save_segimage:
if not os.path.exists(args.seg_save_dir) and not os.path.exists(
args.vis_save_dir):
os.makedirs(args.seg_save_dir)
os.makedirs(args.vis_save_dir)
print("======> test set size:", len(testloader))
my_index = 0
old_temp = ''
for index, batch in enumerate(testloader):
print('%d processd' % (index))
target = batch['target']
#search = batch['search']
temp = batch['seq_name']
args.seq_name = temp[0]
print(args.seq_name)
if old_temp == args.seq_name:
my_index = my_index + 1
else:
my_index = 0
output_sum = 0
for i in range(0, args.sample_range):
search = batch['search' + '_' + str(i)]
search_im = search
#print(search_im.size())
output = model(
Variable(target, volatile=True).cuda(),
Variable(search_im, volatile=True).cuda())
#print(output[0]) # output有两个
output_sum = output_sum + output[0].data[
0, 0].cpu().numpy() #分割那个分支的结果
#np.save('infer'+str(i)+'.npy',output1)
#output2 = output[1].data[0, 0].cpu().numpy() #interp'
output1 = output_sum / args.sample_range
first_image = np.array(
Image.open(args.data_dir + '/JPEGImages/480p/blackswan/00000.jpg'))
original_shape = first_image.shape
output1 = cv2.resize(output1, (original_shape[1], original_shape[0]))
if 0:
original_image = target[0]
#print('image type:',type(original_image.numpy()))
original_image = original_image.numpy()
original_image = original_image.transpose((2, 1, 0))
original_image = cv2.resize(original_image,
(original_shape[1], original_shape[0]))
unary = np.zeros((2, original_shape[0] * original_shape[1]),
dtype='float32')
#unary[0, :, :] = res_saliency/255
#unary[1, :, :] = 1-res_saliency/255
EPSILON = 1e-8
tau = 1.05
crf = dcrf.DenseCRF(original_shape[1] * original_shape[0], 2)
anno_norm = (output1 - np.min(output1)) / (
np.max(output1) - np.min(output1)) #res_saliency/ 255.
n_energy = 1.0 - anno_norm + EPSILON #-np.log((1.0 - anno_norm + EPSILON)) #/ (tau * sigmoid(1 - anno_norm))
p_energy = anno_norm + EPSILON #-np.log(anno_norm + EPSILON) #/ (tau * sigmoid(anno_norm))
#unary = unary.reshape((2, -1))
#print(unary.shape)
unary[1, :] = p_energy.flatten()
unary[0, :] = n_energy.flatten()
crf.setUnaryEnergy(unary_from_softmax(unary))
feats = create_pairwise_gaussian(sdims=(3, 3),
shape=original_shape[:2])
crf.addPairwiseEnergy(feats,
compat=3,
kernel=dcrf.DIAG_KERNEL,
normalization=dcrf.NORMALIZE_SYMMETRIC)
feats = create_pairwise_bilateral(
sdims=(10, 10),
schan=(1, 1, 1), # orgin is 60, 60 5, 5, 5
img=original_image,
chdim=2)
crf.addPairwiseEnergy(feats,
compat=5,
kernel=dcrf.DIAG_KERNEL,
normalization=dcrf.NORMALIZE_SYMMETRIC)
Q = crf.inference(5)
MAP = np.argmax(Q, axis=0)
output1 = MAP.reshape((original_shape[0], original_shape[1]))
mask = (output1 * 255).astype(np.uint8)
#print(mask.shape[0])
mask = Image.fromarray(mask)
if args.dataset == 'voc12':
print(output.shape)
print(size)
output = output[:, :size[0], :size[1]]
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
if args.save_segimage:
seg_filename = os.path.join(args.seg_save_dir,
'{}.png'.format(name[0]))
color_file = Image.fromarray(
voc_colorize(output).transpose(1, 2, 0), 'RGB')
color_file.save(seg_filename)
elif args.dataset == 'davis':
save_dir_res = os.path.join(args.seg_save_dir, 'Results',
args.seq_name)
old_temp = args.seq_name
if not os.path.exists(save_dir_res):
os.makedirs(save_dir_res)
if args.save_segimage:
my_index1 = str(my_index).zfill(5)
seg_filename = os.path.join(save_dir_res,
'{}.png'.format(my_index1))
#color_file = Image.fromarray(voc_colorize(output).transpose(1, 2, 0), 'RGB')
mask.save(seg_filename)
#np.concatenate((torch.zeros(1, 473, 473), mask, torch.zeros(1, 512, 512)),axis = 0)
#save_image(output1 * 0.8 + target.data, args.vis_save_dir, normalize=True)
elif args.dataset == 'cityscapes':
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
if args.save_segimage:
output_color = cityscapes_colorize_mask(output)
output = Image.fromarray(output)
output.save('%s/%s.png' % (args.seg_save_dir, name[0]))
output_color.save('%s/%s_color.png' %
(args.seg_save_dir, name[0]))
else:
print("dataset error")
def main():
args = get_arguments()
print("=====> Configure dataset and model")
configure_dataset_model(args)
print(args)
print("=====> Set GPU for training")
if args.cuda:
print("====> Use gpu id: '{}'".format(args.gpus))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
if not torch.cuda.is_available():
raise Exception(
"No GPU found or Wrong gpu id, please run without --cuda")
model = PSPNet(num_classes=args.num_classes)
saved_state_dict = torch.load(args.restore_from)
model.load_state_dict(convert_state_dict(saved_state_dict["model"]))
model.eval()
model.cuda()
if args.dataset == 'voc12':
testloader = data.DataLoader(VOCDataTestSet(args.data_dir,
args.data_list,
crop_size=(505, 505),
mean=args.img_mean),
batch_size=1,
shuffle=False,
pin_memory=True)
interp = nn.Upsample(size=(505, 505), mode='bilinear')
voc_colorize = VOCColorize()
elif args.dataset == 'cityscapes':
testloader = data.DataLoader(
CityscapesTestDataSet(args.data_dir,
args.data_list,
f_scale=args.f_scale,
mean=args.img_mean),
batch_size=1,
shuffle=False,
pin_memory=True
) # f_sale, meaning resize image at f_scale as input
interp = nn.Upsample(size=(1024, 2048), mode='bilinear') #size = (h,w)
else:
print("dataset error")
data_list = []
if args.save_segimage:
if not os.path.exists(args.seg_save_dir):
os.makedirs(args.seg_save_dir)
print("======> test set size:", len(testloader))
for index, batch in enumerate(testloader):
print('%d processd' % (index))
image, size, name = batch
size = size[0].numpy()
output = model(Variable(image, volatile=True).cuda())
output = interp(output).cpu().data[0].numpy()
if args.dataset == 'voc12':
print(output.shape)
print(size)
output = output[:, :size[0], :size[1]]
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
if args.save_segimage:
seg_filename = os.path.join(args.seg_save_dir,
'{}.png'.format(name[0]))
color_file = Image.fromarray(
voc_colorize(output).transpose(1, 2, 0), 'RGB')
color_file.save(seg_filename)
elif args.dataset == 'cityscapes':
output = output.transpose(1, 2, 0)
output = np.asarray(np.argmax(output, axis=2), dtype=np.uint8)
if args.save_segimage:
output_color = cityscapes_colorize_mask(output)
output = Image.fromarray(output)
output.save('%s/%s.png' % (args.seg_save_dir, name[0]))
output_color.save('%s/%s_color.png' %
(args.seg_save_dir, name[0]))
else:
print("dataset error")
