????????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的口罩检测训练与在线推理的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。