本文紧接前文:
《基于yolov5s实践国际象棋目标检测模型开发》
《yolov5s融合SPD-Conv用于提升小目标和低分辨率图像检测性能实践五子棋检测识别》
首先来看下最终效果:
在我棋类检测系统开发之——五子棋检测那篇博文写完之后就萌生了想做一下基于目标检测模型实现中国象棋检测的项目,但是后来苦于没有找到合适的数据集一直没法去做,在查资料的时候了解到了还有国际象棋,后面爬取了一部分的国际象棋数据集做了一个基于yolov5s的国际象棋检测模型,数据集本身质量不高且数据量很少所以做出来的效果也很一般。
今天专门花了几个小时来构建数据集,虽说是暂时还没有拿来真实对弈棋盘的中国象棋数据,那个我觉得可以直接去4399之类的游戏网站录屏,然后玩几把游戏,把视频切图即可,后面这个也会找时间做一下看看。今天主要是基于生成的中国象棋类的目标检测数据集来开发构建目标检测模型,一方面是早就想做中国象棋检测了,另一方面也是因为想做一点模型改造融合之类的工作,正好两件事情合并在一起做了。
这里基准模型我选择的是yolov5s系列的模型,在backbone加入了se注意力机制和swintransformer机制,形成新的C3SE模块和C3STR模块,模型完整yaml文件如下所示:
#Parameters
nc: 1 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.50 # layer channel multiple
#Anchors
anchors:
- [10,13, 16,30, 33,23] # P3/8
- [30,61, 62,45, 59,119] # P4/16
- [116,90, 156,198, 373,326] # P5/32
#Backbone
backbone:
# [from, number, module, args]
[[-1, 1, Conv, [64, 6, 2, 2]], # 0-P1/2
[-1, 1, Conv, [128, 3, 2]], # 1-P2/4
[-1, 3, C3SE, [128]], #2
[-1, 1, Conv, [256, 3, 2]], # 3-P3/8
[-1, 6, C3SE, [256]], #4
[-1, 1, Conv, [512, 3, 2]], # 5-P4/16
[-1, 9, C3SE, [512]], #6
[-1, 1, Conv, [1024, 3, 2]], # 7-P5/32
[-1, 3, C3STR, [1024]], #8
[-1, 1, SPPF, [1024, 5]], # 9
]
#Head
head:
[[-1, 1, Conv, [512, 1, 1]], #10
[-1, 1, nn.Upsample, [None, 2, 'nearest']], #11
[[-1, 6], 1, Concat, [1]], #12 cat backbone P4
[-1, 3, C3SE, [512, False]], #13
[-1, 1, Conv, [256, 1, 1]], #14
[-1, 1, nn.Upsample, [None, 2, 'nearest']], #15
[[-1, 4], 1, Concat, [1]], #16 cat backbone P3
[-1, 3, C3SE, [256, False]], #17 (P3/8-small)
[-1, 1, Conv, [256, 3, 2]], #18
[[-1, 14], 1, Concat, [1]], #19 cat head P4
[-1, 3, C3SE, [512, False]], #20 (P4/16-medium)
[-1, 1, Conv, [512, 3, 2]], #21
[[-1, 10], 1, Concat, [1]], #22 cat head P5
[-1, 3, C3SE, [1024, False]], #23 (P5/32-large)
[[17, 20, 23], 1, Detect, [nc, anchors]], #24 Detect(P3, P4, P5)
]
接下来看下对应的数据集如下所示:
YOLO格式标注数据文件如下所示:
VOC格式标注数据文件格式如下:
这里一共创建了2000张样本数据。
完成数据集构建之后就可以启动模型训练了,终端日志输出如下所示:
默认执行100epoch的计算,可以看到最终的效果还是可以的。
结果数据如下:
接下来简单看下对应的指标评估可视化,如下所示:
batch检测实例如下所示: