《基于YOLO V5的口罩识别》极市平台算法实训周总结

由于从本科期间开始接触计算机视觉，经常在各大平台查阅相关资料，所以很早就关注极市平台了。无意间在公 众号看到了极市平台的算法实训周打榜报名的推送，不仅可以学习视觉算法以及部署，还能赢取奖励，于是我心动了，报名参加了实训周。

第一波实训导师的直播课我参加了，主要是对平台打榜的介绍，以及实际操作的演示。对新手小白非常友好。大家有的问题导师基本也是有问必答。全程听下来，发现极市平台的软硬件资源提供的还可以，因为曾使用过类似于AutoDL，矩池云的付费在线深度学习平台。对比下来，极市平台在操作的便捷性，平台的稳定性上不输以上付费平台。虽然可选的GPU资源有限，但是打榜过程中GPU的使用基本是免费的，耗费的GPU算力都可以用打榜专属的极力值兑换。

 

​                                                                        口罩识别项目背景

由于导师安排的一些项目琐事，在起初的几天并没有参加打榜。好在考虑到假期，平台给予了时间的延长（不然我也不会在这写感想了哈哈）。言归正传，打榜不限项目，可以从目标检测，图像分割等领域去选择，我选择的是口罩识别，考虑到算法的实时性，采用经典的单阶段的YOLO V5算法。并在极市平台选择Pytorch框架。口罩识别中有20个标注类别，但是实际上检测报警输出的只有六个类别，这是一个小坑。再一个就是极市平台的编码环境和测试环境是分开的。如果在编码环境中训练并保存了模型，在进行训练前最好删除，否则在测试的时候由于ji.py（测试脚本）会默认读取模型保存路径下的第一个文件，会导致读取编码环境保存的文件，最终的得分很低。这两个bug困扰了我相当长一段时间。

 

                                                                     一个容易踩坑的点

再一个，对于模型的测试得分较低的问题。首先为了模型更快地收敛，要选择预训练模型，我先后选择YOLOV5的小预训练模型，和中训练模型，前者的FPS高，但精度低，后者则相反。并且受限于最长12小时的训练时间限制，一般10个Epoch左右训练就会终止。所以我一般选择晚上开始训练，第二天查看结果。

测试代码如下：

yolo_path = '/project/train/src_repo/v5/yolov5'
model_path = '/project/train/models/*/weights/best.pt'
names = ['front_wear', 'front_no_wear', 'front_under_nose_wear', 'front_under_mouth_wear', 'mask_front_wear', 'mask_front_under_nose_wear', 'mask_front_under_mouth_wear', 'side_wear', 'side_no_wear', 'side_under_nose_wear', 'side_under_mouth_wear', 'mask_side_wear', 'mask_side_under_nose_wear', 'mask_side_under_mouth_wear', 'side_back_head_wear', 'side_back_head_no_wear', 'back_head', 'strap', 'front_unknown', 'side_unknown']

target_names=['front_no_wear','front_under_nose_wear','front_under_mouth_wear','side_no_wear','side_under_nose_wear','side_under_mouth_wear']


conf_thres = 0.3
iou_thres = 0.05
prob_thres = 0.26
imgsz = 480

import argparse
import os
import platform
import shutil
import time
from pathlib import Path
import sys
import json
sys.path.insert(1, yolo_path)
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
import numpy as np
import argparse
import time
import cv2
import torch
import torch.backends.cudnn as cudnn
from numpy import random
from utils.augmentations import letterbox
from models.experimental import attempt_load
from utils.general import apply_classifier, check_img_size, check_imshow, check_requirements, check_suffix, colorstr, \
    increment_path, non_max_suppression, print_args, scale_coords, set_logging, \
    strip_optimizer, xyxy2xywh
from utils.plots import Annotator, colors
from utils.torch_utils import select_device, time_sync
from glob import glob


model_path=glob(model_path)[0]
device = '0'
stride = 32


def init():
    # Initialize
    global imgsz, device, stride,model_path
    set_logging()
    device = select_device(device)
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = torch.jit.load(model_path) if 'torchscript' in model_path else attempt_load(model_path, device=device)
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size
    model.eval()
    model.half()  # to FP16
    return model

@torch.no_grad()
def process_image(model, input_image=None, args=None, **kwargs):
    # Padded resize
    img0 = input_image
    img = letterbox(img0, new_shape=imgsz, stride=stride, auto=True)[0]

    # Convert
    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    img = np.ascontiguousarray(img)

    img = torch.from_numpy(img).to(device)
    img = img.half()
#     img = img.float()
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if len(img.shape) == 3:
        img = img[None] 
    pred = model(img, augment=False)[0]
    # Apply NMS
    pred = non_max_suppression(pred, conf_thres, iou_thres, agnostic=False)

    fake_result = {}
    fake_result["algorithm_data"] = {
       "is_alert": False,
       "target_count": 0,
       "target_info": []
   }
    fake_result["model_data"] = {"objects": []}
    # Process detections
    cnt = 0
    for i, det in enumerate(pred):  # detections per image
        gn = torch.tensor(img0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
        if det is not None and len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()
            for *xyxy, conf, cls in det:
                if conf < prob_thres:
                    continue
                if names[int(cls)] in target_names:
                    cnt += 1
                    fake_result["algorithm_data"]["target_info"].append({
                    "x": int(xyxy[0]),
                    "y": int(xyxy[1]),
                    "width": int(xyxy[2])-int(xyxy[0]),
                    "height": int(xyxy[3])-int(xyxy[1]),
                    "confidence":float(conf),
                    "name":names[int(cls)]
                }
                )
                    
                fake_result["model_data"]['objects'].append({
                    "x": int(xyxy[0]),
                    "y": int(xyxy[1]),
                    "width": int(xyxy[2])-int(xyxy[0]),
                    "height": int(xyxy[3])-int(xyxy[1]),
                    "confidence":float(conf),
                    "name":names[int(cls)]
                })
                
    if cnt:
        fake_result ["algorithm_data"]["is_alert"] = True
        fake_result ["algorithm_data"]["target_count"] = cnt
    return json.dumps(fake_result, indent = 4)


if __name__ == '__main__':
    from glob import glob
    # Test API
    
    image_names = glob('/home/data/*/*.jpg')
    predictor = init()
    s = 0
    for image_name in image_names:
        img = cv2.imread(image_name)
        t1 = time.time()
        res = process_image(predictor, img)
        print(res)
        t2 = time.time()
        s += t2 - t1
    print(1/(s/100))

实训周过程中也有不少难忘的事情。首先是其实有很多问题在群里问导师，大部分的实训导师都回复了，有时候问的时间也比较晚，给导师也点个赞。再一个调试封装代码，弄到晚上一点多，虽然最后还有一些问题没解决，不过也算尽力了。同时，经过此次项目，对于YOLOV5的应用和部署都进一步熟悉了，也学到了不少知识，可以说收获满满！