# Copyright (c) OpenMMLab. All rights reserved. import os import os.path as osp import tempfile from unittest import TestCase import cv2 import mmcv import numpy as np import torch from mmengine.structures import PixelData from mmseg.structures import SegDataSample from mmseg.visualization import SegLocalVisualizer class TestSegLocalVisualizer(TestCase): def test_add_datasample(self): h = 10 w = 12 num_class = 2 out_file = 'out_file' image = np.random.randint(0, 256, size=(h, w, 3)).astype('uint8') # test gt_sem_seg gt_sem_seg_data = dict(data=torch.randint(0, num_class, (1, h, w))) gt_sem_seg = PixelData(**gt_sem_seg_data) def test_add_datasample_forward(gt_sem_seg): data_sample = SegDataSample() data_sample.gt_sem_seg = gt_sem_seg with tempfile.TemporaryDirectory() as tmp_dir: seg_local_visualizer = SegLocalVisualizer( vis_backends=[dict(type='LocalVisBackend')], save_dir=tmp_dir) seg_local_visualizer.dataset_meta = dict( classes=('background', 'foreground'), palette=[[120, 120, 120], [6, 230, 230]]) # test out_file seg_local_visualizer.add_datasample(out_file, image, data_sample) assert os.path.exists( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png')) drawn_img = cv2.imread( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png')) assert drawn_img.shape == (h, w, 3) # test gt_instances and pred_instances pred_sem_seg_data = dict( data=torch.randint(0, num_class, (1, h, w))) pred_sem_seg = PixelData(**pred_sem_seg_data) data_sample.pred_sem_seg = pred_sem_seg seg_local_visualizer.add_datasample(out_file, image, data_sample) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h, w * 2, 3)) seg_local_visualizer.add_datasample( out_file, image, data_sample, draw_gt=False) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h, w, 3)) if torch.cuda.is_available(): test_add_datasample_forward(gt_sem_seg.cuda()) test_add_datasample_forward(gt_sem_seg) def test_cityscapes_add_datasample(self): h = 128 w = 256 num_class = 19 out_file = 'out_file_cityscapes' image = mmcv.imread( osp.join( osp.dirname(__file__), '../data/pseudo_cityscapes_dataset/leftImg8bit/val/frankfurt/frankfurt_000000_000294_leftImg8bit.png' # noqa ), 'color') sem_seg = mmcv.imread( osp.join( osp.dirname(__file__), '../data/pseudo_cityscapes_dataset/gtFine/val/frankfurt/frankfurt_000000_000294_gtFine_labelTrainIds.png' # noqa ), 'unchanged') sem_seg = torch.unsqueeze(torch.from_numpy(sem_seg), 0) gt_sem_seg_data = dict(data=sem_seg) gt_sem_seg = PixelData(**gt_sem_seg_data) def test_cityscapes_add_datasample_forward(gt_sem_seg): data_sample = SegDataSample() data_sample.gt_sem_seg = gt_sem_seg with tempfile.TemporaryDirectory() as tmp_dir: seg_local_visualizer = SegLocalVisualizer( vis_backends=[dict(type='LocalVisBackend')], save_dir=tmp_dir) seg_local_visualizer.dataset_meta = dict( classes=('road', 'sidewalk', 'building', 'wall', 'fence', 'pole', 'traffic light', 'traffic sign', 'vegetation', 'terrain', 'sky', 'person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle'), palette=[[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156], [190, 153, 153], [153, 153, 153], [250, 170, 30], [220, 220, 0], [107, 142, 35], [152, 251, 152], [70, 130, 180], [220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70], [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32]]) # test out_file seg_local_visualizer.add_datasample( out_file, image, data_sample, out_file=osp.join(tmp_dir, 'test.png')) self._assert_image_and_shape( osp.join(tmp_dir, 'test.png'), (h, w, 3)) # test gt_instances and pred_instances pred_sem_seg_data = dict( data=torch.randint(0, num_class, (1, h, w))) pred_sem_seg = PixelData(**pred_sem_seg_data) data_sample.pred_sem_seg = pred_sem_seg # test draw prediction with gt seg_local_visualizer.add_datasample(out_file, image, data_sample) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h, w * 2, 3)) # test draw prediction without gt seg_local_visualizer.add_datasample( out_file, image, data_sample, draw_gt=False) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h, w, 3)) if torch.cuda.is_available(): test_cityscapes_add_datasample_forward(gt_sem_seg.cuda()) test_cityscapes_add_datasample_forward(gt_sem_seg) def _assert_image_and_shape(self, out_file, out_shape): assert os.path.exists(out_file) drawn_img = cv2.imread(out_file) assert drawn_img.shape == out_shape def test_add_datasample_depth(self): h = 10 w = 12 out_file = 'out_file' image = np.random.randint(0, 256, size=(h, w, 3)).astype('uint8') # test gt_depth_map gt_depth_map = PixelData(data=torch.rand(1, h, w)) def test_add_datasample_forward_depth(gt_depth_map): data_sample = SegDataSample() data_sample.gt_depth_map = gt_depth_map with tempfile.TemporaryDirectory() as tmp_dir: seg_local_visualizer = SegLocalVisualizer( vis_backends=[dict(type='LocalVisBackend')], save_dir=tmp_dir) seg_local_visualizer.dataset_meta = dict( classes=('background', 'foreground'), palette=[[120, 120, 120], [6, 230, 230]]) # test out_file seg_local_visualizer.add_datasample(out_file, image, data_sample) assert os.path.exists( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png')) drawn_img = cv2.imread( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png')) assert drawn_img.shape == (h * 2, w, 3) # test gt_instances and pred_instances pred_depth_map = PixelData(data=torch.rand(1, h, w)) data_sample.pred_depth_map = pred_depth_map seg_local_visualizer.add_datasample(out_file, image, data_sample) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h * 2, w * 2, 3)) seg_local_visualizer.add_datasample( out_file, image, data_sample, draw_gt=False) self._assert_image_and_shape( osp.join(tmp_dir, 'vis_data', 'vis_image', out_file + '_0.png'), (h * 2, w, 3)) if torch.cuda.is_available(): test_add_datasample_forward_depth(gt_depth_map.cuda()) test_add_datasample_forward_depth(gt_depth_map)