用于可视化的工具主要是是位于utils下的visualizer文件,文件中的Visualizer类用于可视化,里面定义了多种方法,不光可以对数据集标签可视化,也可以对与预测结果进行可视化
直接上代码
from detectron2.data.datasets.coco import load_coco_json
from detectron2.data.detection_utils import read_image
#加载数据集链表字典(数据集应该先被注册过),"smoke_train"是之前注册过的数据集
data_dict = load_coco_json(json_file,image_root,"smoke_train")
from torch.utils.data.sampler import RandomSampler
#在数据集中随机采样一个图片用作数据的可视化
sampler = RandomSampler(data_dict,True,1)
index = next(iter(sampler))
#此图所对应的字典
image_dict = data_dict[index]
#读取图片
img = read_image(image_dict["file_name"],format='BGR')
#创建一个Visualizer对象
from detectron2.utils.visualizer import Visualizer
visualizer = Visualizer(img,metadata={'thing_classes':['smoke']})
img_visual = visualizer.draw_dataset_dict(image_dict)
#Visualizer会返回一个VisImage类
import cv2
cv2.imshow('sample',img_visual.get_image())
cv2.waitKey(0)
可视化结果,0是数据类编号,我用的数据集只有一类,都是smoke