前言
本文仅作学习和交流使用, 如有侵权, 请联系我删除.
点选验证是比较基础的验证之一, 本文将以腾讯防水墙的图标点选验证为例, 一步一步教大家训练自己的模型进行图标点选识别.
两个模型主要用处为: yolov5n检测图标, Siamese孪生神经网络比较图标相似度
注: 电脑配置不好的可以尝试用Google Colab来训练
附上验证图片和目标检测后的结果
准备工作
验证图片下载
具体思路为: 逆向验证码的prehandle得到图片地址, 不断刷新获取新图片地址并且下载, 此处不再赘述
(这里下载不下载sprite都行, 如果下载的话建议把总张数增加)
下载后, 将其保存到一个文件夹中
如果你想省点事少标注一些的话, 大概150张就够了. 如果你要是想要准确度, 那么建议230-300张
数据集构建(标注)
接下来是喜闻乐见 (不情不愿)的标注环节 (标注过程会很枯燥)
安装labelimg
pip install labelimg接下来, 新建一个叫datasets的文件夹, 在其内部创建images和labels两个文件夹, 并将下载好的图片移动到images文件夹中
然后在cmd中启动labelimg
打开之后是这样的, 我们需要用Open Dir打开刚才的images文件夹, Change Save Dir打开刚才的labels文件夹, 并且点击一下save下面的那个按钮, 改为YOLO
然后在View中选择Auto save mode (自动保存)
由于我们只是为了进行目标检测, 不需要分类, 只需要标注一种类型
所以这里在右侧选中Use default label 输入你想要的标签名 (笔者使用object, 要保证没有其他不同类型的标签)
快捷键:
W: 创建一个标注框 (Create RectBox)
A: 上一张图片 (Prev Image)
D: 下一张图片 (Next Image)然后就这样慢慢标注吧, 等到全部弄完后就可以开始训练YOLOv5进行目标检测了
YOLOv5模型
准备工作
下载模型
git clone https://github.com/ultralytics/yolov5.git创建虚拟环境 (非必须, 笔者使用virtualenv), 安装依赖
pip install -r requirements.txt移动数据集的位置
数据集的位置应该移到和yolov5文件夹同级的位置, 此时的文件夹结构应大致如下
Parent (父文件夹)
├── yolov5
└── ...... (省略其他文件/文件夹)
└── train.py
└── detect.py
└── ...... (省略其他文件/文件夹)
└── datasets
└── images
└── labels训练
设置模型
找到yolov5/models/yolov5n.yaml 并在目录里复制一份相同的captcha.yaml
(如果你想用别的也可以, 但本文使用yolov5n作为示例; captcha.yaml 可以是其他的文件名, 如果你修改了, 请将后续所有涉及到captcha.yaml的改为你自己的文件)
打开captcha.yaml, 修改其中的nc为1 (nc是classes的数量)
原文件:
# Parameters
nc: 80 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.25 # layer channel multiple修改后:
# Parameters
nc: 1 # number of classes
depth_multiple: 0.33 # model depth multiple
width_multiple: 0.25 # layer channel multiple设置数据集
找到yolov5/data/coco128.yaml 直接在该文件上修改 (不要问我为什么不复制一份了, 因为我懒)
修改train和val为images 删除掉names里所有的项目, 增加
0: object修改后的文件理论上讲应该是这样的
# ...省略上方注释
path: ../datasets # dataset root dir
train: images # train images (relative to 'path') 128 images
val: images # val images (relative to 'path') 128 images
test: # test images (optional)
# Classes
names:
0: object
# Download script/URL (optional)
download: https://github.com/ultralytics/assets/releases/download/v0.0.0/coco128.zip
修改train.py
修改train.py中argparse的参数 (在parse_opt函数中)
原先参数
parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path")
parser.add_argument("--cfg", type=str, default="", help="model.yaml path")
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
# ...
parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)")将其改为
parser.add_argument("--weights", type=str, default=ROOT / "yolov5n.pt", help="initial weights path")
parser.add_argument("--cfg", type=str, default=ROOT / 'models/captcha.yaml', help="model.yaml path")
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path")
# ...
parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=672, help="train, val image size (pixels)")开始训练!
python train.py训练过程约需要30min-3hours不等, 由电脑配置决定. 显示下图则表示训练完成, 存储到了yolov5/runs/train/exp 模型文件为其中的weights/best.pt
测试
这里我们使用detect.py 进行目标检测, 导出为onnx后就可以调用了
修改detect.py 依旧是parse_opt函数
原始
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path or triton URL")
parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w")改为
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "best.pt", help="model path or triton URL")
parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)")
parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="(optional) dataset.yaml path")
parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[480, 672], help="inference size h,w") # 此为图片大小, 格式 (H, W)将要检测的图片放入data/images 然后运行detect.py 稍等一会就会有如下提示
这时候去runs/detect/exp2就能看到识别后的结果了 (这个路径会变, 以最后执行输出的路径为准)
这时就能做到前言中的效果了
导出onnx并使用
导出
接着改exports.py 还是那个位置
原始
# ......
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)")
parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)")
# ......
parser.add_argument(
"--include",
nargs="+",
default=["torchscript"],
help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle",
)改为
# ......
parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "best.pt", help="model.pt path(s)")
parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[480, 672], help="image (h, w)") # 这里根据你要识别图片的大小来设置
# ......
parser.add_argument(
"--include",
nargs="+",
default=["onnx"],
help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle",
)导出后, 在目录下会出现best.onnx, 这就是导出后的模型了 (其实你应该需要导出两个, 因为还要检测小图)
使用
from PIL import Image
import typing
import cv2
import numpy as np
import onnxruntime
# YOLOv5
class YOLOv5:
def __init__(self, OnnxPath: str, CFThresh: float, IOUThresh: float, Resize: typing.Optional[typing.Tuple[int, int]] = None) -> None:
self.OSession = onnxruntime.InferenceSession(OnnxPath)
self.InputName = [i.name for i in self.OSession.get_inputs()]
self.CFThresh = CFThresh
self.IOUThresh = IOUThresh
self.Resize = Resize
# 推理
def Inference(self, Img: Image.Image) -> np.array:
if self.Resize:
Img = Img.resize(self.Resize)
Img = cv2.cvtColor(np.asarray(Img),cv2.COLOR_RGB2BGR)
OrigImg = Img
Img = np.expand_dims(OrigImg[:,:,::-1].transpose(2,0,1).astype(dtype=np.float32) / 255.0,axis=0)
return self.OSession.run(None, {i: Img for i in self.InputName})[0]
# NMS算法
def NMS(self, Dets: list) -> list[int]:
X1 = Dets[:, 0]
Y1 = Dets[:, 1]
X2 = Dets[:, 2]
Y2 = Dets[:, 3]
Areas = (Y2 - Y1 + 1) * (X2 - X1 + 1)
Keep = []
Index = Dets[:, 4].argsort()[::-1]
while Index.size > 0:
i = Index[0]
Keep.append(i)
X11 = np.maximum(X1[i], X1[Index[1:]])
Y11 = np.maximum(Y1[i], Y1[Index[1:]])
X22 = np.minimum(X2[i], X2[Index[1:]])
Y22 = np.minimum(Y2[i], Y2[Index[1:]])
Overlaps = np.maximum(0, X22 - X11 + 1) * np.maximum(0, Y22 - Y11 + 1)
ious = Overlaps / (Areas[i] + Areas[Index[1:]] - Overlaps)
IDX = np.where(ious <= self.IOUThresh)[0]
Index = Index[IDX + 1]
return Keep
# XYWH -> XYXY
def XYWH_XYXY(self, X: list) -> np.array:
Y = np.copy(X)
Y[:, 0] = X[:, 0] - X[:, 2] / 2
Y[:, 1] = X[:, 1] - X[:, 3] / 2
Y[:, 2] = X[:, 0] + X[:, 2] / 2
Y[:, 3] = X[:, 1] + X[:, 3] / 2
return Y
# 过滤框
def FilterBox(self, OrigBox) -> np.array:
OrigBox=np.squeeze(OrigBox)
Box = OrigBox[OrigBox[..., 4] > self.CFThresh]
CLSInf = Box[..., 5:]
CLS = []
for i in range(len(CLSInf)):
CLS.append(int(np.argmax(CLSInf[i])))
ALLClasses = list(set(CLS))
Output = []
for i in range(len(ALLClasses)):
CurrCLS = ALLClasses[i]
CCB = []
for j in range(len(CLS)):
if CLS[j] == CurrCLS:
Box[j][5] = CurrCLS
CCB.append(Box[j][:6])
CCB = self.XYWH_XYXY(np.array(CCB))
for k in self.NMS(CCB):
Output.append(CCB[k])
Output = np.array(Output)
return Output
# 第二遍NMS
def NMSv2(self, Boxes, Scores) -> list[int]:
x1 = Boxes[:, 0]
y1 = Boxes[:, 1]
x2 = Boxes[:, 2]
y2 = Boxes[:, 3]
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
order = Scores.argsort()[::-1]
Keep = []
while order.size > 0:
i = order[0]
Keep.append(i)
xx1 = np.maximum(x1[i], x1[order[1:]])
yy1 = np.maximum(y1[i], y1[order[1:]])
xx2 = np.minimum(x2[i], x2[order[1:]])
yy2 = np.minimum(y2[i], y2[order[1:]])
w = np.maximum(0, xx2 - xx1 + 1)
h = np.maximum(0, yy2 - yy1 + 1)
inter = w * h
iou = inter / (areas[i] + areas[order[1:]] - inter)
inds = np.where(iou <= self.IOUThresh)[0]
order = order[inds + 1]
return Keep
# 检测
def Detect(self, Img: Image.Image) -> list:
BoxData = self.FilterBox(self.Inference(Img))
Boxes = BoxData[...,:4].astype(np.int32)
return [Boxes[i] for i in self.NMSv2(Boxes, BoxData[...,4])]
DET = YOLOv5('best.onnx', 0.5, 0.3) # 示例这是我封装的一个可以调用YOLOv5模型的类, 使用.Detect进行检测, 只返回Boxes
Simaese孪生神经网络
准备工作
下载模型
git clone https://github.com/2833844911/dianxuan准备数据集
新建一个文件夹, 这个文件夹应有如下目录结构
datasets (这个文件夹, 名字无所谓)
├── images (此文件夹存储小图标)
├── outputs (此文件夹用于存放分类好的图标, 分类完成后内部会有很多文件夹)
├── Tool.py (我制作的标注工具, 一会给出代码)依旧是下载验证图片然后使用刚才训练好的YOLOv5n模型切割检测到的图标并保存到images中 (这次我们需要的不是整张图, 而是一个个小图标) (小图标至少也得两三百张)
像这样:
其中Tool.py的内容
import tkinter as tk
import os
from PIL import Image, ImageTk
def main():
# 创建Tkinter窗口
root = tk.Tk()
root.title("显示图片")
PC = [0]
Ls = os.listdir('images')
# 加载图片
image1 = Image.open('images/' + Ls[PC[0]])
image1 = image1.resize((120, 120))
image = ImageTk.PhotoImage(image1)
# 创建Label小部件并显示图片
label = tk.Label(image=image)
label.pack()
input_label = tk.Label(root, text="Input")
input_label.pack()
entry = tk.Entry(root)
entry.pack()
e = tk.Label(root, text="")
e.pack()
# submit
def submit(*args):
nonlocal image1
P = entry.get()
entry.delete(0, tk.END)
if P:
pass
else:
return
if os.path.exists('outputs/' + P):
pass
else:
os.mkdir('outputs/' + P)
image1.save('outputs/' + P + '/' + Ls[PC[0]])
PC[0] += 1
image1 = Image.open('images/' + Ls[PC[0]])
image1 = image1.resize((120, 120))
image = ImageTk.PhotoImage(image1)
label.config(image=image)
label.image = image
e.config(text=Ls[PC[0]])
submit_button = tk.Button(root, text="提交", command=submit) # 创建按钮，点击时调用get_input函数
entry.bind("<Return>", submit)
submit_button.pack()
# 进入Tkinter事件循环
root.mainloop()
if __name__ == "__main__":
main()这个工具内部有一个输入框和一个按钮, 显示图片后你需要给同类图片给打上一样的标签, 然后按下回车或者按钮就会自动保存并接着下一张
打标签完成后, 删去dianxuan中train和val文件夹内的所有文件, 将outputs中文件夹复制到其中
训练
如果需要, 你可以修改train.py中的Resize部分, 改为你想要的大小
直接使用python执行train.py即可, 训练完成后会将模型保存到bj.pth
导出为onnx并使用
导出代码
import torch.nn as nn
import torch
import torchvision.models as models
import torchvision.transforms as transforms
VGG16 = models.vgg16(pretrained=True)
class Siamese(nn.Module):
def __init__(self, pretrained=True):
super(Siamese, self).__init__()
self.resnet = VGG16.features
self.resnet = self.resnet.eval()
self.resnet.to('cpu')
flat_shape = 512 * 3 * 3
self.fully_connect1 = torch.nn.Linear(flat_shape, 512)
self.fully_connect2 = torch.nn.Linear(512, 1)
self.sgm = nn.Sigmoid()
def forward(self, x1, x2):
x1 = self.resnet(x1)
x2 = self.resnet(x2)
x1 = torch.flatten(x1, 1)
x2 = torch.flatten(x2, 1)
x = torch.abs(x1 - x2)
x = self.fully_connect1(x)
x = self.fully_connect2(x)
x = self.sgm(x)
return x
out_onnx = 'IconCompare.onnx'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
dummy = (torch.randn(1, 3, 120, 120).to(device), torch.randn(1, 3, 120, 120).to(device)) # 注, 此处大小被改过, 改为了120x120, 如果你未修改, 请改回原大小
model = torch.load('bj.pth')
model.eval()
model = model.to(device)
torch_out = torch.onnx.export(model, dummy, out_onnx, input_names=["X", "Y"])
print("finish!")
封装好的类
# ONNX
class ONNXSiamese:
def __init__(self, ONNX: str) -> None:
self.OSession = onnxruntime.InferenceSession(ONNX)
# 比较 返回相似度
def Compare(self, Img0: Image.Image, Img1: Image.Image) -> float:
return self.OSession.run(None, {'X': np.expand_dims(np.transpose(cv2.resize(cv2.cvtColor(cv2.cvtColor(np.asarray(Img0), cv2.COLOR_RGB2BGR), cv2.COLOR_BGR2RGB), (120, 120)).astype(np.float32) / 255, (2, 0, 1)), axis=0), 'Y': np.expand_dims(np.transpose(cv2.resize(cv2.cvtColor(cv2.cvtColor(np.asarray(Img1), cv2.COLOR_RGB2BGR), cv2.COLOR_BGR2RGB), (120, 120)).astype(np.float32) / 255, (2, 0, 1)), axis=0)})[0][0][0]结语
到此这篇文章就结束了, 接下来只要将这两个模型结合一下就可以实现识别点选了.
大概思路是: 先目标检测裁剪出图形, 然后再与小图比较相似度 (我就不给具体代码了)
这篇文章好长啊,,,累死我了