目录
一.开始
本文记录了实验mmdet训练完以后,使用训练得到的权值文件进行图片批量检测,并且将结果图片保存到文件夹的流程
二.代码
1.创建图片名字的txt文件
创建一个.py文件,然后将下列代码放到和图片同一级文件夹内,然后运行即可得到对应的txt文件,内容是图片的名字
# -*- coding:utf-8 -*-
import glob
imageList = glob.glob("*.jpg") # 图片所在文件夹的路径
f = open('val.txt', 'a') # 创建标签文件存放图片文件名
for item in imageList:
# print(item) # images_test/m1_a_018_1.jpg
img_name = item.split('/')[-1] # 图片文件名018.jpg
f.write(img_name + '\n') # 将图片文件名逐行写入txt
print('OK')
2.批量检测代码
使用下列代码,修改对应的参数,运行即可进行批量检测并且保存结果
from argparse import ArgumentParser
import os
from mmdet.apis import inference_detector, init_detector #, show_result_pyplot
import cv2
def show_result_pyplot(model, img, result, score_thr=0.3, fig_size=(15, 10)):
"""Visualize the detection results on the image.
Args:
model (nn.Module): The loaded detector.
img (str or np.ndarray): Image filename or loaded image.
result (tuple[list] or list): The detection result, can be either
(bbox, segm) or just bbox.
score_thr (float): The threshold to visualize the bboxes and masks.
fig_size (tuple): Figure size of the pyplot figure.
"""
if hasattr(model, 'module'):
model = model.module
img = model.show_result(img, result, score_thr=score_thr, show=False)
return img
# plt.figure(figsize=fig_size)
# plt.imshow(mmcv.bgr2rgb(img))
# plt.show()
def main():
# config文件
config_file = './configs/faster_rcnn/faster_rcnn_r50_fpn_1x_coco.py'
# 训练好的模型
checkpoint_file = './work_dirs/faster_rcnn_r50_fpn_1x_coco/epoch_200.pth'
# model = init_detector(config_file, checkpoint_file)
model = init_detector(config_file, checkpoint_file, device='cuda:0')
# 图片路径
img_dir = './data/val/'
# 检测后存放图片路径
out_dir = './faster_rcnn_result/'
if not os.path.exists(out_dir):
os.mkdir(out_dir)
# 测试集的图片名称txt
test_path = './val.txt'
fp = open(test_path, 'r')
test_list = fp.readlines()
count = 0
imgs = []
for test in test_list:
test = test.replace('\n', '')
test = test.split('.')[0] # 如果test里面内容的名字是xxx.jpg,需要这行语句,是因为生成的图片会出现.jpg.jpg,否则不需要。
name = img_dir + test + '.jpg'
count += 1
print('model is processing the {}/{} images.'.format(count, len(test_list)))
# result = inference_detector(model, name)
# model = init_detector(config_file, checkpoint_file, device='cuda:0')
result = inference_detector(model, name)
img = show_result_pyplot(model, name, result, score_thr=0.8)
cv2.imwrite("{}/{}.jpg".format(out_dir, test), img)
if __name__ == '__main__':
main()
三.结语
至此使用mmdet批量检测图片并保存结果的流程就结束了。