基于 YOLOv5 的车牌识别算法实现与训练

5.1 图片识别效果

5.2 视频识别效果

6 部分关键代码

篇幅有限，仅展示部分代码

class Detect(nn.Module): 
    stride =None# strides computed during build
    onnx_dynamic =False# ONNX export parameter
    def __init__(self, nc=80, anchors=(), ch=(), inplace=True):# detection layer
        super().__init__() 
        self.nc = nc # number of classes 
        self.no = nc +5# number of outputs per anchor 
        self.nl =len(anchors)# number of detection layers 
        self.na =len(anchors[0])//2# number of anchors 
        self.grid =[torch.zeros(1)]* self.nl # init grid 
        self.anchor_grid =[torch.zeros(1)]* self.nl # init anchor grid 
        self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl,-1,2))# shape(nl,na,2) 
        self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na,1)for x in ch)# output conv 
        self.inplace = inplace # use in-place ops (e.g. slice assignment)
    def forward(self, x): 
        z =[]# inference output
        for i inrange(self.nl): 
            x[i]= self.m[i](x[i])# conv 
            bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85) 
            x[i]= x[i].view(bs, self.na, self.no, ny, nx).permute(0,1,3,4,2).contiguous()
            ifnot self.training:# inference
                if self.onnx_dynamic or self.grid[i].shape[2:4]!= x[i].shape[2:4]: 
                    self.grid[i], self.anchor_grid[i]= self._make_grid(nx, ny, i) 
                y = x[i].sigmoid()
                if self.inplace: 
                    y[...,0:2]=(y[...,0:2]*2-0.5+ self.grid[i])* self.stride[i]# xy 
                    y[...,2:4]=(y[...,2:4]*2)**2* self.anchor_grid[i]# wh
                else:# for YOLOv5 on AWS Inferentia https://github.com/ultralytics/yolov5/pull/2953 
                    xy =(y[...,0:2]*2-0.5+ self.grid[i])* self.stride[i]# xy 
                    wh =(y[...,2:4]*2)**2* self.anchor_grid[i]# wh 
                    y = torch.cat((xy, wh, y[...,4:]),-1) 
                z.append(y.view(bs,-1, self.no))
        return x if self.training else(torch.cat(z,1), x)
    def _make_grid(self, nx=20, ny=20, i=0): 
        d = self.anchors[i].device 
        if check_version(torch.__version__,'1.10.0'):# torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility 
            yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)], indexing='ij')
        else: 
            yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)]) 
        grid = torch.stack((xv, yv),2).expand((1, self.na, ny, nx,2)).float() 
        anchor_grid =(self.anchors[i].clone()* self.stride[i]) \ .view((1, self.na,1,1,2)).expand((1, self.na, ny, nx,2)).float()
        return grid, anchor_grid 

class Model(nn.Module):
    def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None):# model, input channels, number of classes
        super().__init__()
        ifisinstance(cfg,dict): 
            self.yaml = cfg # model dict
        else:# is *.yaml
            import yaml # for torch hub 
            self.yaml_file = Path(cfg).name 
            withopen(cfg, encoding='ascii', errors='ignore')as f: 
                self.yaml = yaml.safe_load(f)# model dict
        # Define model 
        ch = self.yaml['ch']= self.yaml.get('ch', ch)# input channels
        if nc and nc != self.yaml['nc']: 
            LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}") 
            self.yaml['nc']= nc # override yaml value
        if anchors: 
            LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}') 
            self.yaml['anchors']=round(anchors)# override yaml value 
        self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])# model, save
        list self.names =[str(i)for i inrange(self.yaml['nc'])]# default names 
        self.inplace = self.yaml.get('inplace',True)# Build strides, anchors 
        m = self.model[-1]# Detect()
        ifisinstance(m, Detect): 
            s =256# 2x min stride 
            m.inplace = self.inplace 
            m.stride = torch.tensor([s / x.shape[-2]for x in self.forward(torch.zeros(1, ch, s, s))])# forward 
            m.anchors /= m.stride.view(-1,1,1) 
            check_anchor_order(m) 
            self.stride = m.stride 
            self._initialize_biases()# only run once
        # Init weights, biases 
        initialize_weights(self) 
        self.info() 
        LOGGER.info('')
    def forward(self, x, augment=False, profile=False, visualize=False):
        if augment:return self._forward_augment(x)# augmented inference, Nonereturn self._forward_once(x, profile, visualize)# single-scale inference, train
    def_forward_augment(self, x): 
        img_size = x.shape[-2:]# height, width 
        s =[1,0.83,0.67]# scales 
        f =[None,3,None]# flips (2-ud, 3-lr) 
        y =[]# outputs
        for si, fi inzip(s, f): 
            xi = scale_img(x.flip(fi)if fi else x, si, gs=int(self.stride.max())) 
            yi = self._forward_once(xi)[0]# forward
            # cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save 
            yi = self._descale_pred(yi, fi, si, img_size) 
            y.append(yi) 
        y = self._clip_augmented(y)# clip augmented tails
        return torch.cat(y,1),None# augmented inference, train
    def_forward_once(self, x, profile=False, visualize=False): 
        y, dt =[],[]# outputs
        for m in self.model:
            if m.f !=-1:# if not from previous layer 
                x = y[m.f]ifisinstance(m.f,int)else[x if j ==-1else y[j]for j in m.f]# from earlier layers
            if profile: 
                self._profile_one_layer(m, x, dt) 
            x = m(x)# run 
            y.append(x if m.i in self.save elseNone)# save output
            if visualize: 
                feature_visualization(x, m.type, m.i, save_dir=visualize)
        return x 
    def_descale_pred(self, p, flips, scale, img_size):# de-scale predictions following augmented inference (inverse operation)
        if self.inplace: 
            p[...,:4]/= scale # de-scale
        if flips ==2: 
            p[...,1]= img_size[0]- p[...,1]# de-flip udel
        if flips ==3: 
            p[...,0]= img_size[1]- p[...,0]# de-flip lr
        else: 
            x, y, wh = p[...,0:1]/ scale, p[...,1:2]/ scale, p[...,2:4]/ scale # de-scale
            if flips ==2: 
                y = img_size[0]- y # de-flip udel
            if flips ==3: 
                x = img_size[1]- x # de-flip lr
            p = torch.cat((x, y, wh, p[...,4:]),-1)
        return p 
    def_clip_augmented(self, y):# Clip YOLOv5 augmented inference tails 
        nl = self.model[-1].nl # number of detection layers (P3-P5) 
        g =sum(4** x for x inrange(nl))# grid points 
        e =1# exclude layer count 
        i =(y[0].shape[1]// g)*sum(4** x for x inrange(e))# indices 
        y[0]= y[0][:,:-i]# large 
        i =(y[-1].shape[1]// g)*sum(4**(nl -1- x)for x inrange(e))# indices 
        y[-1]= y[-1][:, i:]# small
        return y 
    def_profile_one_layer(self, m, x, dt): 
        c =isinstance(m, Detect)# is final layer, copy input as inplace fix 
        o = thop.profile(m, inputs=(x.copy()if c else x,), verbose=False)[0]/1E9*2if thop else0# FLOPs 
        t = time_sync()
        for _ inrange(10): 
            m(x.copy()if c else x) 
            dt.append((time_sync()- t)*100)
        if m == self.model[0]: 
            LOGGER.info(f"{'time (ms)':>10s}{'GFLOPs':>10s}{'params':>10s}{'module'}") 
            LOGGER.info(f'{dt[-1]:10.2f}{o:10.2f}{m.np:10.0f}{m.type}')
        if c: 
            LOGGER.info(f"{sum(dt):10.2f}{'-':>10s}{'-':>10s} Total")
    def_initialize_biases(self, cf=None):# initialize biases into Detect(), cf is class frequency
        # https://arxiv.org/abs/1708.02002 section 3.3
        # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. 
        m = self.model[-1]# Detect() module
        for mi, s inzip(m.m, m.stride):# from b = mi.bias.view(m.na,-1)# conv.bias(255) to (3,85) 
            b.data[:,4]+= math.log(8/(640/ s)**2)# obj (8 objects per 640 image) 
            b.data[:,5:]+= math.log(0.6/(m.nc -0.999999))if cf isNoneelse torch.log(cf / cf.sum())# cls 
            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
    def_print_biases(self): 
        m = self.model[-1]# Detect() module
        for mi in m.m:# from b = mi.bias.detach().view(m.na,-1).T # conv.bias(255) to (3,85) 
            LOGGER.info(('%6g Conv2d.bias:'+'%10.3g'*6)%(mi.weight.shape[1],*b[:5].mean(1).tolist(), b[5:].mean()))
    # def _print_weights(self):#
    #     for m in self.model.modules():#
    #         if type(m) is Bottleneck:#
    #             LOGGER.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
    def fuse(self):# fuse model Conv2d() + BatchNorm2d() layers 
        LOGGER.info('Fusing layers... ') 
        for m in self.model.modules():
            ifisinstance(m,(Conv, DWConv))andhasattr(m,'bn'): 
                m.conv = fuse_conv_and_bn(m.conv, m.bn)# update conv
            delattr(m,'bn')# remove batchnorm 
            m.forward = m.forward_fuse # update forward 
        self.info()
        return self 
    defautoshape(self):# add AutoShape module 
        LOGGER.info('Adding AutoShape... ') 
        m = AutoShape(self)# wrap model 
        copy_attr(m, self, include=('yaml','nc','hyp','names','stride'), exclude=())# copy attributes
        return self 
    definfo(self, verbose=False, img_size=640):# print model information 
        model_info(self, verbose, img_size)
    def_apply(self, fn):# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers 
        self =super()._apply(fn) 
        m = self.model[-1]# Detect()
        ifisinstance(m, Detect): 
            m.stride = fn(m.stride) 
            m.grid =list(map(fn, m.grid))
        ifisinstance(m.anchor_grid,list): 
            m.anchor_grid =list(map(fn, m.anchor_grid))
        return self 
    defparse_model(d, ch):# model_dict, input_channels(3) 
        LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10}{'module':<40}{'arguments':<30}") 
        anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] 
        na =(len(anchors[0])//2)ifisinstance(anchors,list)else anchors # number of anchors 
        no = na *(nc +5)# number of outputs = anchors * (classes + 5) 
        layers, save, c2 =[],[], ch[-1]# layers, savelist, ch out
        for i,(f, n, m, args)inenumerate(d['backbone']+ d['head']):# from, number, module, args 
            m =eval(m)ifisinstance(m,str)else m # eval strings
            for j, a inenumerate(args):
                try: 
                    args[j]=eval(a)ifisinstance(a,str)else a # eval strings
                except NameError:pass 
            n = n_ =max(round(n * gd),1)if n >1else n # depth gain
            if m in[Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]: 
                c1, c2 = ch[f], args[0]
                if c2 != no:# if not output 
                    c2 = make_divisible(c2 * gw,8) 
                    args =[c1, c2,*args[1:]]
            if m in[BottleneckCSP, C3, C3TR, C3Ghost]: 
                args.insert(2, n)# number of repeats 
                n =1
            elif m is nn.BatchNorm2d: 
                args =[ch[f]]
            elif m is Concat: 
                c2 =sum(ch[x]for x in f)
            elif m is Detect: 
                args.append([ch[x]for x in f])
            ifisinstance(args[1],int):# number of anchors 
                args[1]=[list(range(args[1]*2))]*len(f)
            elif m is Contract: 
                c2 = ch[f]* args[0]**2
            elif m is Expand: 
                c2 = ch[f]// args[0]**2
            else: 
                c2 = ch[f] 
            m_ = nn.Sequential(*(m(*args)for _ inrange(n)))if n >1else m(*args)# module 
            t =str(m)[8:-2].replace('__main__.','')# module type 
            np =sum(x.numel()for x in m_.parameters())# number params 
            m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params 
            LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f}{t:<40}{str(args):<30}')# print 
            save.extend(x % i for x in([f]ifisinstance(f,int)else f)if x !=-1)# append to save
            layers.append(m_)if i ==0: 
                ch =[] 
                ch.append(c2)
        return nn.Sequential(*layers),sorted(save)