Skip to content

在保持GT框检测的同时检测并输出多边形位置信息。

License

Notifications You must be signed in to change notification settings

Rhine-AI-Lab/Yolo-ArbV2

Repository files navigation

Yolo-ArbV2 在 YOLOv5 基础上进行二次开发。
保持GT框检测功能的同时,新增了额外输出信息,用于检测输出中目标多边形的信息。这样实现了基于矩形Anchor-based的多边形检测功能。

简介

Yolo-ArbV2 在完全保持YOLOv5功能情况下,实现可选多边形信息输出。通过扩展输出维度,额外输出多边形顶点坐标信息,相对box框左上角,归一化至box宽高的信息实现多边形信息的输出。
在整个多边形计算中,从数据结构,数据增强,模型结构,输出结构,坐标后处理,损失计算,NMS过滤都做了相应的调整。
目前只支持输出指定边数(edges参数)的多边形,如果需要识别不同边数多边形,可将数据集最后一个点进行重复处理。

大部分操作及使用方法可参考 YOLOv5 文档 。新增区别信息在下文会介绍。

快速开始

Install

Python>=3.6.0 is required with all requirements.txt installed including PyTorch>=1.7:

$ git clone https://github.com/HRan2004/Yolo-ArbV2
$ cd Yolo-ArbV2
$ pip install -r requirements.txt
Datasets 数据集的准备,大体结构与YOLOv5一致。目录结构如下。
datasets
├ images
│ └ img1,img2...
├ labels
│ └ txt1,txt2...
├ train.txt
└ val.txt 

images与labels文件夹中存放图片与标签,train.txt与val.txt用于分类图片作为训练集或测试集。存放图片路径,一行一个,会自动计算标签路径。


区别在于txt文件中的坐标信息。 当输出边数,即参数edges为4时,用于检测4边形,数据集需要四点坐标。
注意:坐标宽高,全都为相对图片宽高归一化后的数据,即0到1范围内

YOLOv5: class center_x center_y width height

Yolo-ArbV2: class x1 y1 x2 y2 x3 y3 x4 y4
Args

在hyp中新增了5项参数。

edges: 4 # edges of poly for net output
poly: 1.2  # obj loss gain (scale with pixels)
start_poly: 0 # epoch to start use poly loss
poly_out: 0.2 # max percent for poly out of box
poly_loss_smooth: 0.01 # smooth scope for SmoothL1Loss

edges 多边形边数,必填。用于设置模型输出多边形信息的边数,识别四边形则为4,需与数据集一致。设置为0时,模型功能等同于YOLOv5。
poly 多边形框损失。多边形损失权重。
start_poly 开始poly损失的epoch。由于poly相对于box进行归一化输出,建议box准确后再进行训练poly。
poly_out 出框量。poly可溢出box的最大比值。由于数据增强对box框的切割,实际poly位置可能溢出box。
poly_loss_smooth 损失平滑范围,后期将删除该参数。由于输出结果归一化,使用SmoothL1Loss需减少平滑区范围。建议填写0.005。

Opt

在train.py中新增了1项参数。

parser.add_argument('--val_rate', type=int, default=1)

方便于模型初期的debug调试,使用极小型数据集,大量epoch学习时,可适量屏蔽每个epoch后的val环节,用于加速。
设置为0则关闭val,除最后一个epoch,不会进行val。替代原有参数noval。 设置为1则正常,每个epoch进行测试。

理论介绍

Model

模型结构维持不变,仅改变模型输出层维度。后期可能适当增加输出层处理层数。

输出层 每个框:
YOLOv5: x,y,w,h,conf,class1,...,classF
Yolo-ArbV2(edges=N): x,y,w,h,[x1,y1,...,xN,yN],conf,class1,...,classF

每框信息量为: edges*2 + class_num + 5

Processing

对于多边形信息的后处理。poly信息相对于box位置进行归一化输出。
相对于box左上角点,故输出区间为 [-poly_out,1+poly_out]。
输出结果,进行Sigmoid处理后,(1+poly_out2)- poly_out即可得到输出结果。

Loss

损失的处理由于多边形的IOU计算复杂,效率过低。模型设计了多点距离法计算多边形损失。直接计算对应位置输出坐标与真实坐标的距离,将8个参数一同进行SmoothL1Loss损失,得出结果。
为避免点位顺序不同导致的损失异常,模型在导入数据数据增强后(可能包含旋转影响顺序),会自动将数据集处理成,从最高点开始顺时针点位排序的数据集。

新增了SmoothL1LossRange,可自定义损失的平滑范围,提高精确度。

class SmoothL1LossRange(nn.SmoothL1Loss):
    def __init__(self,smooth_range=1.0):
        super().__init__()
        self.sr = smooth_range

    def __call__(self, p, t):
        sr = self.sr
        loss = super().__call__(p/sr, t/sr)*sr
        return loss

新增PolyLoss,主要区别在于处理NaN参数,数据增强中把部分无效点位设置成NaN,此处对应屏蔽这部分涉及到的损失。

class PolyLoss(SmoothL1LossRange):
    def __init__(self,smooth_range=1.0):
        super().__init__(smooth_range)

    def __call__(self, p, t):
        nni = ~torch.isnan(t) # Not NaN index
        return super().__call__(p[nni],t[nni])   
Datasets mosaic 数据增强过程中,跟随box框,共同对poly进行处理。并在最后根据poly生成准确box框(旋转变换会导致box不准确)。

值得一提的是,修复了一项原有Yolov5中对segment生成box框时的bug。
在切割变换中,Yolov5通过在多边形上描绘2000个点,根据剩余点生成box。但是过程中遗漏了终点和起点连接的边线,这导致再切割相邻边线时,无法准确生成box框。

如上图,缺少上边线,并且右边线遭到切割时,生成的box出现了错误。

解决方法也很简单,在描边前进行补充起点在终点后方即可。
def resample_segments(segments, n=500):
    # Up-sample an (n,2) segment
    for i, s in enumerate(segments):
        s = np.concatenate((s, s[0:1, :]), axis=0)
        x = np.linspace(0, len(s) - 1, n) # Debug
        xp = np.arange(len(s))
        segments[i] = np.concatenate([np.interp(x, xp, s[:, i]) for i in range(2)]).reshape(2, -1).T  # segment xy
    return segments
Datasets & Loss 由于固定多边形的点溢出图像,处理时不同于GT框,当四边形,有一个点溢出图片时,我们会发现实际在图像内的图形为五边形,所以并不能对其做适当的处理。
我们采用屏蔽溢出的点的做法,但仍然保持他的占位,数据增强的同时检查并设置为NaN,在后期损失计算的时候,对齐进行屏蔽。

对于特殊情况,比如:四边形在同一条边线上溢出两个点,这时可以计算出图像内部的四边形框。我们对齐做了特殊处理,让模型在局部检测时也能有良好的效果。下面为一部分代码

# utils/general.py
def polygons_check(polys, max_out=0.02):
    if(polys.shape[0]>0):
        edges = polys.shape[1]
        if edges==4:
            out_top    = polys[...,1]<0-max_out
            ...
            for i in range(polys.shape[0]):
                poly = polys[i]
                out_n = np.zeros(4) # top right bottom left
                out_n[0] = out_top[i].sum()
                ...

                for j,num in enumerate(out_n):
                    # Filter
                    ...
                    for j in range(len(out)):
                        if not out[j]:
                            continue
                        fp = poly[j]
                        if not out[pj(j+1)]:
                            sp = poly[pj(j+1)]
                        else:
                            sp = poly[pj(j-1)]
                        np.seterr(divide='ignore')
                        if out_n[0] == 2:
                            tp = [0,0]
                            tp[0] = sp[0]+(fp[0]-sp[0])*(tp[1]-sp[1])/(fp[1]-sp[1])
                        elif out_n[1] == 2:
                            tp = [1,0]
                            tp[1] = sp[1]+(fp[1]-sp[1])*(tp[0]-sp[0])/(fp[0]-sp[0])
                        elif out_n[2] == 2:
                            tp = [0,1]
                            tp[0] = sp[0]+(fp[0]-sp[0])*(tp[1]-sp[1])/(fp[1]-sp[1])
                        else:
                            tp = [0,0]
                            tp[1] = sp[1]+(fp[1]-sp[1])*(tp[0]-sp[0])/(fp[0]-sp[0])
                        polys[i,j]=tp

                    # Fresh out side
                    ...

        nan = float('nan')
        polys[polys[...,0]<-max_out] = [nan,nan]
        polys[polys[...,1]<-max_out] = [nan,nan]
        polys[polys[...,0]>1+max_out] = [nan,nan]
        polys[polys[...,1]>1+max_out] = [nan,nan]
    return polys

为合理计算损失,在数据增强时,会有旋转翻折等情况,数据点位顺序可能错乱,所以在数据增强完毕后,要进行顺序处理。
最终转换为由最高点为起点,顺时针旋转的多边形数据。这使得直接计算对应位置点位距离的损失方法可行了。
以下为比较简洁的批量实现方式

# Make polygons start with the highest point
# and order with clock wise
# Input shape : [polygons_num,edges,2(x,y)]
# Output shape : ==Input shape
def polygons_cw(polys):
    ps = polys.shape
    hpi = np.argmax(polys[..., 1::2], axis=1).reshape(polys.shape[0]) # Highest point index
    lpi,rpi = hpi-1,hpi+1
    np.place(lpi,lpi==-1,ps[1]-1)
    np.place(rpi,rpi==ps[1],0)
    lrpi = np.concatenate(([lpi], [rpi]),0).T.flatten() # Left right point by highest point
    p1 = polys.reshape((-1,2))[hpi+np.arange(ps[0])*ps[1]].repeat(2,axis=0).reshape(-1,2,2)
    p2 = polys.reshape((-1,2))[lrpi+np.arange(ps[0]).repeat(2)*ps[1]].reshape((-1,2,2))
    pc = p2-p1
    d = 90-np.arctan(pc[...,0]/pc[...,1])*180/math.pi
    isCw = d[...,0]<d[...,1] # Is clock wise
    polys_cw = np.zeros_like(polys)
    for i in range(polys.shape[0]):
        hpii = hpi[i]
        if isCw[i]:
            polys_cw[i,:ps[1]-hpii] = polys[i,hpii:]
            polys_cw[i,ps[1]-hpii:] = polys[i,:hpii]
        else:
            polys_cw[i,:hpii+1] = polys[i,hpii::-1]
            polys_cw[i,hpii+1:] = polys[i,:hpii-ps[1]:-1]
    return polys_cw

About

在保持GT框检测的同时检测并输出多边形位置信息。

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages