多标签识别：JoyTag模型的图像标注革命【开源】-CSDN博客

本文链接：https://blog.csdn.net/weixin_44292902/article/details/139100848

公共视觉模型通常会对其训练数据集进行严格过滤，这限制了这些基础模型在广泛概念上的表现，进而限制了表达自由、包容性和多样性。JoyTag通过结合Danbooru 2021数据集和一组手动标记的图像，努力提高模型对不同类型图像的泛化能力。

JoyTag项目的核心功能是自动为图像打标签。用户只需提供一张图片，JoyTag就能输出多达5000个不同标签的预测，这些预测是独立进行的，意味着每个标签的预测不受其他标签的影响。

JoyTag有着广泛的适用性，它不仅能够处理动漫和漫画风格的艺术作品，还能通过辅助数据集的补充，增强对摄影图片和其他类型内容的处理能力。这使得JoyTag在多种图像类型上都具有较好的泛化能力。JoyTag项目特别强调性别积极性和包容性，它减少了对训练数据集的过滤，从而避免了对用户表达自由的限制和潜在的歧视。这种开放性使得JoyTag能够更好地理解和标记多样化的内容。JoyTag的另一个优势是其对标签系统的创新应用。它基于Danbooru的标签体系，但通过手动标记和调整，使得模型能够更准确地适应摄影内容的需要，提高了标签的相关性和准确性。

JoyTag模型是基于ViT（Vision Transformer）架构构建的，具体采用了ViT-B/16的配置。ViT是一种有效的图像分类模型，它将图像分割成多个小块（即"patches"），然后将这些小块作为输入序列进行处理，类似于自然语言处理中的Transformer模型处理文本序列的方式。

模型的输入分辨率为448x448x3，这意味着它可以处理具有这种分辨率的彩色图像。JoyTag模型的参数量为91.5M，这是一个相对适中的参数量，既能够保证模型的复杂度，又不至于过于庞大而导致训练困难。

在输出方面，JoyTag模型进行的是多标签分类，能够为每张图片预测超过5000个不同的标签。这些标签是相互独立的，使得模型可以同时识别图片中的多个特征或对象。

JoyTag模型的训练细节包括了多个方面。它使用了LAMB优化器，这是一种结合了Adam和LAMB优势的优化算法，有助于提高训练效率。此外，模型还采用了FP16混合精度和TF32技术，这些技术可以提高计算效率并减少内存使用。

为了提高模型的泛化能力，JoyTag在训练过程中使用了多种技术，包括Trivial Augment数据增强、Focal loss损失函数，以及学习率的余弦衰减策略。这些技术有助于模型在面对不平衡数据集时保持稳定，并提高其在复杂场景下的表现。

JoyTag模型的训练是在两个阶段进行的。首先，在224x224分辨率下训练了220M个样本，然后重启并在448x448分辨率下训练了额外的440M个样本。这种分阶段的训练策略有助于模型在不同分辨率下都能学习到有效的特征表示。

以下是一个简化的示例代码，展示了如何使用JoyTag模型（代码是一个示例，实际使用时需要根据JoyTag模型的具体实现细节进行调整。代码中省略了图像预处理的具体实现，这部分需要根据模型的要求来完成）：

from PIL import Image
import torch
from pathlib import Path
import torchvision.transforms.functional as TVF

# 加载模型
model = VisionModel.load_model('/path/to/model')
model.eval()
model = model.to('cuda')  # 将模型移动到GPU

# 准备图像的函数
def prepare_image(image, target_size):
    # 省略图像预处理的代码...
    pass

# 预测函数
@torch.no_grad()
def predict(image_path):
    image = Image.open(image_path)
    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()
    
    # 根据阈值筛选标签
    top_tags = [...]  # 假设这是模型训练时得到的顶部标签列表
    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]
    
    return ', '.join(predicted_tags), scores

# 使用模型进行预测
tag_string, scores = predict('test.jpg')
print(tag_string)
for tag, score in sorted(scores.items(), key=lambda x: x[1], reverse=True):
    print(f'{tag}: {score:.3f}')

尽管JoyTag模型在多标签分类任务上取得了不错的成绩，但它仍然存在一些局限性。例如，对于一些微妙的概念，如面部表情，模型可能由于缺乏足够的数据而表现不佳。对于一些主观性强的概念，如胸部大小，模型的预测可能会有所波动，因为Danbooru数据集中的标签并不总是与指定的指南一致。

JoyTag团队正在积极地使用新训练的模型来标记更多图像，以扩展辅助数据集并提高模型在特定领域的性能。他们还鼓励用户报告模型在特定标签上的表现，以便他们能够更好地指导未来的改进工作。这种持续的优化和更新保证了JoyTag项目能够不断进步，满足用户对图像标签化处理的需求。

项目链接：https://github.com/fpgaminer/joytag