mmdetection源码解读（九）

如何展示mmdet的检测结果：

mmdetection/tools/bbox_show.py

import numpy as np
import mmcv
import cv2
import json
from pycocotools.coco import COCO
from mmcv import color_val
from mmdet.apis import show_result
from argparse import ArgumentParser
import os

class_names = ['category_1', 'category_2', 'category_3']

def det_bbox_show(args=None):
    """ 这个是用来展示检测结果的 """
    with open(args.json_file, 'r') as f:
        jsondata = json.load(f)
        results = {}
        for result in jsondata:
            result['bbox'][2] = result['bbox'][0] + result['bbox'][2]
            result['bbox'][3] = result['bbox'][1] + result['bbox'][3]
            result['bbox'].append(result['score'])
            if result['image_id'] in results.keys():
                results[result['image_id']][result['category_id']-1] = np.concatenate((results[result['image_id']][result['category_id']-1], np.array(result['bbox'], dtype=np.float32).reshape(-1,5)))
            else:
                results[result['image_id']] = []
                for i in range(len(class_names)):
                    if result['category_id']-1 == i:
                        results[result['image_id']].append(np.array(result['bbox'], dtype=np.float32).reshape(-1,5))
                    else:
                       results[result['image_id']].append(np.zeros(shape=(0,5),dtype=np.float32))
    for k, v in results.items():
        show_result(args.img_prex + '/' + k + '.jpg', v, class_names, score_thr=0.7, show=False, out_file=args.outfile + '/{}.jpg'.format(k))


def gt_bbox_show(args):
    """这个是用来展示Ground Truth bbox的"""
    coco = COCO(args.json_file)
    img2ann = coco.imgToAnns
    imgids = coco.getImgIds()
    for imgid in imgids:
        filename = args.img_prex + '/{}.jpg'.format(imgid)
        img = mmcv.imread(filename)
        bbox_color = color_val('red')
        text_color = color_val('red')
        anns = img2ann[imgid]
        for ann in anns:
            bbox_int = np.array(ann['bbox']).astype(np.int32)
            left_top = (bbox_int[0], bbox_int[1])
            right_bottom = (bbox_int[0] + bbox_int[2], bbox_int[1] + bbox_int[3])
            label = ann['category_id'] - 1
            cv2.rectangle(img, left_top, right_bottom, bbox_color, thickness=1)
            label_text = class_names[label]
            cv2.putText(img, label_text, (bbox_int[0], bbox_int[1] - 2), cv2.FONT_HERSHEY_COMPLEX, 0.5, text_color)
        mmcv.imwrite(img, args.outfile + '/{}.jpg'.format(imgid))

def main():
    parser = ArgumentParser(description='COCO bbox show Tool')
    parser.add_argument(
        '--gt-bbox-show',
        action='store_true',
        help='show gt bbox'
    )
    parser.add_argument(
        '--json_file',
        type=str,
        help='json file(results or ann file)'
    )
    parser.add_argument(
        '--img_prex',
        type=str,
        help='img dataset prex'
    )
    parser.add_argument(
        '--outfile',
        type=str,
        help='outfile'
    )
    args = parser.parse_args()
    if not os.path.exists(args.outfile):
        os.makedirs(args.outfile)
    if args.gt_bbox_show:
        gt_bbox_show(args)
    else:
        det_bbox_show(args)

if __name__ == '__main__':
    import time
    tim = time.time()
    main()
    print("total time = {}".format(time.time() - tim))

这段程序前提是图片名称和图片id要保持一致，像coco这种图片名前面补好多0的就不能用了，需要debug一下。