1# Copyright 2021 Huawei Technologies Co., Ltd 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15import os 16import numpy as np 17import mindspore.dataset as ds 18from mindspore.dataset.utils.browse_dataset import imshow_det_bbox 19 20 21def test_browse_dataset(): 22 ''' 23 Demo code of visualization on VOC detection dataset. 24 ''' 25 # init 26 DATA_DIR = "../data/dataset/testVOC2012_2" 27 dataset = ds.VOCDataset(DATA_DIR, task="Detection", usage="train", shuffle=False, decode=True, num_samples=3) 28 dataset_iter = dataset.create_dict_iterator(output_numpy=True, num_epochs=1) 29 30 # iter 31 for index, data in enumerate(dataset_iter): 32 image = data["image"] 33 bbox = data["bbox"] 34 label = data["label"] 35 36 masks = np.zeros((4, image.shape[0], image.shape[1])) 37 masks[0][0:500, 0:500] = 1 38 masks[1][1000:1500, 1000:1500] = 2 39 masks[2][0:500, 0:500] = 3 40 masks[3][1000:1500, 1000:1500] = 4 41 segm = masks 42 43 imshow_det_bbox(image, bbox, label, segm, 44 class_names=['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 45 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 46 'sheep', 'sofa', 'train', 'tvmonitor'], 47 win_name="windows 98", 48 wait_time=5, 49 show=False, 50 out_file="test_browse_dataset_{}.jpg".format(str(index))) 51 52 index += 1 53 54 if os.path.exists("test_browse_dataset_0.jpg"): 55 os.remove("test_browse_dataset_0.jpg") 56 if os.path.exists("test_browse_dataset_1.jpg"): 57 os.remove("test_browse_dataset_1.jpg") 58 if os.path.exists("test_browse_dataset_2.jpg"): 59 os.remove("test_browse_dataset_2.jpg") 60 61 62if __name__ == "__main__": 63 test_browse_dataset() 64