Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Internal PR, enable the segement_only flag #3

Merged
merged 3 commits into from
Dec 12, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion demo/image_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,10 @@ def main():
help='Opacity of painted segmentation map. In (0, 1] range.')
parser.add_argument(
'--title', default='result', help='The image identifier.')
parser.add_argument(
'--segementation-only',
action='store_true',
help='Only show the segmentation map.')
args = parser.parse_args()

# build the model from a config file and a checkpoint file
Expand All @@ -38,7 +42,8 @@ def main():
opacity=args.opacity,
draw_gt=False,
show=False if args.out_file is not None else True,
out_file=args.out_file)
out_file=args.out_file,
segmentation_only=args.segementation_only)


if __name__ == '__main__':
Expand Down
2 changes: 2 additions & 0 deletions docs/en/user_guides/visualization.md
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,8 @@ gt_sem_seg = PixelData(**gt_sem_seg_data)
data_sample = SegDataSample()
data_sample.gt_sem_seg = gt_sem_seg

# If you want to show the segement data only (without lable text and background)
# you can use segement_only = True.
seg_local_visualizer = SegLocalVisualizer(
vis_backends=[dict(type='LocalVisBackend')],
save_dir=save_dir)
Expand Down
11 changes: 10 additions & 1 deletion mmseg/apis/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -128,6 +128,7 @@ def show_result_pyplot(model: BaseSegmentor,
wait_time: float = 0,
show: bool = True,
withLabels: Optional[bool] = True,
segmentation_only: bool = False,
save_dir=None,
out_file=None):
"""Visualize the segmentation results on the image.
Expand Down Expand Up @@ -167,10 +168,18 @@ def show_result_pyplot(model: BaseSegmentor,
if save_dir is not None:
mkdir_or_exist(save_dir)
# init visualizer
if withLabels and segmentation_only:
# Issue a warning if withLabels is True and segmentation_only is True
warnings.simplefilter('once')
warnings.warn(
'withLabels is True and segmentation_only is True, '
'withLabels will be set to False')
withLabels = False
visualizer = SegLocalVisualizer(
vis_backends=[dict(type='LocalVisBackend')],
save_dir=save_dir,
alpha=opacity)
alpha=opacity,
segement_only=segmentation_only)
visualizer.dataset_meta = dict(
classes=model.dataset_meta['classes'],
palette=model.dataset_meta['palette'])
Expand Down
11 changes: 10 additions & 1 deletion mmseg/visualization/local_visualizer.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,11 @@ class SegLocalVisualizer(Visualizer):
Defaults to None.
alpha (int, float): The transparency of segmentation mask.
Defaults to 0.8.
segement_only (bool): Only show the segmentation map, by setting this flag
to True, only the segmentation map will be shown, the background image
and the labels will not be shown. Defaults to False.
Warning: If withLabels is True and segmentation_only is True,
withLabels will be set to False. i.e. this flag will override withLabels.

Examples:
>>> import numpy as np
Expand Down Expand Up @@ -73,10 +78,14 @@ def __init__(self,
palette: Optional[List] = None,
dataset_name: Optional[str] = None,
alpha: float = 0.8,
segement_only: bool = False,
**kwargs):
super().__init__(name, image, vis_backends, save_dir, **kwargs)
self.alpha: float = alpha
if segement_only:
self.alpha = 1
self.set_dataset_meta(palette, classes, dataset_name)
self.segement_only = segement_only

def _get_center_loc(self, mask: np.ndarray) -> np.ndarray:
"""Get semantic seg center coordinate.
Expand Down Expand Up @@ -139,7 +148,7 @@ def _draw_sem_seg(self,
for label, color in zip(labels, colors):
mask[sem_seg[0] == label, :] = color

if withLabels:
if withLabels and not self.segement_only:
font = cv2.FONT_HERSHEY_SIMPLEX
# (0,1] to change the size of the text relative to the image
scale = 0.05
Expand Down