mmdet开源项目，有很多人使用，但对于初学者，不知如何调用，尽管能够训练，却有时也为测试代码泛难，趁着项目需要，我个人编写了一份测试代码，共大家学习。

准备：

1.config parameters 参数py文件

2.checkpoint  训练好的权重.pth文件

3.classes  若无，则数字替代

4.测试图片路径  该文件存放需要测试图片，也可以是多个文件下存放，代码会自动寻找

 

详细如下：

 

 

from mmdet.apis import inference_detector, init_detector
import cv2
import os
import numpy as np
import pandas as pd

from tqdm import tqdm
import time
os.environ["CUDA_VISIBLE_DEVICES"] = "7"


class Model():
    def __init__(self, root_config, root_checkpoint, **kwargs):
        self.model = init_detector(root_config, root_checkpoint)  # 模型初始化
        self.thr_ok = kwargs.get('thr_ok', 0.05)
        self.classes = kwargs.get('classes', None)
        self.color = self.get_color()
        self.img_foramt = ['.jpg', '.JPG', '.bmp', '.png']

    def get_color(self):
        color = dict(red=(0, 0, 255),
                     green=(0, 255, 0),
                     blue=(255, 0, 0),
                     cyan=(255, 255, 0),
                     yellow=(0, 255, 255),
                     magenta=(255, 0, 255),
                     white=(255, 255, 255),
                     black=(0, 0, 0))
        return color

    def model_test(self, result, img_name, classes, thr_ok=0.05):
        output_bboxes = []
        json_dict = []
        total_bbox = []
        for id, boxes in enumerate(result):  # loop for categories
            category_id = id + 1
            if len(boxes) != 0:
                for box in boxes:  # loop for bbox
                    conf = box[4]
                    if conf > thr_ok:
                        total_bbox.append(list(box) + [category_id])

        bboxes = np.array(total_bbox)
        best_bboxes = bboxes
        output_bboxes.append(best_bboxes)
        for bbox in best_bboxes:
            coord = [round(i, 2) for i in bbox[:4]]
            conf = bbox[4]
            category = classes[int(bbox[5]) - 1] if classes is not None else int(bbox[5])
            json_dict.append({'img_name': img_name, 'cats': category, 'bbox': coord, 'score': conf})
        det_df = pd.DataFrame(json_dict, columns=['img_name', 'cats', 'bbox', 'score'])
        return det_df

    def single_test(self, root_img):  # 单张图片模型测试
        img = cv2.imread(root_img)
        model_result = inference_detector(self.model, img)
        return model_result

    def run(self, img_root):
        model_result = self.single_test(img_root)
        img_name = self.get_strfile(img_root, pos=-1)
        result_df = self.model_test(model_result, img_name, self.classes, thr_ok=self.thr_ok)
        img_name_lst, cat_lst, box_lst, score_lst = self.pd2lst(result_df)
        return img_name_lst, cat_lst, box_lst, score_lst

    def pd2lst(self, result_df):
        img_name_lst, cat_lst, box_lst, score_lst = [], [], [], []
        if len(result_df) > 0:
            for i in range(len(result_df)):
                img_name_lst.append(result_df.loc[i]['img_name'])
                cat_lst.append(result_df.loc[i]['cats'])
                box_lst.append(result_df.loc[i]['bbox'])
                score_lst.append(result_df.loc[i]['score'])
        return img_name_lst, cat_lst, box_lst, score_lst

    def draw_bbox(self, img, cat_lst, box_lst, score_lst,
                  bbox_color='green',
                  text_color='green',
                  thickness=1,
                  font_scale=0.5
                  ):
        for j, cat in enumerate(cat_lst):
            x1, y1, x2, y2 = np.array(box_lst[j]).astype(np.int32)
            bbox_color_new = self.color[bbox_color]
            cv2.rectangle(img, (x1, y1), (x2, y2), bbox_color_new, thickness=thickness)
            score = round(score_lst[j], 4)
            text_color_new = self.color[text_color]
            label_text = '{}:{}'.format(str(cat), str(score))
            cv2.putText(img, label_text, (x1, y1 - 2), cv2.FONT_HERSHEY_COMPLEX, font_scale, text_color_new)
        return img
    def get_strfile(self, file_str, pos=-1):
        # 得到file_str / or \\ 的最后一个名称
        endstr_f_filestr = file_str.split('\\')[pos] if '\\' in file_str else file_str.split('/')[pos]
        return endstr_f_filestr

    def build_dir(self,out_dir):
        # 构建文件
        if not os.path.exists(out_dir):
            os.makedirs(out_dir)
        return out_dir

    def show_img(self,img):
        import matplotlib.pyplot as plt
        plt.imshow(img)
        plt.show()

def get_files_root(root):
    '''
    :return: 寻找root下面有文件夹的路径，输出所有文件夹绝对路径的列表
    '''
    files_lst = [root]
    result_lst = files_lst
    if build_files(root) == []:
        result_lst = files_lst
    else:
        is_while = True
        files_all_path = [file for file in files_lst]
        while is_while:
            for file_root in files_lst:
                F1 = build_files(file_root)
                for F1 in F1:
                    files_all_path.append(F1)
            is_while = False
            # 排除主文件夹
            record = np.ones((len(files_all_path)))
            for i, F3 in enumerate(files_all_path):
                F3 = files_all_path[i]
                for j, F4 in enumerate(files_all_path):
                    if F3 + '\\' in F4 or F3 + '/' in F4:
                        record[i] = 0
                        break
            # 将需要循环聚集
            files_lst = []
            for i, F3 in enumerate(files_all_path):
                if record[i] == 1:
                    files_lst.append(F3)
            # 判断是否有子文件夹
            for F4 in files_lst:
                file_judge = build_files(F4)
                if file_judge != []:
                    is_while = True
                    break

            result_lst = files_lst

    return result_lst

def build_files(root):
    '''
    :得到该路径下的所有文件
    '''
    files = [os.path.join(root, file) for file in os.listdir(root)]
    files_true = []
    for file in files:
        if not os.path.isfile(file):
            files_true.append(file)
    return files_true


def single_main(model,root_img,work_dir):
    # 一张图片测试所有集合
    img_name_lst, cat_lst, box_lst, score_lst = model.run(root_img)  #
    img = cv2.imread(root_img)
    img = model.draw_bbox(img, cat_lst, box_lst, score_lst)
    file_name = model.get_strfile(root_img, pos=-2)
    out_file = model.build_dir(os.path.join(work_dir, file_name))
    img_name = model.get_strfile(root_img, pos=-1)
    cv2.imwrite(os.path.join(out_file, img_name), img)

def main(root,model,work_dir):
    root_files=get_files_root(root)
    num=0
    for file_path in tqdm(root_files):
        for name in tqdm(os.listdir(file_path)):
            if name[-4:] in model.img_foramt:
                root_img=os.path.join(file_path,name)
                single_main(model,root_img,work_dir)
                num+=1

    print('num of images:',num)



if __name__ == '__main__':
    root_config = '/data/sdv3/tangjun/xmtm/xmtm_pointer/code/model_new_meter/parameters.py'
    root_checkpoint = '/data/sdv3/tangjun/xmtm/xmtm_pointer/code/model_new_meter/model.pth'
    root = '/data/sdv3/tangjun/xmtm/xmtm_pointer/data/data_0512/data_step_two/train_step2/train'  # 只需要测试文件夹 
    work_dir='/data/sdv3/cj/First_Blood/tj/code/mmdet50/123/78/90'
    info = {'classes': None}

    time_start = time.time()
    model = Model(root_config, root_checkpoint, **info)  # 类实列化，也是初始化
    # img_name_lst, cat_lst, box_lst, score_lst = model.run(root_img) #单张图片的预测

    main(root, model, work_dir)

    time_end=time.time()

    time_gap=time_end-time_start
    print('time gap:',time_gap)