标签：定义数据 Generic args RFB fd device path 1MB size

上一篇帖子我们介绍了如何修改源代码适应我们自己定义的数据集做一个多分类的目标检测任务。

Ultra-Light-Fast-Generic-Face-Detector-1MB的自定义数据集训练 | 多分类修改指南_Andrwin的博客-CSDN博客

训练好的权重会被保存在 项目根目录/models/train-version-RFB 这个文件夹里面，你可以根据训练日志选择损失最小的一个权重。

第一件事是新建一个分类标签列表，看过上一篇博客的朋友可能不理解为什么又要新建一个列表，因为原始代码对于列表读取的分隔符不一致，本着对原始代码最小改动的原则新建一个文件是最简单的方式：

首先在项目根目录\models\这个目录里面新建一个user_labels.txt

BACKGROUND
classA
classB
calssC
classD

需要注意的地方是第一行加入默认分类，注意全部字母大写，满足与预训练模型保持一致。

其他部分不再以英文逗号作为行结束符，直接换行就可以。

推理的代码主要是围绕detect_imgs.py修改的，由于代码量很少，我直接放出全部代码：

"""
This code is used to batch detect images in a folder.
"""
import argparse
import os
import sys

import cv2

from vision.ssd.config.fd_config import define_img_size

parser = argparse.ArgumentParser(
    description='detect_imgs')

parser.add_argument('--net_type', default="RFB", type=str,
                    help='The network architecture ,optional: RFB (higher precision) or slim (faster)')
parser.add_argument('--input_size', default=320, type=int,
                    help='define network input size,default optional value 128/160/320/480/640/1280')
parser.add_argument('--threshold', default=0.6, type=float,
                    help='score threshold')
parser.add_argument('--candidate_size', default=1500, type=int,
                    help='nms candidate size')
parser.add_argument('--path', default="test_data", type=str,
                    help='imgs dir')
parser.add_argument('--test_device', default="cuda:0", type=str,
                    help='cuda:0 or cpu')
args = parser.parse_args()
define_img_size(args.input_size)  # must put define_img_size() before 'import create_mb_tiny_fd, create_mb_tiny_fd_predictor'

from vision.ssd.mb_tiny_fd import create_mb_tiny_fd, create_mb_tiny_fd_predictor
from vision.ssd.mb_tiny_RFB_fd import create_Mb_Tiny_RFB_fd, create_Mb_Tiny_RFB_fd_predictor

result_path = "./detect_imgs_results"
label_path = "./models/user_labels.txt"
test_device = args.test_device

class_names = [name.strip() for name in open(label_path).readlines()]

if args.net_type == 'slim':
    model_path = "models/pretrained/version-slim-320.pth"
    # model_path = "models/pretrained/version-slim-640.pth"
    net = create_mb_tiny_fd(len(class_names), is_test=True, device=test_device)
    predictor = create_mb_tiny_fd_predictor(net, candidate_size=args.candidate_size, device=test_device)
elif args.net_type == 'RFB':
    model_path = "models/train-version-RFB/RFB-Epoch-199-Loss-0.3564375452697277.pth"
    # model_path = "models/pretrained/version-RFB-320.pth"
    # model_path = "models/pretrained/version-RFB-640.pth"
    net = create_Mb_Tiny_RFB_fd(len(class_names), is_test=True, device=test_device)
    predictor = create_Mb_Tiny_RFB_fd_predictor(net, candidate_size=args.candidate_size, device=test_device)
else:
    print("The net type is wrong!")
    sys.exit(1)
net.load(model_path)

if not os.path.exists(result_path):
    os.makedirs(result_path)
listdir = os.listdir(args.path)
sum = 0
for file_path in listdir:
    img_path = os.path.join(args.path, file_path)
    orig_image = cv2.imread(img_path)
    image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
    boxes, labels, probs = predictor.predict(image, args.candidate_size / 2, args.threshold)
    print(labels)
    sum += boxes.size(0)
    for i in range(boxes.size(0)):
        box = boxes[i, :]
        label = class_names[labels[i]]
        color = (255,255,255)
        cv2.rectangle(orig_image, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), color, 2)
        cv2.putText(orig_image, label, (int(box[0]), int(box[1] - 10)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)
    # cv2.putText(orig_image, str(boxes.size(0)), (30, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)
    cv2.imwrite(os.path.join(result_path, file_path), orig_image)
    print(f"Found {len(probs)} faces. The output image is {result_path}")
print(sum)

如果自己修改那么有如下位置需要注意：

第17行 '--input_size', default=640，把这个预设尺寸改为320，因为预训练权重就是320

第34行 改为自定义的分类标签列表 注意看它的分割形式

第23行 改为自己的测试图片文件夹路径

第44行 改为自己训练结束后的最佳权重

第68行 如果解除注释的话注意让box[0]和box[1]-10转为int类型，因为puttext函数不允许坐标值存在小数

由于数据比较敏感我就不放推理结果图了，效果我觉得还可以，媲美YOLO3.

可以看到推理结果非常快 一张图差不多4毫秒。

标签：定义数据,Generic,args,RFB,fd,device,path,1MB,size
来源： https://blog.csdn.net/Andrwin/article/details/120413366

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。