diff --git a/.gitignore b/.gitignore index ce1d561..95d67c8 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,8 @@ result/ plate/ wan/ label_imgs/ +pretrained_model/ +cfg1/ # 不忽略下面指定的文件类型 !*.cpp !*.h diff --git a/README.md b/README.md index 674c036..cb46a35 100644 --- a/README.md +++ b/README.md @@ -14,9 +14,9 @@ **12 民航车牌** **测试demo:** - +以yolov7-lite-s 为例: ``` -python detect_rec_plate.py --detect_model weights/plate_detect.pt --rec_model weights/plate_rec.pth --source imgs --output result +python detect_rec_plate.py --detect_model weights/yolov7-lite-s.pt --rec_model weights/plate_rec.pth --source imgs --output result ``` 测试文件夹imgs，结果保存再 result 文件夹中 @@ -45,9 +45,9 @@ python detect_rec_plate.py --detect_model weights/plate_detect.pt --rec_model w ``` 3. 训练 - +以yolov7-lite-s 为例: ``` - python train.py --batch-size 32 --data data/plate.yaml --img 640 640 --cfg cfg/yolov7-lite-e-plate.yaml --weights weights/yolov7-lite-e.pt --name yolov7 --hyp data/hyp.face.yaml + python train.py --batch-size 32 --data data/plate.yaml --img 640 640 --cfg cfg/yolov7-lite-s.yaml --weights weights/yolov7-lite-s.pt --name yolov7 --hyp data/hyp.face.yaml ``` 结果存在run文件夹中 @@ -72,7 +72,7 @@ python detect_rec_plate.py --detect_model weights/plate_detect.pt --rec_model w ## References * [https://github.com/derronqi/yolov7-face](https://github.com/derronqi/yolov7-face) +* [https://github.com/qinggangwu/yolov7-pose_Npoint_Ncla](https://github.com/qinggangwu/yolov7-pose_Npoint_Ncla) * [https://github.com/WongKinYiu/yolov7](https://github.com/WongKinYiu/yolov7) * [https://github.com/TexasInstruments/edgeai-yolov5/tree/yolo-pose](https://github.com/TexasInstruments/edgeai-yolov5/tree/yolo-pose) -* [https://github.com/qinggangwu/yolov7-pose_Npoint_Ncla](https://github.com/qinggangwu/yolov7-pose_Npoint_Ncla) * [https://github.com/WongKinYiu/yolov7/tree/pose](https://github.com/WongKinYiu/yolov7/tree/pose) diff --git a/cfg/yolov5s-face.yaml b/cfg/yolov5s-face.yaml deleted file mode 100644 index 91500a7..0000000 --- a/cfg/yolov5s-face.yaml +++ /dev/null @@ -1,49 +0,0 @@ -# parameters -nc: 1 # number of classes -nkpt: 5 # number of keypoints -depth_multiple: 0.33 # model depth multiple -width_multiple: 0.50 # layer channel multiple -dw_conv_kpt: False - -anchors: - - [4,5, 6,8, 10,12] # P3/8 - - [15,19, 23,30, 39,52] # P4/16 - - [72,97, 123,164, 209,297] # P5/32 - -# YOLOv5 v6.0 backbone -backbone: - # [from, number, module, args] - [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 - [-1, 1, Conv, [128, 3, 2]], # 1-P2/4 - [-1, 3, C3, [128]], - [-1, 1, Conv, [256, 3, 2]], # 3-P3/8 - [-1, 6, C3, [256]], - [-1, 1, Conv, [512, 3, 2]], # 5-P4/16 - [-1, 9, C3, [512]], - [-1, 1, Conv, [1024, 3, 2]], # 7-P5/32 - [-1, 3, C3, [1024]], - [-1, 1, SPPF, [1024, 3]], # 9 - ] - -# YOLOv5 v6.0 head -head: - [[-1, 1, Conv, [512, 1, 1]], - [-1, 1, nn.Upsample, [None, 2, 'nearest']], - [[-1, 6], 1, Concat, [1]], # cat backbone P4 - [-1, 3, C3, [512, False]], # 13 - - [-1, 1, Conv, [256, 1, 1]], - [-1, 1, nn.Upsample, [None, 2, 'nearest']], - [[-1, 4], 1, Concat, [1]], # cat backbone P3 - [-1, 3, C3, [256, False]], # 17 (P3/8-small) - - [-1, 1, Conv, [256, 3, 2]], - [[-1, 14], 1, Concat, [1]], # cat head P4 - [-1, 3, C3, [512, False]], # 20 (P4/16-medium) - - [-1, 1, Conv, [512, 3, 2]], - [[-1, 10], 1, Concat, [1]], # cat head P5 - [-1, 3, C3, [1024, False]], # 23 (P5/32-large) - - [[17, 20, 23], 1, IKeypoint, [nc, anchors, nkpt]], # Detect(P3, P4, P5) - ] diff --git a/cfg/yolov7-face.yaml b/cfg/yolov7-face.yaml index 5b176a3..61b7391 100644 --- a/cfg/yolov7-face.yaml +++ b/cfg/yolov7-face.yaml @@ -1,9 +1,9 @@ # parameters -nc: 1 # number of classes -nkpt: 5 # number of keypoints +nc: 2 # number of classes +nkpt: 4 # number of keypoints depth_multiple: 1.0 # model depth multiple width_multiple: 1.0 # layer channel multiple -dw_conv_kpt: True +dw_conv_kpt: False anchors: - [4,5, 6,8, 10,12] # P3/8 @@ -73,7 +73,7 @@ backbone: # yolov7 head head: - [[-1, 1, SPPCSPC, [512]], # 51 + [[-1, 1, SPPFCSPC, [512]], # 51 [-1, 1, Conv, [256, 1, 1]], [-1, 1, nn.Upsample, [None, 2, 'nearest']], diff --git a/cfg/yolov7-lite-e-plate.yaml b/cfg/yolov7-lite-e-plate.yaml deleted file mode 100644 index 011bdcb..0000000 --- a/cfg/yolov7-lite-e-plate.yaml +++ /dev/null @@ -1,46 +0,0 @@ -# parameters -nc: 2 # number of classes -nkpt: 4 # number of keypoints -depth_multiple: 1.0 # model depth multiple -width_multiple: 1.0 # layer channel multiple -dw_conv_kpt: True - -anchors: - - [4,5, 6,8, 10,12] # P3/8 - - [15,19, 23,30, 39,52] # P4/16 - - [72,97, 123,164, 209,297] # P5/32 - -# custom backbone -backbone: - # [from, number, module, args] - [ [ -1, 1, conv_bn_relu_maxpool, [ 32 ] ], # 0-P2/4 - [ -1, 1, Shuffle_Block, [ 116, 2 ] ], # 1-P3/8 - [ -1, 3, Shuffle_Block, [ 116, 1 ] ], # 2 - [ -1, 1, Shuffle_Block, [ 232, 2 ] ], # 3-P4/16 - [ -1, 7, Shuffle_Block, [ 232, 1 ] ], # 4 - [ -1, 1, Shuffle_Block, [ 464, 2 ] ], # 5-P5/32 - [ -1, 1, Shuffle_Block, [ 464, 1 ] ], # 6 - ] - -# v5lite-e head -head: - [ [ -1, 1, Conv, [ 96, 1, 1 ] ], - [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], - [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P4 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 10 - - [ -1, 1, Conv, [ 96, 1, 1 ] ], - [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], - [ [ -1, 2 ], 1, Concat, [ 1 ] ], # cat backbone P3 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 14 (P3/8-small) - - [-1, 1, DWConvblock, [96, 3, 2]], - [ [ -1, 11 ], 1, ADD, [ 1 ] ], # cat head P4 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 17 (P4/16-medium) - - [ -1, 1, DWConvblock, [ 96, 3, 2 ] ], - [ [ -1, 7 ], 1, ADD, [ 1 ] ], # cat head P5 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 20 (P5/32-large) - - [ [ 14, 17, 20 ], 1, IKeypoint, [nc, anchors, nkpt]], # Detect(P3, P4, P5) - ] diff --git a/cfg/yolov7-lite-e.yaml b/cfg/yolov7-lite-e.yaml deleted file mode 100644 index 011bdcb..0000000 --- a/cfg/yolov7-lite-e.yaml +++ /dev/null @@ -1,46 +0,0 @@ -# parameters -nc: 2 # number of classes -nkpt: 4 # number of keypoints -depth_multiple: 1.0 # model depth multiple -width_multiple: 1.0 # layer channel multiple -dw_conv_kpt: True - -anchors: - - [4,5, 6,8, 10,12] # P3/8 - - [15,19, 23,30, 39,52] # P4/16 - - [72,97, 123,164, 209,297] # P5/32 - -# custom backbone -backbone: - # [from, number, module, args] - [ [ -1, 1, conv_bn_relu_maxpool, [ 32 ] ], # 0-P2/4 - [ -1, 1, Shuffle_Block, [ 116, 2 ] ], # 1-P3/8 - [ -1, 3, Shuffle_Block, [ 116, 1 ] ], # 2 - [ -1, 1, Shuffle_Block, [ 232, 2 ] ], # 3-P4/16 - [ -1, 7, Shuffle_Block, [ 232, 1 ] ], # 4 - [ -1, 1, Shuffle_Block, [ 464, 2 ] ], # 5-P5/32 - [ -1, 1, Shuffle_Block, [ 464, 1 ] ], # 6 - ] - -# v5lite-e head -head: - [ [ -1, 1, Conv, [ 96, 1, 1 ] ], - [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], - [ [ -1, 4 ], 1, Concat, [ 1 ] ], # cat backbone P4 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 10 - - [ -1, 1, Conv, [ 96, 1, 1 ] ], - [ -1, 1, nn.Upsample, [ None, 2, 'nearest' ] ], - [ [ -1, 2 ], 1, Concat, [ 1 ] ], # cat backbone P3 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 14 (P3/8-small) - - [-1, 1, DWConvblock, [96, 3, 2]], - [ [ -1, 11 ], 1, ADD, [ 1 ] ], # cat head P4 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 17 (P4/16-medium) - - [ -1, 1, DWConvblock, [ 96, 3, 2 ] ], - [ [ -1, 7 ], 1, ADD, [ 1 ] ], # cat head P5 - [ -1, 1, DWConvblock, [96, 3, 1] ], # 20 (P5/32-large) - - [ [ 14, 17, 20 ], 1, IKeypoint, [nc, anchors, nkpt]], # Detect(P3, P4, P5) - ] diff --git a/cfg/yolov7-lite-s.yaml b/cfg/yolov7-lite-s.yaml new file mode 100644 index 0000000..dc01a41 --- /dev/null +++ b/cfg/yolov7-lite-s.yaml @@ -0,0 +1,47 @@ +# parameters +nc: 2 # number of classes +nkpt: 4 # number of keypoints +depth_multiple: 1.0 # model depth multiple +width_multiple: 1.0 # layer channel multiple +dw_conv_kpt: False + +anchors: + - [4,5, 6,8, 10,12] # P3/8 + - [15,19, 23,30, 39,52] # P4/16 + - [72,97, 123,164, 209,297] # P5/32 + +# custom backbone +backbone: + # [from, number, module, args] + [ [ -1, 1, StemBlock, [32, 3, 2] ], # 0-P2/4 + [ -1, 1, Shuffle_Block, [96, 2]], # 1-P3/8 + [ -1, 3, Shuffle_Block, [96, 1]], # 2 + [ -1, 1, Shuffle_Block, [192, 2]], # 3-P4/16 + [ -1, 7, Shuffle_Block, [192, 1]], # 4 + [ -1, 1, Shuffle_Block, [384, 2]], # 5-P5/32 + [ -1, 3, Shuffle_Block, [384, 1]], # 6 + [ -1, 1, SPPF, [384, 5]], + ] + +# v5lite-e head +head: + [ [ -1, 1, Conv, [96, 1, 1]], + [ -1, 1, nn.Upsample, [ None, 2, 'nearest']], + [[ -1, 4], 1, Concat, [1]], # cat backbone P4 + [ -1, 1, DWConvblock, [96, 3, 1]], # 11 + + [ -1, 1, Conv, [96, 1, 1]], + [ -1, 1, nn.Upsample, [ None, 2, 'nearest']], + [[ -1, 2], 1, Concat, [1]], # cat backbone P3 + [ -1, 1, DWConvblock, [96, 3, 1] ], # 15 (P3/8-small) + + [-1, 1, DWConvblock, [96, 3, 2]], + [[ -1, 12], 1, ADD, [1]], # cat head P4 + [ -1, 1, DWConvblock, [96, 3, 1]], # 18 (P4/16-medium) + + [ -1, 1, DWConvblock, [96, 3, 2]], + [[ -1, 8], 1, ADD, [1]], # cat head P5 + [ -1, 1, DWConvblock, [96, 3, 1]], # 21 (P5/32-large) + + [[ 15, 18, 21], 1, IKeypoint, [nc, anchors, nkpt]], # Detect(P3, P4, P5) + ] diff --git a/cfg/yolov7-lite-t.yaml b/cfg/yolov7-lite-t.yaml new file mode 100644 index 0000000..accdf2a --- /dev/null +++ b/cfg/yolov7-lite-t.yaml @@ -0,0 +1,47 @@ +# parameters +nc: 2 # number of classes +nkpt: 4 # number of keypoints +depth_multiple: 1.0 # model depth multiple +width_multiple: 1.0 # layer channel multiple +dw_conv_kpt: False + +anchors: + - [4,5, 6,8, 10,12] # P3/8 + - [15,19, 23,30, 39,52] # P4/16 + - [72,97, 123,164, 209,297] # P5/32 + +# custom backbone +backbone: + # [from, number, module, args] + [ [ -1, 1, StemBlock, [16, 3, 2] ], # 0-P2/4 + [ -1, 1, Shuffle_Block, [48, 2]], # 1-P3/8 + [ -1, 2, Shuffle_Block, [48, 1]], # 2 + [ -1, 1, Shuffle_Block, [96, 2]], # 3-P4/16 + [ -1, 5, Shuffle_Block, [96, 1]], # 4 + [ -1, 1, Shuffle_Block, [192, 2]], # 5-P5/32 + [ -1, 2, Shuffle_Block, [192, 1]], # 6 + [ -1, 1, SPPF, [192, 5]], + ] + +# v5lite-e head +head: + [ [ -1, 1, Conv, [48, 1, 1]], + [ -1, 1, nn.Upsample, [ None, 2, 'nearest']], + [[ -1, 4], 1, Concat, [1]], # cat backbone P4 + [ -1, 1, DWConvblock, [48, 3, 1]], # 11 + + [ -1, 1, Conv, [48, 1, 1]], + [ -1, 1, nn.Upsample, [ None, 2, 'nearest']], + [[ -1, 2], 1, Concat, [1]], # cat backbone P3 + [ -1, 1, DWConvblock, [48, 3, 1] ], # 15 (P3/8-small) + + [-1, 1, DWConvblock, [48, 3, 2]], + [[ -1, 12], 1, ADD, [1]], # cat head P4 + [ -1, 1, DWConvblock, [48, 3, 1]], # 18 (P4/16-medium) + + [ -1, 1, DWConvblock, [48, 3, 2]], + [[ -1, 8], 1, ADD, [1]], # cat head P5 + [ -1, 1, DWConvblock, [48, 3, 1]], # 21 (P5/32-large) + + [[ 15, 18, 21], 1, IKeypoint, [nc, anchors, nkpt]], # Detect(P3, P4, P5) + ] diff --git a/cfg/yolov7-tiny-face.yaml b/cfg/yolov7-tiny-face.yaml index a4a1c8d..5ba29b9 100644 --- a/cfg/yolov7-tiny-face.yaml +++ b/cfg/yolov7-tiny-face.yaml @@ -1,6 +1,6 @@ # parameters -nc: 1 # number of classes -nkpt: 5 # number of keypoints +nc: 2 # number of classes +nkpt: 4 # number of keypoints depth_multiple: 1.0 # model depth multiple width_multiple: 1.0 # layer channel multiple dw_conv_kpt: True @@ -53,9 +53,9 @@ backbone: head: [[-1, 1, Conv, [256, 1, 1, None, 1, nn.LeakyReLU(0.1)]], [-2, 1, Conv, [256, 1, 1, None, 1, nn.LeakyReLU(0.1)]], - [-1, 1, SP, [5]], - [-2, 1, SP, [9]], - [-3, 1, SP, [13]], + [-1, 1, SPF, [5]], + [-2, 1, SPF, [9]], + [-3, 1, SPF, [13]], [[-1, -2, -3, -4], 1, Concat, [1]], [-1, 1, Conv, [256, 1, 1, None, 1, nn.LeakyReLU(0.1)]], [[-1, -7], 1, Concat, [1]], diff --git a/cfg/yolov7s-face.yaml b/cfg/yolov7s-face.yaml new file mode 100644 index 0000000..c648c31 --- /dev/null +++ b/cfg/yolov7s-face.yaml @@ -0,0 +1,139 @@ +# parameters +nc: 2 # number of classes +nkpt: 4 # number of keypoints +depth_multiple: 1.0 # model depth multiple +width_multiple: 0.4 # layer channel multiple +dw_conv_kpt: False + +anchors: + - [4,5, 6,8, 10,12] # P3/8 + - [15,19, 23,30, 39,52] # P4/16 + - [72,97, 123,164, 209,297] # P5/32 + +# yolov7 backbone +backbone: + # [from, number, module, args] + [[-1, 1, Conv, [64, 3, 2]], # 0-P1/2 + [-1, 1, Conv, [64, 3, 1]], + + [-1, 1, Conv, [128, 3, 2]], # 2-P2/4 + [-1, 1, Conv, [64, 1, 1]], + [-2, 1, Conv, [64, 1, 1]], + [-1, 1, Conv, [64, 3, 1]], + [-1, 1, Conv, [64, 3, 1]], + [-1, 1, Conv, [64, 3, 1]], + [-1, 1, Conv, [64, 3, 1]], + [[-1, -3, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 10 + + [-1, 1, MP, []], + [-1, 1, Conv, [128, 1, 1]], + [-3, 1, Conv, [128, 1, 1]], + [-1, 1, Conv, [128, 3, 2]], + [[-1, -3], 1, Concat, [1]], # 15-P3/8 + [-1, 1, Conv, [128, 1, 1]], + [-2, 1, Conv, [128, 1, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [[-1, -3, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [512, 1, 1]], # 23 + + [-1, 1, MP, []], + [-1, 1, Conv, [256, 1, 1]], + [-3, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [256, 3, 2]], + [[-1, -3], 1, Concat, [1]], # 28-P4/16 + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [[-1, -3, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [1024, 1, 1]], # 36 + + [-1, 1, MP, []], + [-1, 1, Conv, [512, 1, 1]], + [-3, 1, Conv, [512, 1, 1]], + [-1, 1, Conv, [512, 3, 2]], + [[-1, -3], 1, Concat, [1]], # 41-P5/32 + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [[-1, -3, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [1024, 1, 1]], # 49 + ] + +# yolov7 head +head: + [[-1, 1, SPPF, [512, 3]], # 50 + + [-1, 1, Conv, [256, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [36, 1, Conv, [256, 1, 1]], # route backbone P4 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 62 + + [-1, 1, Conv, [128, 1, 1]], + [-1, 1, nn.Upsample, [None, 2, 'nearest']], + [23, 1, Conv, [128, 1, 1]], # route backbone P3 + [[-1, -2], 1, Concat, [1]], + + [-1, 1, Conv, [128, 1, 1]], + [-2, 1, Conv, [128, 1, 1]], + [-1, 1, Conv, [64, 3, 1]], + [-1, 1, Conv, [64, 3, 1]], + [-1, 1, Conv, [64, 3, 1]], + [-1, 1, Conv, [64, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [128, 1, 1]], # 74 + + [-1, 1, MP, []], + [-1, 1, Conv, [128, 1, 1]], + [-3, 1, Conv, [128, 1, 1]], + [-1, 1, Conv, [128, 3, 2]], + [[-1, -3, 62], 1, Concat, [1]], + + [-1, 1, Conv, [256, 1, 1]], + [-2, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [-1, 1, Conv, [128, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [256, 1, 1]], # 87 + + [-1, 1, MP, []], + [-1, 1, Conv, [256, 1, 1]], + [-3, 1, Conv, [256, 1, 1]], + [-1, 1, Conv, [256, 3, 2]], + [[-1, -3, 50], 1, Concat, [1]], + + [-1, 1, Conv, [512, 1, 1]], + [-2, 1, Conv, [512, 1, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [-1, 1, Conv, [256, 3, 1]], + [[-1, -2, -3, -4, -5, -6], 1, Concat, [1]], + [-1, 1, Conv, [512, 1, 1]], # 100 + + [74, 1, Conv, [256, 1, 1]], + [87, 1, Conv, [512, 1, 1]], + [100, 1, Conv, [1024,1, 1]], + + [[101, 102, 103], 1, IKeypoint, [nc, anchors, nkpt]], # Detect(P3, P4, P5) + ] diff --git a/data/plate.yaml b/data/plate.yaml index 47b9b6e..59946f6 100644 --- a/data/plate.yaml +++ b/data/plate.yaml @@ -9,7 +9,7 @@ # train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/] train: /mnt/Gpan/Mydata/pytorchPorject/datasets/ccpd/train_detect -val: /mnt/Gpan/Mydata/pytorchPorject/datasets/ccpd/val_detect +val: /mnt/Gpan/Mydata/pytorchPorject/datasets/ccpd/val_detect/ #val: /ssd_1t/derron/yolov5-face/data/widerface/train/ # 4952 images # number of classes diff --git a/detect.py b/detect.py index 43f7592..c366c39 100644 --- a/detect.py +++ b/detect.py @@ -166,7 +166,7 @@ def detect(opt): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') + parser.add_argument('--weights', nargs='+', type=str, default='yolov7-lite-s.pt', help='model.pt path(s)') parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam # parser.add_argument('--img-size', nargs= '+', type=int, default=640, help='inference size (pixels)') parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)') @@ -190,7 +190,7 @@ if __name__ == '__main__': parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)') parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels') parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences') - parser.add_argument('--kpt-label', type=int, default=4, help='number of keypoints') + parser.add_argument('--kpt-label', type=int, default=5, help='number of keypoints') opt = parser.parse_args() print(opt) check_requirements(exclude=('tensorboard', 'pycocotools', 'thop')) diff --git a/detect_rec_plate.py b/detect_rec_plate.py index f995da7..3e009cd 100644 --- a/detect_rec_plate.py +++ b/detect_rec_plate.py @@ -66,10 +66,16 @@ def get_plate_rec_landmark(img, xyxy, conf, landmarks, class_num,device,plate_re class_label= int(class_num) #车牌的的类型0代表单牌，1代表双层车牌 roi_img = four_point_transform(img,landmarks_np) #透视变换得到车牌小图 - if class_label: #判断是否是双层车牌，是双牌的话进行分割后然后拼接 + # cv2.imwrite("roi.jpg",roi_img) + # roi_img_h = roi_img.shape[0] + # roi_img_w = roi_img.shape[1] + # if roi_img_w/roi_img_h<3: + # class_label= + # h_w_r = roi_img_w/roi_img_h + if class_label : #判断是否是双层车牌，是双牌的话进行分割后然后拼接 roi_img=get_split_merge(roi_img) plate_number = get_plate_result(roi_img,device,plate_rec_model) #对车牌小图进行识别 - # cv2.imwrite("roi.jpg",roi_img) + result_dict['rect']=rect result_dict['landmarks']=landmarks_np.tolist() result_dict['plate_no']=plate_number @@ -92,7 +98,7 @@ def detect_Recognition_plate(model, orgimg, device,plate_rec_model,img_size): if img.ndimension() == 3: img = img.unsqueeze(0) pred = model(img)[0] - pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, kpt_label=4) + pred = non_max_suppression(pred, conf_thres=conf_thres, iou_thres=iou_thres, kpt_label=4,agnostic=True) for i, det in enumerate(pred): if len(det): # Rescale boxes from img_size to im0 size @@ -124,10 +130,10 @@ def draw_result(orgimg,dict_list): height_area = result['roi_height'] landmarks=result['landmarks'] - result = result['plate_no']+" "+"{:.2f}".format(result['score']) + result = result['plate_no'] result_str+=result+" " cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),(0,0,255),2) #画框 - if len(result)>=7: + if len(result)>6: for i in range(4): #关键点 cv2.circle(orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1) @@ -138,7 +144,7 @@ def draw_result(orgimg,dict_list): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--detect_model', nargs='+', type=str, default='runs/train/yolov75/weights/best.pt', help='model.pt path(s)') + parser.add_argument('--detect_model', nargs='+', type=str, default='weights/yolov7-lite-s.pt', help='model.pt path(s)') parser.add_argument('--rec_model', type=str, default='weights/plate_rec.pth', help='model.pt path(s)') parser.add_argument('--source', type=str, default='wan', help='source') # file/folder, 0 for webcam # parser.add_argument('--img-size', nargs= '+', type=int, default=640, help='inference size (pixels)') diff --git a/get_small_pic.py b/get_small_pic.py index 86fc40a..4de0c39 100644 --- a/get_small_pic.py +++ b/get_small_pic.py @@ -1,4 +1,4 @@ -from detect_rec_plate import detect_Recognition_plate,attempt_load,init_model,allFilePath,cv_imread,draw_result,four_point_transform,order_points +from detect_rec_plate_macao import detect_Recognition_plate,attempt_load,init_model,allFilePath,cv_imread,draw_result,four_point_transform,order_points import os import argparse import torch @@ -23,13 +23,13 @@ def is_car_number(pattern, string): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--detect_model', nargs='+', type=str, default=r'runs/train/yolov75/weights/best.pt', help='model.pt path(s)') + parser.add_argument('--detect_model', nargs='+', type=str, default=r'runs/train/yolov711/weights/best.pt', help='model.pt path(s)') parser.add_argument('--rec_model', type=str, default=r'weights/plate_rec.pth', help='model.pt path(s)') - parser.add_argument('--source', type=str, default=r'/mnt/Gpan/BaiduNetdiskDownload/VehicleColour/VehicleColour/', help='source') # file/folder, 0 for webcam + parser.add_argument('--source', type=str, default=r'/mnt/Gpan/Mydata/pytorchPorject/datasets/macao_plate/download/', help='source') # file/folder, 0 for webcam # parser.add_argument('--source', type=str, default=r'test2', help='source') # parser.add_argument('--img-size', nargs= '+', type=int, default=640, help='inference size (pixels)') parser.add_argument('--img_size', type=int, default=640, help='inference size (pixels)') - parser.add_argument('--output', type=str, default=r'/mnt/Gpan/BaiduNetdiskDownload/VehicleColour/result1', help='source') + parser.add_argument('--output', type=str, default=r'/mnt/Gpan/Mydata/pytorchPorject/datasets/macao_plate/result/', help='source') parser.add_argument('--kpt-label', type=int, default=4, help='number of keypoints') device =torch.device("cuda" if torch.cuda.is_available() else "cpu") # device = torch.device("cpu") @@ -74,7 +74,8 @@ if __name__ == '__main__': landmarks_np = np.array(landmarks).reshape(-1,2) img_roi = four_point_transform(img,landmarks_np) plate_no= result_['plate_no'] - if len(plate_no)<6 or not is_car_number(pattern_str,plate_no): + # if len(plate_no)<6 or not is_car_number(pattern_str,plate_no): + if len(plate_no)<6: continue height = result_['roi_height'] # if height<48: diff --git a/imgs/Quicker_20220930_180856.png b/imgs/Quicker_20220930_180856.png deleted file mode 100644 index eed5017..0000000 Binary files a/imgs/Quicker_20220930_180856.png and /dev/null differ diff --git a/imgs/Quicker_20220930_180919.png b/imgs/Quicker_20220930_180919.png deleted file mode 100644 index 42abc0f..0000000 Binary files a/imgs/Quicker_20220930_180919.png and /dev/null differ diff --git a/imgs/Quicker_20220930_180938.png b/imgs/Quicker_20220930_180938.png deleted file mode 100644 index 1e34d8c..0000000 Binary files a/imgs/Quicker_20220930_180938.png and /dev/null differ diff --git a/imgs/Quicker_20220930_181044.png b/imgs/Quicker_20220930_181044.png deleted file mode 100644 index cde51c1..0000000 Binary files a/imgs/Quicker_20220930_181044.png and /dev/null differ diff --git a/imgs/minghang.jpg b/imgs/minghang.jpg deleted file mode 100644 index ebdaaf0..0000000 Binary files a/imgs/minghang.jpg and /dev/null differ diff --git a/imgs/minghang.png b/imgs/minghang.png new file mode 100644 index 0000000..048c739 Binary files /dev/null and b/imgs/minghang.png differ diff --git a/models/common.py b/models/common.py index 67126bf..d2b6766 100644 --- a/models/common.py +++ b/models/common.py @@ -42,6 +42,15 @@ class SP(nn.Module): def forward(self, x): return self.m(x) +class SPF(nn.Module): + def __init__(self, k=3, s=1): + super(SPF, self).__init__() + self.n = (k - 1) // 2 + self.m = nn.Sequential(*[nn.MaxPool2d(kernel_size=3, stride=s, padding=1) for _ in range(self.n)]) + + def forward(self, x): + return self.m(x) + class ImplicitA(nn.Module): def __init__(self, channel): @@ -302,6 +311,27 @@ class SPPCSPC(nn.Module): y2 = self.cv2(x) return self.cv7(torch.cat((y1, y2), dim=1)) +class SPPFCSPC(nn.Module): + def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=5): + super(SPPFCSPC, self).__init__() + c_ = int(2 * c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c1, c_, 1, 1) + self.cv3 = Conv(c_, c_, 3, 1) + self.cv4 = Conv(c_, c_, 1, 1) + self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2) + self.cv5 = Conv(4 * c_, c_, 1, 1) + self.cv6 = Conv(c_, c_, 3, 1) + self.cv7 = Conv(2 * c_, c2, 1, 1) + + def forward(self, x): + x1 = self.cv4(self.cv3(self.cv1(x))) + x2 = self.m(x1) + x3 = self.m(x2) + y1 = self.cv6(self.cv5(torch.cat((x1,x2,x3, self.m(x3)),1))) + y2 = self.cv2(x) + return self.cv7(torch.cat((y1, y2), dim=1)) + class SPPF(nn.Module): # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13)) @@ -389,13 +419,30 @@ class Concat(nn.Module): return torch.cat(x, self.d) # yolov7-lite +class StemBlock(nn.Module): + def __init__(self, c1, c2, k=3, s=2, p=None, g=1, act=True): + super(StemBlock, self).__init__() + self.stem_1 = Conv(c1, c2, k, s, p, g, act) + self.stem_2a = Conv(c2, c2 // 2, 1, 1, 0) + self.stem_2b = Conv(c2 // 2, c2, 3, 2, 1) + self.stem_2p = nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True) + self.stem_3 = Conv(c2 * 2, c2, 1, 1, 0) + + def forward(self, x): + stem_1_out = self.stem_1(x) + stem_2a_out = self.stem_2a(stem_1_out) + stem_2b_out = self.stem_2b(stem_2a_out) + stem_2p_out = self.stem_2p(stem_1_out) + out = self.stem_3(torch.cat((stem_2b_out,stem_2p_out),1)) + return out + class conv_bn_relu_maxpool(nn.Module): def __init__(self, c1, c2): # ch_in, ch_out super(conv_bn_relu_maxpool, self).__init__() self.conv = nn.Sequential( nn.Conv2d(c1, c2, kernel_size=3, stride=2, padding=1, bias=False), nn.BatchNorm2d(c2), - nn.ReLU(inplace=True), + nn.SiLU(inplace=True), ) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) @@ -404,23 +451,23 @@ class conv_bn_relu_maxpool(nn.Module): class DWConvblock(nn.Module): "Depthwise conv + Pointwise conv" - def __init__(self, in_channels, out_channels, k, s): super(DWConvblock, self).__init__() self.p = k // 2 - self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=k, stride=s, padding=self.p, groups=in_channels, - bias=False) + self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=k, stride=s, padding=self.p, groups=in_channels, bias=False) self.bn1 = nn.BatchNorm2d(in_channels) + self.act1 = nn.SiLU(inplace=True) self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) + self.act2 = nn.SiLU(inplace=True) def forward(self, x): x = self.conv1(x) x = self.bn1(x) - x = F.relu(x) + x = self.act1(x) x = self.conv2(x) x = self.bn2(x) - x = F.relu(x) + x = self.act2(x) return x class ADD(nn.Module): @@ -438,14 +485,10 @@ def channel_shuffle(x, groups): channels_per_group = num_channels // groups # reshape - x = x.view(batchsize, groups, - channels_per_group, height, width) - + x = x.view(batchsize, groups, channels_per_group, height, width) x = torch.transpose(x, 1, 2).contiguous() - # flatten x = x.view(batchsize, -1, height, width) - return x class Shuffle_Block(nn.Module): @@ -465,19 +508,19 @@ class Shuffle_Block(nn.Module): nn.BatchNorm2d(inp), nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), - nn.ReLU(inplace=True), + nn.SiLU(inplace=True), ) self.branch2 = nn.Sequential( nn.Conv2d(inp if (self.stride > 1) else branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), - nn.ReLU(inplace=True), + nn.SiLU(inplace=True), self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1), nn.BatchNorm2d(branch_features), nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(branch_features), - nn.ReLU(inplace=True), + nn.SiLU(inplace=True), ) @staticmethod diff --git a/models/common_ori.py b/models/common_ori.py new file mode 100644 index 0000000..67126bf --- /dev/null +++ b/models/common_ori.py @@ -0,0 +1,696 @@ +# This file contains modules common to various models + +import math +from copy import copy +from pathlib import Path + +import numpy as np +import pandas as pd +import requests +import torch +import torch.nn as nn +from PIL import Image +from torch.cuda import amp +import torch.nn.functional as F + +from utils.datasets import letterbox +from utils.general import non_max_suppression, non_max_suppression_export, make_divisible, scale_coords, increment_path, xyxy2xywh, save_one_box +from utils.plots import colors, plot_one_box +from utils.torch_utils import time_synchronized + + +def autopad(k, p=None): # kernel, padding + # Pad to 'same' + if p is None: + p = k // 2 if isinstance(k, int) else [x // 2 for x in k] # auto-pad + return p + +class MP(nn.Module): + def __init__(self, k=2): + super(MP, self).__init__() + self.m = nn.MaxPool2d(kernel_size=k, stride=k) + + def forward(self, x): + return self.m(x) + + +class SP(nn.Module): + def __init__(self, k=3, s=1): + super(SP, self).__init__() + self.m = nn.MaxPool2d(kernel_size=k, stride=s, padding=k // 2) + + def forward(self, x): + return self.m(x) + + +class ImplicitA(nn.Module): + def __init__(self, channel): + super(ImplicitA, self).__init__() + self.channel = channel + self.implicit = nn.Parameter(torch.zeros(1, channel, 1, 1)) + nn.init.normal_(self.implicit, std=.02) + + def forward(self, x): + return self.implicit.expand_as(x) + x + + +class ImplicitM(nn.Module): + def __init__(self, channel): + super(ImplicitM, self).__init__() + self.channel = channel + self.implicit = nn.Parameter(torch.ones(1, channel, 1, 1)) + nn.init.normal_(self.implicit, mean=1., std=.02) + + def forward(self, x): + return self.implicit.expand_as(x) * x + + +class ReOrg(nn.Module): + def __init__(self): + super(ReOrg, self).__init__() + + def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) + return torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1) + + +def DWConv(c1, c2, k=1, s=1, act=True): + # Depthwise convolution + return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act) + + +class Conv(nn.Module): + # Standard convolution + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups + super(Conv, self).__init__() + self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g, bias=False) + self.bn = nn.BatchNorm2d(c2) + if act != "ReLU": + self.act = nn.SiLU() if act is True else (act if isinstance(act, nn.Module) else nn.Identity()) + else: + self.act = nn.ReLU(inplace=True) + + def forward(self, x): + return self.act(self.bn(self.conv(x))) + + def fuseforward(self, x): + return self.act(self.conv(x)) + +class TransformerLayer(nn.Module): + # Transformer layer https://arxiv.org/abs/2010.11929 (LayerNorm layers removed for better performance) + def __init__(self, c, num_heads): + super().__init__() + self.q = nn.Linear(c, c, bias=False) + self.k = nn.Linear(c, c, bias=False) + self.v = nn.Linear(c, c, bias=False) + self.ma = nn.MultiheadAttention(embed_dim=c, num_heads=num_heads) + self.fc1 = nn.Linear(c, c, bias=False) + self.fc2 = nn.Linear(c, c, bias=False) + + def forward(self, x): + x = self.ma(self.q(x), self.k(x), self.v(x))[0] + x + x = self.fc2(self.fc1(x)) + x + return x + + +class TransformerBlock(nn.Module): + # Vision Transformer https://arxiv.org/abs/2010.11929 + def __init__(self, c1, c2, num_heads, num_layers): + super().__init__() + self.conv = None + if c1 != c2: + self.conv = Conv(c1, c2) + self.linear = nn.Linear(c2, c2) # learnable position embedding + self.tr = nn.Sequential(*[TransformerLayer(c2, num_heads) for _ in range(num_layers)]) + self.c2 = c2 + + def forward(self, x): + if self.conv is not None: + x = self.conv(x) + b, _, w, h = x.shape + p = x.flatten(2) + p = p.unsqueeze(0) + p = p.transpose(0, 3) + p = p.squeeze(3) + e = self.linear(p) + x = p + e + + x = self.tr(x) + x = x.unsqueeze(3) + x = x.transpose(0, 3) + x = x.reshape(b, self.c2, w, h) + return x + + +class Bottleneck(nn.Module): + # Standard bottleneck + def __init__(self, c1, c2, shortcut=True, g=1, e=0.5, act=True): # ch_in, ch_out, shortcut, groups, expansion + super(Bottleneck, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1, act=act) + self.cv2 = Conv(c_, c2, 3, 1, g=g, act=act) + self.add = shortcut and c1 == c2 + + def forward(self, x): + return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) + + +class BottleneckCSP(nn.Module): + # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion + super(BottleneckCSP, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) + self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) + self.cv4 = Conv(2 * c_, c2, 1, 1) + self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) + self.act = nn.SiLU() + self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) + + def forward(self, x): + y1 = self.cv3(self.m(self.cv1(x))) + y2 = self.cv2(x) + return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) + + +class BottleneckCSPF(nn.Module): + # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion + super(BottleneckCSPF, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) + #self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False) + self.cv4 = Conv(2 * c_, c2, 1, 1) + self.bn = nn.BatchNorm2d(2 * c_) # applied to cat(cv2, cv3) + self.act = nn.SiLU() + self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) + + def forward(self, x): + y1 = self.m(self.cv1(x)) + y2 = self.cv2(x) + return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1)))) + + +class BottleneckCSP2(nn.Module): + # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks + def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): # ch_in, ch_out, number, shortcut, groups, expansion + super(BottleneckCSP2, self).__init__() + c_ = int(c2) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = nn.Conv2d(c_, c_, 1, 1, bias=False) + self.cv3 = Conv(2 * c_, c2, 1, 1) + self.bn = nn.BatchNorm2d(2 * c_) + self.act = nn.SiLU() + self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) + + def forward(self, x): + x1 = self.cv1(x) + y1 = self.m(x1) + y2 = self.cv2(x1) + return self.cv3(self.act(self.bn(torch.cat((y1, y2), dim=1)))) + + +class C3(nn.Module): + # CSP Bottleneck with 3 convolutions + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5, act=True): # ch_in, ch_out, number, shortcut, groups, expansion + super(C3, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1, act=act) + self.cv2 = Conv(c1, c_, 1, 1, act=act) + self.cv3 = Conv(2 * c_, c2, 1, act=act) # act=FReLU(c2) + self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0, act=act) for _ in range(n)]) + # self.m = nn.Sequential(*[CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)]) + + def forward(self, x): + return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1)) + + +class C3TR(C3): + # C3 module with TransformerBlock() + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): + super().__init__(c1, c2, n, shortcut, g, e) + c_ = int(c2 * e) + self.m = TransformerBlock(c_, c_, 4, n) + + +class SPP(nn.Module): + # Spatial pyramid pooling layer used in YOLOv3-SPP + def __init__(self, c1, c2, k=(3, 3, 3)): + print(k) + super(SPP, self).__init__() + c_ = c1 // 2 # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) + num_3x3_maxpool = [] + max_pool_module_list = [] + for pool_kernel in k: + assert (pool_kernel-3)%2==0; "Required Kernel size cannot be implemented with kernel_size of 3" + num_3x3_maxpool = 1 + (pool_kernel-3)//2 + max_pool_module_list.append(nn.Sequential(*num_3x3_maxpool*[nn.MaxPool2d(kernel_size=3, stride=1, padding=1)])) + #max_pool_module_list[-1] = nn.ModuleList(max_pool_module_list[-1]) + self.m = nn.ModuleList(max_pool_module_list) + + #self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) + + + def forward(self, x): + x = self.cv1(x) + return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) + + +class SPPCSP(nn.Module): + # CSP SPP https://github.com/WongKinYiu/CrossStagePartialNetworks + def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)): + super(SPPCSP, self).__init__() + c_ = int(2 * c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False) + self.cv3 = Conv(c_, c_, 3, 1) + self.cv4 = Conv(c_, c_, 1, 1) + self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) + self.cv5 = Conv(4 * c_, c_, 1, 1) + self.cv6 = Conv(c_, c_, 3, 1) + self.bn = nn.BatchNorm2d(2 * c_) + self.act = nn.SiLU() + self.cv7 = Conv(2 * c_, c2, 1, 1) + + def forward(self, x): + x1 = self.cv4(self.cv3(self.cv1(x))) + y1 = self.cv6(self.cv5(torch.cat([x1] + [m(x1) for m in self.m], 1))) + y2 = self.cv2(x) + return self.cv7(self.act(self.bn(torch.cat((y1, y2), dim=1)))) + + +class SPPCSPC(nn.Module): + # CSP SPP https://github.com/WongKinYiu/CrossStagePartialNetworks + def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5, k=(5, 9, 13)): + super(SPPCSPC, self).__init__() + c_ = int(2 * c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c1, c_, 1, 1) + self.cv3 = Conv(c_, c_, 3, 1) + self.cv4 = Conv(c_, c_, 1, 1) + self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) + self.cv5 = Conv(4 * c_, c_, 1, 1) + self.cv6 = Conv(c_, c_, 3, 1) + self.cv7 = Conv(2 * c_, c2, 1, 1) + + def forward(self, x): + x1 = self.cv4(self.cv3(self.cv1(x))) + y1 = self.cv6(self.cv5(torch.cat([x1] + [m(x1) for m in self.m], 1))) + y2 = self.cv2(x) + return self.cv7(torch.cat((y1, y2), dim=1)) + +class SPPF(nn.Module): + # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher + def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13)) + super().__init__() + c_ = c1 // 2 # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c_ * 4, c2, 1, 1) + self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2) + + def forward(self, x): + x = self.cv1(x) + y1 = self.m(x) + y2 = self.m(y1) + return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1)) + +class Focus(nn.Module): + # Focus wh information into c-space + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups + super(Focus, self).__init__() + self.contract = Contract(gain=2) + self.conv = Conv(c1 * 4, c2, k, s, p, g, act) + + def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) + if hasattr(self, "contract"): + x = self.contract(x) + elif hasattr(self, "conv_slice"): + x = self.conv_slice(x) + else: + x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1) + return self.conv(x) + +class ConvFocus(nn.Module): + # Focus wh information into c-space + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups + super(ConvFocus, self).__init__() + slice_kernel = 3 + slice_stride = 2 + self.conv_slice = Conv(c1, c1*4, slice_kernel, slice_stride, p, g, act) + self.conv = Conv(c1 * 4, c2, k, s, p, g, act) + + def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) + if hasattr(self, "conv_slice"): + x = self.conv_slice(x) + else: + x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1) + x = self.conv(x) + return x + + +class Contract(nn.Module): + # Contract width-height into channels, i.e. x(1,64,80,80) to x(1,256,40,40) + def __init__(self, gain=2): + super().__init__() + self.gain = gain + + def forward(self, x): + N, C, H, W = x.size() # assert (H / s == 0) and (W / s == 0), 'Indivisible gain' + s = self.gain + x = x.view(N, C, H // s, s, W // s, s) # x(1,64,40,2,40,2) + x = x.permute(0, 3, 5, 1, 2, 4).contiguous() # x(1,2,2,64,40,40) + return x.view(N, C * s * s, H // s, W // s) # x(1,256,40,40) + + +class Expand(nn.Module): + # Expand channels into width-height, i.e. x(1,64,80,80) to x(1,16,160,160) + def __init__(self, gain=2): + super().__init__() + self.gain = gain + + def forward(self, x): + N, C, H, W = x.size() # assert C / s ** 2 == 0, 'Indivisible gain' + s = self.gain + x = x.view(N, s, s, C // s ** 2, H, W) # x(1,2,2,16,80,80) + x = x.permute(0, 3, 4, 1, 5, 2).contiguous() # x(1,16,80,2,80,2) + return x.view(N, C // s ** 2, H * s, W * s) # x(1,16,160,160) + + +class Concat(nn.Module): + # Concatenate a list of tensors along dimension + def __init__(self, dimension=1): + super(Concat, self).__init__() + self.d = dimension + + def forward(self, x): + return torch.cat(x, self.d) + +# yolov7-lite +class conv_bn_relu_maxpool(nn.Module): + def __init__(self, c1, c2): # ch_in, ch_out + super(conv_bn_relu_maxpool, self).__init__() + self.conv = nn.Sequential( + nn.Conv2d(c1, c2, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(c2), + nn.ReLU(inplace=True), + ) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) + + def forward(self, x): + return self.maxpool(self.conv(x)) + +class DWConvblock(nn.Module): + "Depthwise conv + Pointwise conv" + + def __init__(self, in_channels, out_channels, k, s): + super(DWConvblock, self).__init__() + self.p = k // 2 + self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=k, stride=s, padding=self.p, groups=in_channels, + bias=False) + self.bn1 = nn.BatchNorm2d(in_channels) + self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False) + self.bn2 = nn.BatchNorm2d(out_channels) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = F.relu(x) + x = self.conv2(x) + x = self.bn2(x) + x = F.relu(x) + return x + +class ADD(nn.Module): + # Stortcut a list of tensors along dimension + def __init__(self, alpha=0.5): + super(ADD, self).__init__() + self.a = alpha + + def forward(self, x): + x1, x2 = x[0], x[1] + return torch.add(x1, x2, alpha=self.a) + +def channel_shuffle(x, groups): + batchsize, num_channels, height, width = x.data.size() + channels_per_group = num_channels // groups + + # reshape + x = x.view(batchsize, groups, + channels_per_group, height, width) + + x = torch.transpose(x, 1, 2).contiguous() + + # flatten + x = x.view(batchsize, -1, height, width) + + return x + +class Shuffle_Block(nn.Module): + def __init__(self, inp, oup, stride): + super(Shuffle_Block, self).__init__() + + if not (1 <= stride <= 3): + raise ValueError('illegal stride value') + self.stride = stride + + branch_features = oup // 2 + assert (self.stride != 1) or (inp == branch_features << 1) + + if self.stride > 1: + self.branch1 = nn.Sequential( + self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1), + nn.BatchNorm2d(inp), + nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(branch_features), + nn.ReLU(inplace=True), + ) + + self.branch2 = nn.Sequential( + nn.Conv2d(inp if (self.stride > 1) else branch_features, + branch_features, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(branch_features), + nn.ReLU(inplace=True), + self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1), + nn.BatchNorm2d(branch_features), + nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(branch_features), + nn.ReLU(inplace=True), + ) + + @staticmethod + def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False): + return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i) + + def forward(self, x): + if self.stride == 1: + x1, x2 = x.chunk(2, dim=1) + out = torch.cat((x1, self.branch2(x2)), dim=1) + else: + out = torch.cat((self.branch1(x), self.branch2(x)), dim=1) + + out = channel_shuffle(out, 2) + + return out + +# end of yolov7-lite + +class NMS(nn.Module): + # Non-Maximum Suppression (NMS) module + iou = 0.45 # IoU threshold + classes = None # (optional list) filter by class + + def __init__(self, conf=0.25, kpt_label=False): + super(NMS, self).__init__() + self.conf=conf + self.kpt_label = kpt_label + + + def forward(self, x): + return non_max_suppression(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes, kpt_label=self.kpt_label) + +class NMS_Export(nn.Module): + # Non-Maximum Suppression (NMS) module used while exporting ONNX model + iou = 0.45 # IoU threshold + classes = None # (optional list) filter by class + + def __init__(self, conf=0.001, kpt_label=False): + super(NMS_Export, self).__init__() + self.conf = conf + self.kpt_label = kpt_label + + def forward(self, x): + return non_max_suppression_export(x[0], conf_thres=self.conf, iou_thres=self.iou, classes=self.classes, kpt_label=self.kpt_label) + + + +class autoShape(nn.Module): + # input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS + conf = 0.25 # NMS confidence threshold + iou = 0.45 # NMS IoU threshold + classes = None # (optional list) filter by class + + def __init__(self, model): + super(autoShape, self).__init__() + self.model = model.eval() + + def autoshape(self): + print('autoShape already enabled, skipping... ') # model already converted to model.autoshape() + return self + + @torch.no_grad() + def forward(self, imgs, size=640, augment=False, profile=False): + # Inference from various sources. For height=640, width=1280, RGB images example inputs are: + # filename: imgs = 'data/images/zidane.jpg' + # URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg' + # OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(640,1280,3) + # PIL: = Image.open('image.jpg') # HWC x(640,1280,3) + # numpy: = np.zeros((640,1280,3)) # HWC + # torch: = torch.zeros(16,3,320,640) # BCHW (scaled to size=640, 0-1 values) + # multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images + + t = [time_synchronized()] + p = next(self.model.parameters()) # for device and type + if isinstance(imgs, torch.Tensor): # torch + with amp.autocast(enabled=p.device.type != 'cpu'): + return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference + + # Pre-process + n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images + shape0, shape1, files = [], [], [] # image and inference shapes, filenames + for i, im in enumerate(imgs): + f = f'image{i}' # filename + if isinstance(im, str): # filename or uri + im, f = np.asarray(Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im)), im + elif isinstance(im, Image.Image): # PIL Image + im, f = np.asarray(im), getattr(im, 'filename', f) or f + files.append(Path(f).with_suffix('.jpg').name) + if im.shape[0] < 5: # image in CHW + im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1) + im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input + s = im.shape[:2] # HWC + shape0.append(s) # image shape + g = (size / max(s)) # gain + shape1.append([y * g for y in s]) + imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update + shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape + x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad + x = np.stack(x, 0) if n > 1 else x[0][None] # stack + x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW + x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32 + t.append(time_synchronized()) + + with amp.autocast(enabled=p.device.type != 'cpu'): + # Inference + y = self.model(x, augment, profile)[0] # forward + t.append(time_synchronized()) + + # Post-process + y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS + for i in range(n): + scale_coords(shape1, y[i][:, :4], shape0[i]) + + t.append(time_synchronized()) + return Detections(imgs, y, files, t, self.names, x.shape) + + +class Detections: + # detections class for YOLOv5 inference results + def __init__(self, imgs, pred, files, times=None, names=None, shape=None): + super(Detections, self).__init__() + d = pred[0].device # device + gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations + self.imgs = imgs # list of images as numpy arrays + self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) + self.names = names # class names + self.files = files # image filenames + self.xyxy = pred # xyxy pixels + self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels + self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized + self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized + self.n = len(self.pred) # number of images (batch size) + self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms) + self.s = shape # inference BCHW shape + + def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')): + for i, (im, pred) in enumerate(zip(self.imgs, self.pred)): + str = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' + if pred is not None: + for c in pred[:, -1].unique(): + n = (pred[:, -1] == c).sum() # detections per class + str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string + if show or save or render or crop: + for *box, conf, cls in pred: # xyxy, confidence, class + label = f'{self.names[int(cls)]} {conf:.2f}' + if crop: + save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i]) + else: # all others + plot_one_box(box, im, label=label, color=colors(cls)) + + im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np + if pprint: + print(str.rstrip(', ')) + if show: + im.show(self.files[i]) # show + if save: + f = self.files[i] + im.save(save_dir / f) # save + print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n') + if render: + self.imgs[i] = np.asarray(im) + + def print(self): + self.display(pprint=True) # print results + print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % self.t) + + def show(self): + self.display(show=True) # show results + + def save(self, save_dir='runs/hub/exp'): + save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir + self.display(save=True, save_dir=save_dir) # save results + + def crop(self, save_dir='runs/hub/exp'): + save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/hub/exp', mkdir=True) # increment save_dir + self.display(crop=True, save_dir=save_dir) # crop results + print(f'Saved results to {save_dir}\n') + + def render(self): + self.display(render=True) # render results + return self.imgs + + def pandas(self): + # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0]) + new = copy(self) # return copy + ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns + cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns + for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]): + a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update + setattr(new, k, [pd.DataFrame(x, columns=c) for x in a]) + return new + + def tolist(self): + # return a list of Detections objects, i.e. 'for result in results.tolist():' + x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)] + for d in x: + for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: + setattr(d, k, getattr(d, k)[0]) # pop out of list + return x + + def __len__(self): + return self.n + + +class Classify(nn.Module): + # Classification head, i.e. x(b,c1,20,20) to x(b,c2) + def __init__(self, c1, c2, k=1, s=1, p=None, g=1): # ch_in, ch_out, kernel, stride, padding, groups + super(Classify, self).__init__() + self.aap = nn.AdaptiveAvgPool2d(1) # to x(b,c1,1,1) + self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p), groups=g) # to x(b,c2,1,1) + self.flat = nn.Flatten() + + def forward(self, x): + z = torch.cat([self.aap(y) for y in (x if isinstance(x, list) else [x])], 1) # cat if list + return self.flat(self.conv(z)) # flatten to x(b,c2) diff --git a/models/yolo.py b/models/yolo.py index 9ae629a..6043b72 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -531,7 +531,7 @@ def parse_model(d, ch): # model_dict, input_channels(3) n = max(round(n * gd), 1) if n > 1 else n # depth gain if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, ConvFocus, CrossConv, BottleneckCSP, - C3, C3TR, BottleneckCSPF, BottleneckCSP2, SPPCSP, SPPCSPC, SPPF, conv_bn_relu_maxpool, Shuffle_Block, DWConvblock]: + C3, C3TR, BottleneckCSPF, BottleneckCSP2, SPPCSP, SPPCSPC, SPPF, conv_bn_relu_maxpool, Shuffle_Block, DWConvblock,StemBlock]: c1, c2 = ch[f], args[0] if c2 != no: # if not output c2 = make_divisible(c2 * gw, 8) diff --git a/ncnn/ncnn_export.py b/ncnn/ncnn_export.py new file mode 100644 index 0000000..45ecf5c --- /dev/null +++ b/ncnn/ncnn_export.py @@ -0,0 +1,79 @@ +"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats + +Usage: + $ export PYTHONPATH="$PWD" && python models/export.py --weights yolov5s.pt --img 640 --batch 1 +""" + +import argparse +import sys +import time +from pathlib import Path + +sys.path.append(Path(__file__).parent.parent.absolute().__str__()) # to run '$ python *.py' files in subdirectories + +import onnx +import torch +import torch.nn as nn +from torch.utils.mobile_optimizer import optimize_for_mobile +import cv2 +import numpy as np + +import models +from models.experimental import attempt_load +from utils.activations import Hardswish, SiLU +from utils.general import colorstr, check_img_size, check_requirements, file_size, set_logging +from utils.torch_utils import select_device + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--weights', type=str, default='runs/train/yolov714/weights/best.pt', help='weights path') + parser.add_argument('--img-size', nargs='+', type=int, default=[320, 320], help='image size') # height, width + parser.add_argument('--batch-size', type=int, default=1, help='batch size') + parser.add_argument('--grid', action='store_true', help='export Detect() layer grid') + parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') + opt = parser.parse_args() + opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand + print(opt) + set_logging() + t = time.time() + + # Load PyTorch model + device = select_device(opt.device) + model = attempt_load(opt.weights, map_location=device) # load FP32 model + labels = model.names + + # Checks + gs = int(max(model.stride)) # grid size (max stride) + opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples + + # Input + img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection + + # Update model + for k, m in model.named_modules(): + m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility + if isinstance(m, models.common.Conv): # assign export-friendly activations + if isinstance(m.act, nn.Hardswish): + m.act = Hardswish() + elif isinstance(m.act, nn.SiLU): + m.act = SiLU() + # elif isinstance(m, models.yolo.Detect): + # m.forward = m.forward_export # assign forward (optional) + model.model[-1].export = True # set Detect() layer grid export + for _ in range(2): + y = model(img) # dry runs + output_names = None + + print(f'starting export with onnx {onnx.__version__}...') + f = opt.weights.replace('.pt', '.onnx') # filename + torch.onnx.export(model, img, f, verbose=False, opset_version=11, input_names=['data'], + output_names=['stride_' + str(int(x)) for x in model.stride]) + + + # Checks + model_onnx = onnx.load(f) # load onnx model + onnx.checker.check_model(model_onnx) # check onnx model + # print(onnx.helper.printable_graph(model_onnx.graph)) # print + + # Finish + print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.') diff --git a/onnx/yolov7_plate_onnx_infer.py b/onnx/yolov7_plate_onnx_infer.py index e8142e5..caa0fa2 100644 --- a/onnx/yolov7_plate_onnx_infer.py +++ b/onnx/yolov7_plate_onnx_infer.py @@ -214,7 +214,7 @@ def draw_result(orgimg,dict_list): for i in range(4): #关键点 cv2.circle(orgimg, (int(landmarks[i][0]), int(landmarks[i][1])), 5, clors[i], -1) cv2.rectangle(orgimg,(rect_area[0],rect_area[1]),(rect_area[2],rect_area[3]),(0,0,255),2) #画框 - if len(result)>=1: + if len(result)<0: orgimg=cv2ImgAddText(orgimg,result,rect_area[0]-height_area,rect_area[1]-height_area-10,(255,0,0),height_area) print(result_str) return orgimg diff --git a/plate_recognition/plateNet.py b/plate_recognition/plateNet.py index 1911efc..cc1649e 100644 --- a/plate_recognition/plateNet.py +++ b/plate_recognition/plateNet.py @@ -14,7 +14,7 @@ class myNet_ocr(nn.Module): # self.loc = nn.MaxPool2d((2, 2), (5, 1), (0, 1),ceil_mode=True) # self.loc = nn.AvgPool2d((2, 2), (5, 2), (0, 1),ceil_mode=False) self.loc = nn.MaxPool2d((5, 2), (1, 1),(0,1),ceil_mode=False) - self.newCnn=nn.Conv2d(256,num_classes,1,1) + self.newCnn=nn.Conv2d(cfg[-1],num_classes,1,1) # self.newBn=nn.BatchNorm2d(num_classes) def make_layers(self, cfg, batch_norm=False): layers = [] diff --git a/plate_recognition/plate_rec.py b/plate_recognition/plate_rec.py index 1e6b4ae..ad14654 100644 --- a/plate_recognition/plate_rec.py +++ b/plate_recognition/plate_rec.py @@ -20,7 +20,7 @@ def allFilePath(rootPath,allFIleList): else: allFilePath(os.path.join(rootPath,temp),allFIleList) device = torch.device('cuda') if torch.cuda.is_available() else torch.device("cpu") -plateName=r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航深0123456789ABCDEFGHJKLMNPQRSTUVWXYZ" +plateName=r"#京沪津渝冀晋蒙辽吉黑苏浙皖闽赣鲁豫鄂湘粤桂琼川贵云藏陕甘青宁新学警港澳挂使领民航危0123456789ABCDEFGHJKLMNPQRSTUVWXYZ险品" mean_value,std_value=(0.588,0.193) def decodePlate(preds): pre=0 @@ -64,7 +64,7 @@ def init_model(device,model_path): check_point = torch.load(model_path,map_location=device) model_state=check_point['state_dict'] cfg=check_point['cfg'] - model_path = os.sep.join([sys.path[0],model_path]) + # model_path = os.sep.join([sys.path[0],model_path]) model = myNet_ocr(num_classes=78,export=True,cfg=cfg) model.load_state_dict(model_state) diff --git a/train.py b/train.py index e637364..ade29d3 100644 --- a/train.py +++ b/train.py @@ -480,11 +480,11 @@ def train(hyp, opt, device, tb_writer=None): if __name__ == '__main__': parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default='weights/yolov7-lite-e.pt', help='initial weights path') - parser.add_argument('--cfg', type=str, default='cfg/yolov7-lite-e-plate.yaml', help='model.yaml path') + parser.add_argument('--weights', type=str, default='yolov7-tiny.pt', help='initial weights path') + parser.add_argument('--cfg', type=str, default='cfg/yolov7-tiny-face.yaml', help='model.yaml path') parser.add_argument('--data', type=str, default='data/plate.yaml', help='data.yaml path') parser.add_argument('--hyp', type=str, default='data/hyp.face.yaml', help='hyperparameters path') - parser.add_argument('--epochs', type=int, default=120) + parser.add_argument('--epochs', type=int, default=200) parser.add_argument('--batch-size', type=int, default=32, help='total batch size for all GPUs') parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') parser.add_argument('--rect', action='store_true', help='rectangular training') diff --git a/weights/plate_detect.pt b/weights/plate_detect.pt deleted file mode 100644 index e3867c1..0000000 Binary files a/weights/plate_detect.pt and /dev/null differ diff --git a/weights/plate_rec.pth b/weights/plate_rec.pth index 38104c7..5c2c7af 100644 Binary files a/weights/plate_rec.pth and b/weights/plate_rec.pth differ diff --git a/weights/yolov7-lite-e.pt b/weights/yolov7-lite-e.pt deleted file mode 100644 index 4a9eb43..0000000 Binary files a/weights/yolov7-lite-e.pt and /dev/null differ diff --git a/weights/yolov7-lite-s.pt b/weights/yolov7-lite-s.pt new file mode 100644 index 0000000..ffeeb44 Binary files /dev/null and b/weights/yolov7-lite-s.pt differ diff --git a/weights/yolov7-lite-t.pt b/weights/yolov7-lite-t.pt new file mode 100644 index 0000000..fe98975 Binary files /dev/null and b/weights/yolov7-lite-t.pt differ