Skip to content
Snippets Groups Projects
Commit 9cbb4b12 authored by David de la Iglesia Castro's avatar David de la Iglesia Castro Committed by GitHub
Browse files

Add opacity option to show_result (#425)

parent 17220103
No related branches found
No related tags found
No related merge requests found
......@@ -15,6 +15,11 @@ def main():
'--palette',
default='cityscapes',
help='Color palette used for segmentation map')
parser.add_argument(
'--opacity',
type=float,
default=0.5,
help='Opacity of painted segmentation map. In (0, 1] range.')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
......@@ -22,7 +27,12 @@ def main():
# test a single image
result = inference_segmentor(model, args.img)
# show the results
show_result_pyplot(model, args.img, result, get_palette(args.palette))
show_result_pyplot(
model,
args.img,
result,
get_palette(args.palette),
opacity=args.opacity)
if __name__ == '__main__':
......
......@@ -166,7 +166,8 @@ result = inference_segmentor(model, img)
# visualize the results in a new window
model.show_result(img, result, show=True)
# or save the visualization results to image files
model.show_result(img, result, out_file='result.jpg')
# you can change the opacity of the painted segmentation map in (0, 1].
model.show_result(img, result, out_file='result.jpg', opacity=0.5)
# test a video and show the results
video = mmcv.VideoReader('video.mp4')
......
......@@ -98,7 +98,12 @@ def inference_segmentor(model, img):
return result
def show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10)):
def show_result_pyplot(model,
img,
result,
palette=None,
fig_size=(15, 10),
opacity=0.5):
"""Visualize the segmentation results on the image.
Args:
......@@ -109,10 +114,14 @@ def show_result_pyplot(model, img, result, palette=None, fig_size=(15, 10)):
map. If None is given, random palette will be generated.
Default: None
fig_size (tuple): Figure size of the pyplot figure.
opacity(float): Opacity of painted segmentation map.
Default 0.5.
Must be in (0, 1] range.
"""
if hasattr(model, 'module'):
model = model.module
img = model.show_result(img, result, palette=palette, show=False)
img = model.show_result(
img, result, palette=palette, show=False, opacity=opacity)
plt.figure(figsize=fig_size)
plt.imshow(mmcv.bgr2rgb(img))
plt.show()
......@@ -35,7 +35,8 @@ def single_gpu_test(model,
data_loader,
show=False,
out_dir=None,
efficient_test=False):
efficient_test=False,
opacity=0.5):
"""Test with single GPU.
Args:
......@@ -46,7 +47,9 @@ def single_gpu_test(model,
the directory to save output results.
efficient_test (bool): Whether save the results as local numpy files to
save CPU memory during evaluation. Default: False.
opacity(float): Opacity of painted segmentation map.
Default 0.5.
Must be in (0, 1] range.
Returns:
list: The prediction results.
"""
......@@ -82,7 +85,8 @@ def single_gpu_test(model,
result,
palette=dataset.PALETTE,
show=show,
out_file=out_file)
out_file=out_file,
opacity=opacity)
if isinstance(result, list):
if efficient_test:
......
......@@ -212,7 +212,8 @@ class BaseSegmentor(nn.Module):
win_name='',
show=False,
wait_time=0,
out_file=None):
out_file=None,
opacity=0.5):
"""Draw `result` over `img`.
Args:
......@@ -229,7 +230,9 @@ class BaseSegmentor(nn.Module):
Default: False.
out_file (str or None): The filename to write the image.
Default: None.
opacity(float): Opacity of painted segmentation map.
Default 0.5.
Must be in (0, 1] range.
Returns:
img (Tensor): Only if not `show` or `out_file`
"""
......@@ -246,13 +249,14 @@ class BaseSegmentor(nn.Module):
assert palette.shape[0] == len(self.CLASSES)
assert palette.shape[1] == 3
assert len(palette.shape) == 2
assert 0 < opacity <= 1.0
color_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
for label, color in enumerate(palette):
color_seg[seg == label, :] = color
# convert to BGR
color_seg = color_seg[..., ::-1]
img = img * 0.5 + color_seg * 0.5
img = img * (1 - opacity) + color_seg * opacity
img = img.astype(np.uint8)
# if out_file specified, do not show image in window
if out_file is not None:
......
......@@ -55,6 +55,11 @@ def parse_args():
choices=['none', 'pytorch', 'slurm', 'mpi'],
default='none',
help='job launcher')
parser.add_argument(
'--opacity',
type=float,
default=0.5,
help='Opacity of painted segmentation map. In (0, 1] range.')
parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args()
if 'LOCAL_RANK' not in os.environ:
......@@ -123,7 +128,7 @@ def main():
if not distributed:
model = MMDataParallel(model, device_ids=[0])
outputs = single_gpu_test(model, data_loader, args.show, args.show_dir,
efficient_test)
efficient_test, args.opacity)
else:
model = MMDistributedDataParallel(
model.cuda(),
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment