# Copyright (c) OpenMMLab. All rights reserved. import argparse import os.path as osp import numpy as np from mmcv.image import imwrite from mmengine.config import Config, DictAction from mmengine.registry import init_default_scope from mmengine.utils import ProgressBar from mmdet.registry import DATASETS, VISUALIZERS from mmdet.structures.bbox import BaseBoxes def parse_args(): parser = argparse.ArgumentParser(description='Browse a dataset') parser.add_argument('config', help='train config file path') parser.add_argument( '--output-dir', '-o', default=None, type=str, help='If there is no display interface, you can save it') parser.add_argument('--not-show', default=False, action='store_true') parser.add_argument('--show-num', '-n', type=int, default=30) parser.add_argument('--shuffle', default=False, action='store_true') parser.add_argument( '--show-interval', type=float, default=0, help='the interval of show (s)') parser.add_argument( '--cfg-options', nargs='+', action=DictAction, help='override some settings in the used config, the key-value pair ' 'in xxx=yyy format will be merged into config file. If the value to ' 'be overwritten is a list, it should be like key="[a,b]" or key=a,b ' 'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" ' 'Note that the quotation marks are necessary and that no white space ' 'is allowed.') args = parser.parse_args() return args def draw_all_character(visualizer, characters, w): start_index = 2 y_index = 5 for char in characters: if isinstance(char, str): visualizer.draw_texts( str(char), positions=np.array([start_index, y_index]), colors=(0, 0, 0), font_families='monospace') start_index += len(char) * 8 else: visualizer.draw_texts( str(char[0]), positions=np.array([start_index, y_index]), colors=char[1], font_families='monospace') start_index += len(char[0]) * 8 if start_index > w - 10: start_index = 2 y_index += 15 drawn_text = visualizer.get_image() return drawn_text def main(): args = parse_args() cfg = Config.fromfile(args.config) if args.cfg_options is not None: cfg.merge_from_dict(args.cfg_options) assert args.show_num > 0 # register all modules in mmdet into the registries init_default_scope(cfg.get('default_scope', 'mmdet')) dataset = DATASETS.build(cfg.train_dataloader.dataset) visualizer = VISUALIZERS.build(cfg.visualizer) visualizer.dataset_meta = dataset.metainfo dataset_index = list(range(len(dataset))) if args.shuffle: import random random.shuffle(dataset_index) progress_bar = ProgressBar(len(dataset)) for i in dataset_index[:args.show_num]: item = dataset[i] img = item['inputs'].permute(1, 2, 0).numpy() data_sample = item['data_samples'].numpy() gt_instances = data_sample.gt_instances tokens_positive = data_sample.tokens_positive gt_labels = gt_instances.labels base_name = osp.basename(item['data_samples'].img_path) name, extension = osp.splitext(base_name) out_file = osp.join(args.output_dir, name + '_' + str(i) + extension) if args.output_dir is not None else None img = img[..., [2, 1, 0]] # bgr to rgb gt_bboxes = gt_instances.get('bboxes', None) if gt_bboxes is not None and isinstance(gt_bboxes, BaseBoxes): gt_instances.bboxes = gt_bboxes.tensor print(data_sample.text) dataset_mode = data_sample.dataset_mode if dataset_mode == 'VG': max_label = int(max(gt_labels) if len(gt_labels) > 0 else 0) palette = np.random.randint(0, 256, size=(max_label + 1, 3)) bbox_palette = [tuple(c) for c in palette] # bbox_palette = get_palette('random', max_label + 1) colors = [bbox_palette[label] for label in gt_labels] visualizer.set_image(img) for label, bbox, color in zip(gt_labels, gt_bboxes, colors): visualizer.draw_bboxes( bbox, edge_colors=color, face_colors=color, alpha=0.3) visualizer.draw_bboxes(bbox, edge_colors=color, alpha=1) drawn_img = visualizer.get_image() new_image = np.ones((100, img.shape[1], 3), dtype=np.uint8) * 255 visualizer.set_image(new_image) gt_tokens_positive = [ tokens_positive[label] for label in gt_labels ] split_by_character = [char for char in data_sample.text] characters = [] start_index = 0 end_index = 0 for w in split_by_character: end_index += len(w) is_find = False for i, positive in enumerate(gt_tokens_positive): for p in positive: if start_index >= p[0] and end_index <= p[1]: characters.append([w, colors[i]]) is_find = True break if is_find: break if not is_find: characters.append([w, (0, 0, 0)]) start_index = end_index drawn_text = draw_all_character(visualizer, characters, img.shape[1]) drawn_img = np.concatenate((drawn_img, drawn_text), axis=0) else: gt_labels = gt_instances.labels text = data_sample.text label_names = [] for label in gt_labels: label_names.append(text[ tokens_positive[label][0][0]:tokens_positive[label][0][1]]) gt_instances.label_names = label_names data_sample.gt_instances = gt_instances visualizer.add_datasample( base_name, img, data_sample, draw_pred=False, show=False, wait_time=0, out_file=None) drawn_img = visualizer.get_image() new_image = np.ones((100, img.shape[1], 3), dtype=np.uint8) * 255 visualizer.set_image(new_image) characters = [char for char in text] drawn_text = draw_all_character(visualizer, characters, img.shape[1]) drawn_img = np.concatenate((drawn_img, drawn_text), axis=0) if not args.not_show: visualizer.show( drawn_img, win_name=base_name, wait_time=args.show_interval) if out_file is not None: imwrite(drawn_img[..., ::-1], out_file) progress_bar.update() if __name__ == '__main__': main()