FasterRCNN所有类别放在一张图上并用不同的颜色框

FasterRCNN 所有类别的框放在一张图上

在使用faster RCNN进行目标检测的时候发现不同类别的框不在同一张图上。所以进行一下的修改。
修改tools/demo.py
首先增加一个color map，使每个类别都有不同的颜色框。

COLORMAP = {
     'aeroplane': 'k', 'bicycle': 'y', 'bird': 'crimson', 'boat': 'coral', 'bottle': 'aqua', 'bus': 'm', 'car': 'c', 'cat': 'brown', 'chair': 'grey', 'cow': 'b', 'diningtable': 'teal', 'dog': 'navy', 'horse': 'wheat', 'motorbike': 'g', 'person': 'r', 'pottedplant': 'peru', 'sheep': 'skyblue', 'sofa': 'sienna', 'train': 'yellow', 'tvmonitor': 'olive'}

将不同类别的框放在同一张图上：

def vis_detections(ax, im, image_name, class_name, dets, thresh=0.5):
    """
        将ax放入参数
    """
    inds = np.where(dets[:, -1] >= thresh)[0]
    if len(inds) == 0:
        return
        
    for i in inds:
        bbox = dets[i, :4]
        score = dets[i, -1]

        ax.add_patch(
            plt.Rectangle((bbox[0], bbox[1]),
                          bbox[2] - bbox[0],
                          bbox[3] - bbox[1], fill=False,
                          edgecolor=COLORMAP[class_name], linewidth=3.5)
            ) 
            #edgecolor使用与类别名对应的colormap中的颜色
        ax.text(bbox[0], bbox[1] - 2,
                '{:s} {:.3f}'.format(class_name, score),
                bbox=dict(facecolor='blue', alpha=0.5),
                fontsize=14, color='white')

    ax.set_title(('{} detections with '
                  'p({} | box) >= {:.1f}').format(class_name, class_name,
                                                  thresh),
                  fontsize=14) 
    # plt.axis('off')
    # plt.tight_layout()
    # plt.draw()
    

def demo(sess, net, image_name):
    """Detect object classes in an image using pre-computed object proposals."""
    dataset_dir = './dataset' #dataset路径
    # Load the demo image
    im_file = os.path.join(dataset_dir, image_name+'.png')
    im = cv2.imread(im_file)

    # Detect all object classes and regress object bounds
    timer = Timer()
    timer.tic()
    scores, boxes = im_detect(sess, net, im)
    timer.toc()
    print('Detection took {:.3f}s for {:d} object proposals'.format(timer.total_time, boxes.shape[0]))  

    # Visualize detections for each class
    CONF_THRESH = 0.8
    NMS_THRESH = 0.3

    im = im[:, :, (2, 1, 0)]
    fig, ax = plt.subplots(figsize=(12, 12))
    ax.imshow(im, aspect='equal')

    # iou_score=[]
    for cls_ind, cls in enumerate(CLASSES[1:]):
        cls_ind += 1 # because we skipped background
        cls_boxes = boxes[:, 4*cls_ind:4*(cls_ind + 1)]
        cls_scores = scores[:, cls_ind]
        dets = np.hstack((cls_boxes, cls_scores[:, np.newaxis])).astype(np.float32)
        keep = nms(dets, NMS_THRESH)
        dets = dets[keep, :]
        vis_detections(ax, im,image_name, cls, dets, thresh=CONF_THRESH)
        plt.axis('off')
        plt.tight_layout()
        plt.draw()
    #Save images
    save_dir = './result'
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    if SAVE_VIS_IMAGES:
            save_path = os.path.join(save_dir, image_name+ '.png')
    plt.savefig(save_path)

参考了以下这篇：
https://blog.csdn.net/weixin_42111393/article/details/82940681