使用mmdetection檢測並存儲結果

使用voc格式樣本訓練

修改mmdet/datasets/voc.py文件中類別,注意格式別出錯

CLASSES = ('bird-nest',)

然後修改配置文件,其中img_scale是你輸入圖片的寬高,num_classes=數據集類別數+1(背景),其他超參按照官網教程修改即可。

如果要做精度評價的話,要記得修改mmdet/core/evaluation/class_names.py中的類別跟數據集的類別一致。

檢測並保存結果(添加nms)

由於我在超算上運行,所以沒有做結果顯示,只是保存爲了jpg
首先是在根目錄下新建一個demo.py文件

import mmcv
import os
import numpy as np
from mmcv.runner import load_checkpoint
from mmdet.models import build_detector
from mmdet.apis import init_detector, inference_detector, show_result, draw_result

model = init_detector('configs/pascal_voc/faster_rcnn_r50_fpn_1x_voc0712.py', 'work_dirs/faster_rcnn_r50_fpn_1x_voc0712/epoch_8.pth', device='cuda:0')

input_dir = 'data/VOCdevkit/VOC2007/JPEGImages/'
out_dir = 'results/'

if not os.path.exists(out_dir):
    os.mkdir(out_dir)


def py_cpu_nms(dets, thresh):
    """Pure Python NMS baseline."""
    dets = np.array(dets)
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)

    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter + 0.000000000001)
        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep

def nms_cpu(total_detections, classnames, thresh=0.5):
    for i in range(len(classnames)):
        keep = py_cpu_nms(total_detections[i], thresh)
        total_detections[i] = total_detections[i][keep]
    return total_detections


files = os.listdir(input_dir)
if len(files) != 0:
    for file in files:
        name = os.path.splitext(file)[0]
        print ('detecting: ' + name)
        img = mmcv.imread(os.path.join(input_dir, file))
        img_resize = mmcv.imresize(img, (2000, 1500))
        result = inference_detector(model, img_resize)
        result = nms_cpu(result, model.CLASSES, 0.1)
        result_img = draw_result(img_resize, result, model.CLASSES, score_thr=0.3)
        out_file_path = os.path.join(out_dir, name + '.jpg')
        mmcv.imwrite(result_img, out_file_path)

然後是在mmdet/apis/inference.py中添加一個函數draw_result,這個不填寫也可以(直接調用show_result),我懶得改了

def draw_result(img, result, class_names, score_thr=0.3):

    return show_result(
        img, result, class_names, score_thr=score_thr, show=False)

如果添加了draw_result函數一定記得修改mmdet/apis/init.py文件

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章