掘金 人工智能 06月28日 18:08
某象图标点选,训练模型记录、踩坑、流程
index_new5.html
../../../zaker_core/zaker_tpl_static/wap/tpl_guoji1.html

 

本文分享了通过YOLOv8和CLIP模型进行图标识别与坐标获取的实践过程,并探讨了在CLIP识别效果不理想的情况下,转而使用孪生网络(预训练ResNet18+三元组损失)进行相似匹配的方案。文章详细介绍了数据集构建、模型训练、以及最终效果,并提供了对该技术方案的深入思考。

💡 **技术选型与流程**:项目采用了YOLOv8进行图标检测,CLIP进行图标识别。由于CLIP识别效果不佳,最终转向了孪生网络方案,通过相似匹配进行识别。

🖼️ **数据集构建**:文章提及了自动生成和标注数据集的方法,为模型训练提供了基础。这对于提高模型性能至关重要。

🛠️ **模型训练与优化**:详细介绍了YOLOv8的训练命令和YAML配置,以及孪生网络(ResNet18)的训练过程,包括三元组损失和早停策略。这展示了模型训练的实践细节。

✅ **最终效果与性能**:孪生网络的识别速度达到20-40ms,满足了实际应用的需求。文章还提供了最终效果图,展示了项目的成果。

目标:

输入一张包含多个图标的图片和指定的点击顺序,输出这些图标按顺序排列的精确坐标,以便进行下一步操作。

例如,给定这样一张图片:

题目要求:

训练的模型最终能够达到这样的一个效果

最终根据返回的坐标来做下一步处理:

这就是本次的目的。

开始

首先选择的技术路线为:

YOLOv8负责检测(找出图标在哪) + CLIP负责识别(认出图标是谁)

随后通过标注等方式拿到对应的数据集,划分训练集和验证集,

数据集获取

开始是打算一张图一张图的用labelimg来手动标注,但是那样实在有些枯燥,便想了个法子,自动生成以及标注,当然,这一切的核心肯定是得有原始图标在的,部分标注文件如下:

训练

yolo

划分好之后便可以开始训练了,这里使用的yolov8,直接使用这样的一行命令在对应目录执行:

yolo train model=yolov8n.pt data=char_dataset.yaml epochs=100 imgsz=640 batch=16 name=char_detector

yaml 配置如下图:

这里如果没有关于yolo训练方法的前置知识的话,可能看的有点懵,为啥要用这个yaml,那行命令执行了就能开始训练了??

简单说明一下,那行命令:

相当于别人封装设计好的一个开关,打开后会自动按照流程来执行后边的一些操作,例如读取训练集,验证集,前向传播,反向传播,验证等等,包括但不限于:

所以说,看似只是执行一行命令,实际上是触发了一整套自动化训练流程,大道至简~

这期间可能有各种的报错,最常见的就两种了,

一是提示说 yolo命令不存在,这个可以直接通过安装ultralytics来解决

pip install ultralytics

安装过后再次执行

yolo train model=yolov8n.pt data=char_dataset.yaml epochs=100 imgsz=640 batch=16 name=char_detector

就可以开始训练了,再之后就是第一次训练时,会从github上下载个对应训练命令里边指定model来作为初始权重,例如本次是model=yolov8n.pt,即指的是使用 YOLOv8 的 nano 模型作为初始权重,如果碰见到这里显示Download死活不动了,可以试着挂下科学或者直接去执行命令之后显示的链接里边下载好,然后放到执行命令的目录,即可跳过这一步骤,后续会用到的其他模型也同理。

训练结束后本目录会有个runs文件夹,并且一路点进去,模型会保存进weights文件夹内

train文件夹里边则保留了本次训练过程中各种信息,以及验证(考试试卷)的结果:

能够看到其几乎能够在任何情况都精准的框选到对应的图标,当然这些仅限于训练集里边有的,如果没有的,例如这种:

那就会超出它的认知,产生些难以捉摸的现象:

认错或者不知道这个该框不该框都是很正常的事情,这还是借助yolov8 这种比较成熟的项目的基础上。

随便拿几张要测的实际的图给它试了下,没啥问题了,就开始下一步 —— 认出图标是谁

clip

对所有可能出现的图标使用clip视觉模型来生成对应的特征向量,并保存成图标特征库文件(.pt)来供后续的识别使用

import torchimport clipfrom PIL import Imageimport osfrom tqdm import tqdmCONFIG = {    'source_icons_dir': '图标所在目录',    'output_file': 'icon_feature_library.pt'}def create_library():    print("正在加载 CLIP 模型...")    device = "cuda" if torch.cuda.is_available() else "cpu"    model, preprocess = clip.load("ViT-B/32", device=device)    print(f"模型已加载到 {device}")    icon_files = [f for f in os.listdir(CONFIG['source_icons_dir']) if f.endswith('.png')]    if not icon_files:        print(f"错误: 目录 '{CONFIG['source_icons_dir']}' 中没有找到 .png 图标文件。")        return    library = {        'ids': [],        'features': []    }    with torch.no_grad():        for filename in tqdm(icon_files):            icon_id = os.path.splitext(filename)[0]            file_path = os.path.join(CONFIG['source_icons_dir'], filename)            try:                image = Image.open(file_path).convert("RGB")                image_input = preprocess(image).unsqueeze(0).to(device)                feature = model.encode_image(image_input)                library['ids'].append(icon_id)                library['features'].append(feature)            except Exception as e:                print(f"处理文件 {filename} 时出错: {e}")    library['features'] = torch.cat(library['features'], dim=0)    torch.save(library, CONFIG['output_file'])    print(f"【正确的】特征库已成功保存到: {CONFIG['output_file']}")    print(f"共处理了 {len(library['ids'])} 个图标。")if __name__ == '__main__':    create_library()

经过以上种种步骤之后,终于到了验证阶段,按理来说,我们能够几乎完美的框选住所有的图标,又 使用了强大且也比较成熟的clip模型来为每个图标生成几乎独一无二的特征向量,正常是可以框选住且 识别出它是哪张图的对吧,可实际验证下来,竟然对每张不同的图的识别结果都是一样的,全都是相同名字。

这里真的要一头雾水了,不过,

既然YOLO的“眼睛”正常了,那问题必然出在CLIP上,毕竟本项目对于yolo的任务来说,只是单纯的框选出来并且裁切交由clip来认出来这是谁。

经过各种调试最终放弃clip了...

转孪生网络吧,说人话就是相似匹配,看他们是不是一个胚子里出的。

先是尝试自己训练自建,搭建了简单的几层CNN网络,经过一段时间的训练,终于完成,满怀期待的去预测,结果还不如clip....

最终使用预训练ResNet18,并且搭配三元组损失,增强数据集,再让其从0学习,经过200轮(29轮早停)的训练,最终:

识别速度还是可以的,基本在20-40ms之间。

在随后的请求方面,也是经历了点点小插曲,最终效果图

如果大家对项目中的数据集自动生成脚本、孪生网络训练的完整代码等细节感兴趣,可以加入我的星球获取全部资料和源码,那里还有更多,仅供学习交流哈

t.zsxq.com/GEIze

🧠 除此之外目前星球内已整理的核心干货:

✅ 某里 v2验证全流程详解

✅ 某里 v2 最新3.9 sg & pe的动态提取,加密data一步到位

✅ 顶象 ac 参数构造,AST逆向提取自动化脚本

✅ 某验 AST 解混淆插件,适配多层三元结构还原(已实战)

✅ 阿里231 递归函数解密链还原,魔改字符串还原算法

✅ 东航 refer__1036 加解密,含 req+res 解密算法

Fish AI Reader

Fish AI Reader

AI辅助创作,多种专业模板,深度分析,高质量内容生成。从观点提取到深度思考,FishAI为您提供全方位的创作支持。新版本引入自定义参数,让您的创作更加个性化和精准。

FishAI

FishAI

鱼阅,AI 时代的下一个智能信息助手,助你摆脱信息焦虑

联系邮箱 441953276@qq.com

相关标签

YOLOv8 CLIP 孪生网络 图标识别 坐标获取
相关文章