Source code for mmseg.apis.test
import os.path as osp
import pickle
import shutil
import tempfile
import mmcv
import numpy as np
import torch
import torch.distributed as dist
from mmcv.image import tensor2imgs
from mmcv.runner import get_dist_info
import pdb
import cv2, os, sys, matplotlib.pyplot as plt
from PIL import Image
sys.path.append('/home/xshadow/LAM_Demo')
from ModelZoo.utils import load_as_tensor, Tensor2PIL, PIL2Tensor, _add_batch_one
from ModelZoo import get_model, load_model, print_network
from SaliencyModel.utils import vis_saliency, vis_saliency_kde, click_select_position, grad_abs_norm, grad_norm, prepare_images, make_pil_grid, blend_input
from SaliencyModel.utils import cv2_to_pil, pil_to_cv2, gini
from SaliencyModel.attributes import attr_grad, attr_id, attr_gabor_generator
from SaliencyModel.BackProp import I_gradient, attribution_objective, Path_gradient_GTA
from SaliencyModel.BackProp import saliency_map_PG as saliency_map
from SaliencyModel.BackProp import GaussianBlurPath, LinearPath, CosinePath
from SaliencyModel.utils import grad_norm, IG_baseline, interpolation, isotropic_gaussian_kernel
import pdb
def np2tmp(array, temp_file_name=None):
"""Save ndarray to local numpy file.
Args:
array (ndarray): Ndarray to save.
temp_file_name (str): Numpy file name. If 'temp_file_name=None', this
function will generate a file name with tempfile.NamedTemporaryFile
to save ndarray. Default: None.
Returns:
str: The numpy file name.
"""
if temp_file_name is None:
temp_file_name = tempfile.NamedTemporaryFile(
suffix='.npy', delete=False).name
np.save(temp_file_name, array)
return temp_file_name
[docs]def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
efficient_test=False):
"""Test with single GPU.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
window_size = 8 # Define windoes_size of D
#w = 403 # The x coordinate of your select patch, 125 as an example
#h = 161 # The y coordinate of your select patch, 160 as an example
# And check the red box
# Is your selected patch this one? If not, adjust the `w` and `h`.
#w, h = 1603,497
w, h = 297,630
w,h = 181,128
w,h = 1193,377
w,h = 463,611
w,h = 301,617
w,h = 645,699
w,h = 929,112
w,h = 161,531
w,h = 780,163
w,h = 247,122
w,h = 192,379
w,h = 743,554
w,h = 560,498
w,h = 597,664
w,h = 51,894
#w,h = 387,609
w,h = 237,967
w,h = 694,879
w,h = 743,554
w,h = 597,664
fold = 100; alpha = 0.4
attr_objective = attribution_objective(attr_id, h, w, window=window_size)
gaus_blur_path_func = LinearPath(fold)
#gaus_blur_path_func = CosinePath(fold)
for i, data in enumerate(data_loader):
#with torch.no_grad():
#result = model(return_loss=False, **data)
#result = model(return_loss=False, img=[data['img']], img_metas=[data['img_metas']])
interpolated_grad_numpy, result_numpy, interpolated_numpy = Path_gradient_GTA(data, model, attr_objective, gaus_blur_path_func, cuda=True)
img_lr = data['img'][0].cpu().numpy().squeeze()
img_lr = cv2.imread('/home/xshadow/GTA_height/images/validation_t/img--1036_-1540_299_15_278_0_0_k0.jpg')
img_lr = cv2.imread('/home/xshadow/height_data/images/validation_syt/OMA_329_001_RGB.tif')
#img_lr = cv2.resize(img_lr, (960,540))
#img_lr = np.moveaxis(img_lr,0,2)
position_pil = cv2_to_pil(img_lr)
grad_numpy, result = saliency_map(interpolated_grad_numpy, result_numpy)
abs_normed_grad_numpy = grad_abs_norm(grad_numpy)
saliency_image_abs = vis_saliency(abs_normed_grad_numpy, zoomin=2)
#saliency_image_kde = vis_saliency_kde(abs_normed_grad_numpy)
#pdb.set_trace()
blend_abs_and_input = cv2_to_pil(pil_to_cv2(saliency_image_abs) * (1.0 - alpha) + img_lr * alpha)
#blend_kde_and_input = cv2_to_pil(pil_to_cv2(saliency_image_kde) * (1.0 - alpha) + pil_to_cv2(img_lr.resize(img_hr.size)) * alpha)
#'''
pil = make_pil_grid(
[position_pil,
saliency_image_abs,
blend_abs_and_input]
)
#'''
pil.save('pppil.png')
saliency_image_abs.save('pig_test.png')
pdb.set_trace()
simage = result.detach().cpu().numpy().squeeze()
simage = (simage - simage.min()) / (simage.max() - simage.min()) * 255
simage = simage.astype(np.uint8)
simage = cv2.applyColorMap(simage, cv2.COLORMAP_JET)
cv2.imwrite('tt.png', simage)
if isinstance(result, list):
if efficient_test:
result = [np2tmp(_) for _ in result]
results.extend(result)
else:
if efficient_test:
result = np2tmp(result)
results.append(result)
#batch_size = data['img'][0].size(0)
batch_size=1
for _ in range(batch_size):
prog_bar.update()
return results
def single_gpu_test_ori(model,
data_loader,
show=False,
out_dir=None,
efficient_test=False):
"""Test with single GPU.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
#result = model(return_loss=False, **data)
result = model(return_loss=False, img=[data['img']], img_metas=[data['img_metas']])
#pdb.set_trace()
if isinstance(result, list):
if efficient_test:
result = [np2tmp(_) for _ in result]
results.extend(result)
else:
if efficient_test:
result = np2tmp(result)
results.append(result)
#batch_size = data['img'][0].size(0)
batch_size=1
for _ in range(batch_size):
prog_bar.update()
return results
[docs]def multi_gpu_test(model,
data_loader,
tmpdir=None,
gpu_collect=False,
efficient_test=False):
"""Test model with multiple gpus.
This method tests model with multiple gpus and collects the results
under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
it encodes results to gpu tensors and use gpu communication for results
collection. On cpu mode it saves the results on different gpus to 'tmpdir'
and collects them by the rank 0 worker.
Args:
model (nn.Module): Model to be tested.
data_loader (utils.data.Dataloader): Pytorch data loader.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode.
gpu_collect (bool): Option to use either gpu or cpu to collect results.
efficient_test (bool): Whether save the results as local numpy files to
save CPU memory during evaluation. Default: False.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
for i, data in enumerate(data_loader):
with torch.no_grad():
result = model(return_loss=False, rescale=True, **data)
if isinstance(result, list):
if efficient_test:
result = [np2tmp(_) for _ in result]
results.extend(result)
else:
if efficient_test:
result = np2tmp(result)
results.append(result)
if rank == 0:
batch_size = data['img'][0].size(0)
for _ in range(batch_size * world_size):
prog_bar.update()
# collect results from all ranks
if gpu_collect:
results = collect_results_gpu(results, len(dataset))
else:
results = collect_results_cpu(results, len(dataset), tmpdir)
return results
def collect_results_cpu(result_part, size, tmpdir=None):
"""Collect results with CPU."""
rank, world_size = get_dist_info()
# create a tmp dir if it is not specified
if tmpdir is None:
MAX_LEN = 512
# 32 is whitespace
dir_tensor = torch.full((MAX_LEN, ),
32,
dtype=torch.float32,
device='cuda')
if rank == 0:
tmpdir = tempfile.mkdtemp()
tmpdir = torch.tensor(
bytearray(tmpdir.encode()), dtype=torch.float32, device='cuda')
dir_tensor[:len(tmpdir)] = tmpdir
dist.broadcast(dir_tensor, 0)
tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
else:
mmcv.mkdir_or_exist(tmpdir)
# dump the part result to the dir
mmcv.dump(result_part, osp.join(tmpdir, 'part_{}.pkl'.format(rank)))
dist.barrier()
# collect all parts
if rank != 0:
return None
else:
# load results of all parts from tmp dir
part_list = []
for i in range(world_size):
part_file = osp.join(tmpdir, 'part_{}.pkl'.format(i))
part_list.append(mmcv.load(part_file))
# sort the results
ordered_results = []
for res in zip(*part_list):
ordered_results.extend(list(res))
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
# remove tmp dir
shutil.rmtree(tmpdir)
return ordered_results
def collect_results_gpu(result_part, size):
"""Collect results with GPU."""
rank, world_size = get_dist_info()
# dump result part to tensor with pickle
part_tensor = torch.tensor(
bytearray(pickle.dumps(result_part)), dtype=torch.float32, device='cuda')
# gather all result part tensor shape
shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
shape_list = [shape_tensor.clone() for _ in range(world_size)]
dist.all_gather(shape_list, shape_tensor)
# padding result part tensor to max length
shape_max = torch.tensor(shape_list).max()
part_send = torch.zeros(shape_max, dtype=torch.float32, device='cuda')
part_send[:shape_tensor[0]] = part_tensor
part_recv_list = [
part_tensor.new_zeros(shape_max) for _ in range(world_size)
]
# gather all result part
dist.all_gather(part_recv_list, part_send)
if rank == 0:
part_list = []
for recv, shape in zip(part_recv_list, shape_list):
part_list.append(
pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
# sort the results
ordered_results = []
for res in zip(*part_list):
ordered_results.extend(list(res))
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
return ordered_results