多标签AI图像标记模型 支持5000+标签
JoyTag是基于ViT-B/16架构的AI视觉模型,专用于图像多标签分类。采用Danbooru标记体系,支持5000多个标签,适用于手绘和摄影等多种图像类型。模型在0.4阈值下F1分数达0.578,能为每张图像生成独立标签预测。可用于diffusion模型训练等多种应用场景。
JoyTag是一个最先进的人工智能视觉模型,用于图像标记,特别注重性积极性和包容性。它使用Danbooru标记方案,但适用于广泛的图像类型,从手绘到摄影作品。
输入一张图像,它会为5000多个不同标签给出预测。这是一个多标签模型,因此每个标签的预测是相互独立的,不像单类预测视觉模型。这使得自动"标记"图像成为可能,这对于广泛的应用非常有用,包括在缺少文本对的图像上训练扩散模型。
按上述方式下载模型。使用示例:
from Models import VisionModel from PIL import Image import torch.amp.autocast_mode from pathlib import Path import torch import torchvision.transforms.functional as TVF path = '/home/.../joytag/models' # 更改为您下载模型的位置 THRESHOLD = 0.4 model = VisionModel.load_model(path) model.eval() model = model.to('cuda') with open(Path(path) / 'top_tags.txt', 'r') as f: top_tags = [line.strip() for line in f.readlines() if line.strip()] def prepare_image(image: Image.Image, target_size: int) -> torch.Tensor: # 将图像填充为正方形 image_shape = image.size max_dim = max(image_shape) pad_left = (max_dim - image_shape[0]) // 2 pad_top = (max_dim - image_shape[1]) // 2 padded_image = Image.new('RGB', (max_dim, max_dim), (255, 255, 255)) padded_image.paste(image, (pad_left, pad_top)) # 调整图像大小 if max_dim != target_size: padded_image = padded_image.resize((target_size, target_size), Image.BICUBIC) # 转换为张量 image_tensor = TVF.pil_to_tensor(padded_image) / 255.0 # 标准化 image_tensor = TVF.normalize(image_tensor, mean=[0.48145466, 0.4578275, 0.40821073], std=[0.26862954, 0.26130258, 0.27577711]) return image_tensor @torch.no_grad() def predict(image: Image.Image): image_tensor = prepare_image(image, model.image_size) batch = { 'image': image_tensor.unsqueeze(0).to('cuda'), } with torch.amp.autocast_mode.autocast('cuda', enabled=True): preds = model(batch) tag_preds = preds['tags'].sigmoid().cpu() scores = {top_tags[i]: tag_preds[0][i] for i in range(len(top_tags))} predicted_tags = [tag for tag, score in scores.items() if score > THRESHOLD] tag_string = ', '.join(predicted_tags) return tag_string, scores image = Image.open('test.jpg') tag_string, scores = predict(image) print(tag_string) for tag, score in sorted(scores.items(), key=lambda x: x[1], reverse=True): print(f'{tag}: {score:.3f}')
大多数公共视觉模型在训练数据集上进行了严格的过滤。这意味着当今的基础视觉模型在基本层面上对广泛的概念表现较弱。这限制了表达自由、包容性和多样性。它还限制了机器学习模型对我们世界的基本理解。JoyTag团队认为,人类用户应该有自由表达自己的权利,不应受到任意和反复无常的内容过滤歧视。JoyTag团队还认为,机器学习模型应该对世界有广泛、深入和包容的理解。这并不排除使用后训练对齐来减少模型中的偏见,但确实排除了使用会减少模型理解世界能力或用户表达能力的过滤或对齐。
当前的JoyTag模型是在Danbooru 2021数据集和一组手动标记的图像的组合上训练的,以扩展模型在danbooru领域之外的泛化能力。使用Danbooru数据集作为主要数据源是因为其规模(超过400万张人工标记的图像)、质量和标签的多样性。Danbooru使用的标记系统范围广泛且定义明确。然而,Danbooru数据集在内容多样性方面有限;它主要关注动漫/漫画风格的艺术。例如,数据集中只有0.3%是摄影图像。为了解决这个问题,JoyTag团队手动标记了一小部分来自互联网的图像,重点关注主要数据集中未很好表示的照片和其他内容。
最新的模型版本在整个数据集(包括照片)上实现了0.578的F1分数。在对训练或验证过程中未见过的图像进行手动测试时,该模型表现一致,证明了其良好的泛化能力。
JoyTag模型基于ViT架构,带有CNN干和GAP头。
用于训练JoyTag的脚本和流程在training
目录中共享,以防对他人有用。
目前模型主要受数据限制,如果超出当前训练方案或模型规模,就会出现过拟合。基于L/14@224的模型在过拟 合前能训练到0.51的F1分数。H/14模型也是如此,即使增加了StochDepth,也只有小幅F1提升就过拟合了。这是可以预料的,因为B/16@448的计算成本与L/14@224相似,但参数更少,因此似乎提供了一种在保持良好正则化的同时扩展模型的方法。
在我的测试中,无论是220M还是440M规模,MixUp和标签平滑都没有带来改进。Mixup提供了最强的正则化,但最终的F1分数反而低于不使用它的情况。根据扫描结果,当前的学习率、权重衰减和stochdepth率在220M范围内似乎接近最优。我们没有对448部分的训练设置进行实验,因此那里可能存在更优的学习率。
Trivial augment比RandAugment和其他数据增强方案效果更好,为防止过拟合提供了最佳解决方案。其他增强方法可能有助于进一步正则化并允许更多训练。
过拟合时,模型的F1tend趋于继续改善,而验证损失则急剧下降。
在220M扫描中,与普通BCE损失相比,Focal loss似乎没有显著差异。按标签频率加权也未能改善结果。
损失乘数对稳定训练至关重要,这可能是由于标签数量和极端分布与典型BCE任务相比导致损失过低,GradScaler无法处理。简单地将损失乘以128或256就足以稳定训练。
AdamW也有效,可以提供很小的F1改进,但LAMB结果在训练曲线上更稳定,且不需要调整预热计划。
对于B/16模型,使用CNN stem而不是传统的ViT补丁stem可以大幅提升F1。对于L/14模型,这种改进似乎较小。这方面可能需要更多研究。
虽然0.578的F1分数对于这样一个高度不平衡的多类任务来说已经很好,但还远非完美。JoyTag模型对NSFW概念的理解远超各种CLIP等基础模型,但在面部表情(笑、微笑、咧嘴笑等)等数据不足的细微概念上仍有失误。
一些概念是主观的,因此模型在这些概念上倾向 于摇摆不定,比如乳房大小。Danbooru对如何标记这些有明确的指导方针,但数据集本身并不完全符合这些指导方针。
额外数据的缺乏意味着当前模型对不常见的数据,如各种类型的服装和时尚,处理能力较弱。例如,"围裙"这个标签在辅助数据集中表现不佳,因此模型在处理它时遇到困难。
"水印"等标签仍然较弱,部分原因是缺乏优质数据,部分原因是即使在448x448分辨率下,许多水印对模型来说也太小而无法识别。
Danbooru标记系统主要针对动漫/漫画风格的艺术,因此尽管JoyTag团队认为该系统足以满足大多数摄影内容的需求,但其原始使用指南必须为摄影内容和模型需求进行"翻译"。例如,"鼻子"标签仅在Danbooru数据集中不到0.5%的情况下使用,因为它被指定为仅在鼻子"画得"比正常"更突出"时使用。这如何转化为摄影内容?JoyTag团队已尽最大努力翻译这些指南,但不平衡和冲突的标记规则使得这样的标签目前仍然存在问题。目前,JoyTag团队倾向于根据简单的"是否可见?"规则来标记照片,这被认为对模型、训练和下游任务来说更直接。"突出"的概念更适合用"*_focus"之类的标签来处理。
JoyTag团队的目标是让模型能够以同等技能处理各种内容,并专注于保持辅助数据集的多样性以减少偏见。然而,这是一场持续的战斗,永远是一项进行中的工作。例如,模型在处理"非常深的肤色"等标签时仍然存在困难。如果你发现模型在某个标签上表现不佳,请开启一个问题,让我们知道你的经历,以便我们更好地指导我们的努力。
除了上面报告的平均F1外,更详细的每个标签指标可在full-metrics.txt
中找到。所有指标都是在32,768 大小的验证集上报告的。使用PHash确保验证集不包含任何出现在训练数据中的图像。
在Validation Arena中可以找到与其他模型在未见过的图像上的比较。
这是初始版本。JoyTag团队正忙于使用新训练的模型来协助标记更多图像,以扩展辅助数据集并改善模型当前的弱点。
特别感谢SmilingWolf,感谢他们在开创性地训练基于danbooru数据集的视觉模型方面所做的工作。
字节跳动发布的AI编程神器IDE
Trae是一种自适应的集成开发环境(IDE),通过自动化和多元协作改变开发流程。利用Trae,团队能够更快速、精确地编写和部署代码,从而提高编程效率和项目交付速度。Trae具备上下文感知和代码自动完成功能,是提升开发效率的理想工具。
全能AI智能助手,随时解答生活与工作的多样问题
问小白,由元石科技研发的AI智能助手,快速准确地解答各种生活和工作问题,包括但不限于搜索、规划和社交互动,帮助用户在日常生活中提高效率,轻松管理个人事务。
实时语音翻译/同声传译工具
Transly是一个多场景的AI大语言模型驱动的同声传译、专业翻译助手,它拥有超精准的音频识别翻译能力,几乎零延迟的使用体验和支持多国语言可以让你带它走遍全球,无论你是留学生、商务人士、韩剧美剧爱好者,还是出国游玩、多国会议、跨国追星等等,都可以满足你所有需要同传的场景需求,线上线下通用,扫除语言障碍,让全世界的语言交流不再有国界。
一键生成PPT和Word,让学习生活更轻松
讯飞智文是一个利用 AI 技术的项目,能够帮助用户生成 PPT 以及各类文档。无论是商业领域的市场分析报告、年度目标制定,还是学生群 体的职业生涯规划、实习避坑指南,亦或是活动策划、旅游攻略等内容,它都能提供支持,帮助用户精准表达,轻松呈现各种信息。