基于mindspore的口罩检测训练与在线推理

mindspore安装地址：https://www.mindspore.cn/install

mindspore开源模型库：https://gitee.com/mindspore/models

测试平台为昇腾Atlas800训练服务器，Ubuntu18.04系统，搭载昇腾910AI加速芯片，FP16算力为320T。

数据集为VOC_MASK口罩数据集，在使用数据集之前需要对数据集文件进行转换，将voc标注文件由xml转为coco的json格式。

转换脚本：

import os
import glob
import json
import shutil
import numpy as np
import xml.etree.ElementTree as ETSTART_BOUNDING_BOX_ID = 1def get(root, name):return root.findall(name)def get_and_check(root, name, length):vars = root.findall(name)if len(vars) == 0:raise NotImplementedError('Can not find %s in %s.'%(name, root.tag))if length > 0 and len(vars) != length:raise NotImplementedError('The size of %s is supposed to be %d, but is %d.'%(name, length, len(vars)))if length == 1:vars = vars[0]return varsdef convert(xml_list, json_file):json_dict = {"info":['none'], "license":['none'], "images": [], "annotations": [], "categories": []}categories = pre_define_categories.copy()bnd_id = START_BOUNDING_BOX_IDall_categories = {}for index, line in enumerate(xml_list):xml_f = linetree = ET.parse(xml_f)root = tree.getroot()filename = os.path.basename(xml_f)[:-4] + ".jpg"image_id = indexsize = get_and_check(root, 'size', 1)width = int(get_and_check(size, 'width', 1).text)height = int(get_and_check(size, 'height', 1).text)image = {'file_name': filename, 'height': height, 'width': width, 'id':image_id}json_dict['images'].append(image)for obj in get(root, 'object'):category = get_and_check(obj, 'name', 1).textif category in all_categories:all_categories[category] += 1else:all_categories[category] = 1if category not in categories:if only_care_pre_define_categories:continuenew_id = len(categories) + 1print("[warning] category '{}' not in 'pre_define_categories'({}), create new id: {} automatically".format(category, pre_define_categories, new_id))categories[category] = new_idcategory_id = categories[category]bndbox = get_and_check(obj, 'bndbox', 1)xmin = int(float(get_and_check(bndbox, 'xmin', 1).text))ymin = int(float(get_and_check(bndbox, 'ymin', 1).text))xmax = int(float(get_and_check(bndbox, 'xmax', 1).text))ymax = int(float(get_and_check(bndbox, 'ymax', 1).text))assert(xmax > xmin), "xmax <= xmin, {}".format(line)assert(ymax > ymin), "ymax <= ymin, {}".format(line)o_width = abs(xmax - xmin)o_height = abs(ymax - ymin)ann = {'area': o_width*o_height, 'iscrowd': 0, 'image_id':image_id, 'bbox':[xmin, ymin, o_width, o_height],'category_id': category_id, 'id': bnd_id, 'ignore': 0,'segmentation': []}json_dict['annotations'].append(ann)bnd_id = bnd_id + 1for cate, cid in categories.items():cat = {'supercategory': 'none', 'id': cid, 'name': cate}json_dict['categories'].append(cat)json_fp = open(json_file, 'w')json_str = json.dumps(json_dict)json_fp.write(json_str)json_fp.close()print("------------create {} done--------------".format(json_file))print("find {} categories: {} -->>> your pre_define_categories {}: {}".format(len(all_categories), all_categories.keys(), len(pre_define_categories), pre_define_categories.keys()))print("category: id --> {}".format(categories))print(categories.keys())print(categories.values())if __name__ == '__main__':# xml标注文件夹xml_dir = 'Annotations'# 训练数据的josn文件save_json_train = 'annotation/train.json'# 验证数据的josn文件save_json_val = 'annotation/val.json'# 类别classes = ['face','mask']pre_define_categories = {}for i, cls in enumerate(classes):pre_define_categories[cls] = i + 1print(pre_define_categories)only_care_pre_define_categories = True# 训练数据集比例 train_ratio = 0.9xml_list = glob.glob(xml_dir + "/*.xml")  xml_list = np.sort(xml_list)np.random.seed(100)np.random.shuffle(xml_list)train_num = int(len(xml_list)*train_ratio)print('训练样本数目是 {}'.format(train_num))print('测试样本数目是 {}'.format(len(xml_list) - train_num))xml_list_train = xml_list[:train_num]xml_list_val = xml_list[train_num:]# 对训练数据集对应的xml进行coco转换   convert(xml_list_train, save_json_train)# 对验证数据集的xml进行coco转换convert(xml_list_val, save_json_val)

模型采用mindspore官方仓中的yolov5，训练之前需要修改配置文件和tain.py脚本，官方demo采用的是coco2017训练集，相关设置参考coco进行修改。将分类数修改为num_classes=2，标签修改为labels: [ 'face', 'mask']，coco_ids修改为coco_ids: [1, 2]。

训练脚本：

import os
import time
import mindspore as ms
import mindspore.nn as nn
import mindspore.communication as comm
from mindspore import context
import cv2from src.yolo import YOLOV5, YoloWithLossCell
from src.logger import get_logger
from src.util import AverageMeter, get_param_groups, cpu_affinity
from src.lr_scheduler import get_lr
from src.yolo_dataset import create_yolo_dataset
from src.initializer import default_recurisive_init, load_yolov5_paramsfrom model_utils.config import config
from model_utils.device_adapter import get_device_idms.set_seed(1)def init_distribute():comm.init()config.rank = comm.get_rank()config.group_size = comm.get_group_size()ms.set_auto_parallel_context(parallel_mode=ms.ParallelMode.DATA_PARALLEL, gradients_mean=True,device_num=config.group_size)def train_preprocess():if config.lr_scheduler == 'cosine_annealing' and config.max_epoch > config.T_max:config.T_max = config.max_epochconfig.lr_epochs = list(map(int, config.lr_epochs.split(',')))config.data_root = os.path.join(config.data_dir, config.train_img_dir)config.annFile = os.path.join(config.data_dir, config.train_json_file)device_id = get_device_id()context.set_context(mode=context.GRAPH_MODE, device_target=config.device_target, device_id=device_id)if config.is_distributed:# init distributedinit_distribute()# for promoting performance in GPU deviceif config.device_target == "GPU" and config.bind_cpu:cpu_affinity(config.rank, min(config.group_size, config.device_num))# logger module is managed by config, it is used in other function. e.x. config.logger.info("xxx")config.logger = get_logger(config.output_dir, config.rank)config.logger.save_args(config)def run_train():train_preprocess()loss_meter = AverageMeter('loss')dict_version = {'yolov5s': 0, 'yolov5m': 1, 'yolov5l': 2, 'yolov5x': 3}network = YOLOV5(is_training=True, version=dict_version[config.yolov5_version])# default is kaiming-normaldefault_recurisive_init(network)load_yolov5_params(config, network)network = YoloWithLossCell(network)ds = create_yolo_dataset(image_dir=config.data_root, anno_path=config.annFile, is_training=True,batch_size=config.per_batch_size, device_num=config.group_size,rank=config.rank, config=config)config.logger.info('Finish loading dataset')steps_per_epoch = ds.get_dataset_size()lr = get_lr(config, steps_per_epoch)opt = nn.Momentum(params=get_param_groups(network), momentum=config.momentum, learning_rate=ms.Tensor(lr),weight_decay=config.weight_decay, loss_scale=config.loss_scale)network = nn.TrainOneStepCell(network, opt, config.loss_scale // 2)network.set_train()data_loader = ds.create_tuple_iterator(do_copy=False)first_step = Truet_end = time.time()for epoch_idx in range(config.max_epoch):for step_idx, data in enumerate(data_loader):images = data[0]input_shape = images.shape[2:4]input_shape = ms.Tensor(tuple(input_shape[::-1]), ms.float32)loss = network(images, data[2], data[3], data[4], data[5], data[6],data[7], input_shape)loss_meter.update(loss.asnumpy())# it is used for loss, performance output per config.log_interval steps.if (epoch_idx * steps_per_epoch + step_idx) % config.log_interval == 0:time_used = time.time() - t_endif first_step:fps = config.per_batch_size * config.group_size / time_usedper_step_time = time_used * 1000first_step = Falseelse:fps = config.per_batch_size * config.log_interval * config.group_size / time_usedper_step_time = time_used / config.log_interval * 1000config.logger.info('epoch[{}], iter[{}], {}, fps:{:.2f} imgs/sec, ''lr:{}, per step time: {}ms'.format(epoch_idx + 1, step_idx + 1,loss_meter, fps, lr[step_idx], per_step_time))t_end = time.time()loss_meter.reset()if config.rank == 0:ckpt_name = os.path.join(config.output_dir, "yolov5_{}_{}.ckpt".format(epoch_idx + 1, steps_per_epoch))ms.save_checkpoint(network, ckpt_name)config.logger.info('==========end training===============')if __name__ == "__main__":run_train()

完成训练配置后，进入scripts文件夹，运行run_standalone_train.sh脚本，开始模型训练。训练过程观察输出日志中loss变化情况。训练完成后，取出保存的ckpt文件进行在线推理测试。

推理流程：

1、使用mindspore构建模型；

2、将权重文件ckpt加载进网络；

3、预处理图像，送入模型进行推理；

4、获取推理结果，yolov5推理会返回三个列表，分别对应80*80,40*40,20*20三个不同尺度的预测框，每个预测框包含分类数+5个结果，这里模型只有戴口罩和不带口罩两种检测结果，那么一个检测框对应输出的是7个值(中心点x坐标，中心点y坐标，检测框宽，检测框高，检测框置信度，分类预测列表);

5、对检测框进行预处理，将坐标转换为检测框的左上角和右下角坐标；

6、进行非极大值抑制，删除无效检测框；

7、将检测结果绘制到原图。

mindspore在线推理脚本：

import os
import numpy as np
import mindspore as ms
from src.yolo import YOLOV5
import cv2
import matplotlib.pyplot as pltdef nms(pred, conf_thres, iou_thres):# 置信度抑制，小于置信度阈值则删除conf = pred[..., 4] > conf_thresbox = pred[conf == True]# 类别获取cls_conf = box[..., 5:]cls = []for i in range(len(cls_conf)):cls.append(int(np.argmax(cls_conf[i])))# 获取类别total_cls = list(set(cls))  #删除重复项，获取出现的类别标签列表,example=[0, 17]output_box = []   #最终输出的预测框# 不同分类候选框置信度for i in range(len(total_cls)):clss = total_cls[i]   #当前类别标签# 从所有候选框中取出当前类别对应的所有候选框cls_box = []for j in range(len(cls)):if cls[j] == clss:box[j][5] = clsscls_box.append(box[j][:6])cls_box = np.array(cls_box)box_conf = cls_box[..., 4]   #取出候选框置信度box_conf_sort = np.argsort(box_conf)   #获取排序后索引max_conf_box = cls_box[box_conf_sort[len(box_conf) - 1]]output_box.append(max_conf_box)   #将置信度最高的候选框输出为第一个预测框cls_box = np.delete(cls_box, 0, 0)  #删除置信度最高的候选框while len(cls_box) > 0:max_conf_box = output_box[len(output_box) - 1]     #将输出预测框列表最后一个作为当前最大置信度候选框del_index = []for j in range(len(cls_box)):current_box = cls_box[j]      #当前预测框interArea = getInter(max_conf_box, current_box)    #当前预测框与最大预测框交集iou = getIou(max_conf_box, current_box, interArea)  # 计算交并比if iou > iou_thres:del_index.append(j)   #根据交并比确定需要移出的索引cls_box = np.delete(cls_box, del_index, 0)   #删除此轮需要移出的候选框if len(cls_box) > 0:output_box.append(cls_box[0])cls_box = np.delete(cls_box, 0, 0)return output_box#计算并集
def getIou(box1, box2, inter_area):box1_area = box1[2] * box1[3]box2_area = box2[2] * box2[3]union = box1_area + box2_area - inter_areaiou = inter_area / unionreturn iou#计算交集
def getInter(box1, box2):box1_x1, box1_y1, box1_x2, box1_y2 = box1[0], box1[1], \box1[0] + box1[2], box1[1] + box1[3]box2_x1, box2_y1, box2_x2, box2_y2 = box2[0], box2[1], \box2[0] + box2[2], box2[1] + box2[3]if box1_x1 > box2_x2 or box1_x2 < box2_x1:return 0if box1_y1 > box2_y2 or box1_y2 < box2_y1:return 0x_list = [box1_x1, box1_x2, box2_x1, box2_x2]x_list = np.sort(x_list)x_inter = x_list[2] - x_list[1]y_list = [box1_y1, box1_y2, box2_y1, box2_y2]y_list = np.sort(y_list)y_inter = y_list[2] - y_list[1]inter = x_inter * y_interreturn interdef draw(img, pred):img_ = img.copy()if len(pred):for detect in pred:x1 = int(detect[0])y1 = int(detect[1])x2 = int(detect[0] + detect[2])y2 = int(detect[1] + detect[3])score = detect[4]cls = detect[5]labels = ['no_mask', 'mask']print(x1, y1, x2, y2, score, cls)img_ = cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 1)text = labels[int(cls)] + ':' + str(score)cv2.putText(img, text, (x1, y1 + 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 1,)return img_def load_parameters(network, filename):param_dict = ms.load_checkpoint(filename)param_dict_new = {}for key, values in param_dict.items():if key.startswith('moments.'):continueelif key.startswith('yolo_network.'):param_dict_new[key[13:]] = valueselse:param_dict_new[key] = valuesms.load_param_into_net(network, param_dict_new)def main(ckpt_file, img):orig_h, orig_w = img.shape[:2]ms.set_context(mode=ms.GRAPH_MODE, device_target='CPU', device_id=0)dict_version = {'yolov5s': 0, 'yolov5m': 1, 'yolov5l': 2, 'yolov5x': 3}network = YOLOV5(is_training=False, version=dict_version['yolov5s'])if os.path.isfile(ckpt_file):load_parameters(network, ckpt_file)else:raise FileNotFoundError(f"{ckpt_file} is not a filename.")network.set_train(False)input_shape = ms.Tensor(tuple([640, 640]), ms.float32)img = cv2.resize(img, (640, 640), cv2.INTER_LINEAR)img = img[:, :, ::-1].transpose((2, 0, 1))img = img / 255.img = np.expand_dims(img, axis=0)image = np.concatenate((img[..., ::2, ::2], img[..., 1::2, ::2],img[..., ::2, 1::2], img[..., 1::2, 1::2]), axis=1)image = ms.Tensor(image, dtype=ms.float32)output_big, output_me, output_small = network(image, input_shape)output_big = output_big.asnumpy()output_me = output_me.asnumpy()output_small = output_small.asnumpy()output_small = np.squeeze(output_small)output_small = np.reshape(output_small, [19200, 7])output_me = np.squeeze(output_me)output_me = np.reshape(output_me, [4800, 7])output_big = np.squeeze(output_big)output_big = np.reshape(output_big, [1200, 7])result = np.vstack([output_small, output_me, output_big])for i in range(len(result)):x = result[i][0] * orig_wy = result[i][1] * orig_hw = result[i][2] * orig_wh = result[i][3] * orig_hx_top_left = x - w / 2.y_top_left = y - h / 2.x_left, y_left = max(0, x_top_left), max(0, y_top_left)wi, hi = min(orig_w, w), min(orig_h, h)result[i][0], result[i][1], result[i][2], result[i][3] = x_left, y_left, wi, hireturn resultif __name__ == '__main__':img = cv2.imread('test_00000025.jpg')pred = main('yolov5_mask.ckpt', img)pred = nms(pred, 0.6, 0.4)ret_img = draw(img, pred)ret_img = ret_img[:, :, ::-1]plt.imshow(ret_img)plt.show()

基于mindspore的口罩检测训练与在线推理相关推荐

基于Opencv-python人脸口罩检测（附完整代码）
目录一.开发环境二.设计要求三.设计原理四.程序代码五.结果展示六.结论一.开发环境 python 3.6.6 opencv-python 4.5.1 二.设计要求 · 1.使用open ...
基于openMV的口罩检测
基于openMV的口罩检测什么是openmv openmv可以做什么学习过程好久没有更新了,这学期过得有点狼狈,暑假留校做一辆stm32智能小车,会不时更新一些我学习到的东西,今天分享一下我这两 ...
《Python与硬件项目案例》— 基于Python的口罩检测与指纹识别签到系统设计
<Python与硬件项目案例>- 基于Python的口罩检测与指纹识别签到系统设计目录 <Python与硬件项目案例>- 基于Python的口罩检测与指纹识别签到系统设计 1 ...
基于yolov5佩戴口罩检测项目代码
以下是一个基于YOLOv5模型进行口罩检测的Python代码示例: import cv2 import torch from yolov5.models.experimental import att ...
基于python keras口罩检测人脸检测佩戴口罩可读视频可读摄像头实时视频流
人脸识别技术已经非常普及啦,现在戴口罩的脸支付宝也可以识别,据报道阿里现在正在尝试主导人脸识别技术的某些标准.在商业上大多数公司会选择国内AI大咖,比如百度智能云.阿里智慧云.华为云.腾讯云等等.这些 ...
基于yolov5的口罩检测
1项目的克隆和必要的环境依赖项目地址:ultralytics/yolov5: YOLOv5
Yolov5：强大到你难以想象──新冠疫情下的口罩检测
初识Yolov5是看到一个视频可以检测街道上所有的行人,并实时框选出来.之后学习了CNN卷积神经网络,在完成一个项目需求时,发现卷积神经网络在切割图像方面仍然不太好用.于是我想到了之前看到的Yolov ...
yolov3模型识别不出训练图片_技术实践丨基于MindSpore框架Yolov3-darknet模型的篮球动作检测体验...
摘要:通过对篮球动作的分类训练及识别检测实例的讲解和体验,使我们了解了Yolov3模型的原理.架构等基本知识,为日后的深入学习奠定了基础. 背靠全新的设计理念,华为云推出了 MindSpore深度学习 ...
基于OpenCV训练口罩检测数据集并测试
以下内容是利用opencv自带的训练器opencv_traincascade.exe与opencv_createsamples.exe,来对口罩数据集进行训练.内容是自己操作过程中的笔记,可能会有些杂 ...

基于mindspore的口罩检测训练与在线推理

基于mindspore的口罩检测训练与在线推理相关推荐

最新文章

热门文章