标签:

• 图像分类
• PyTorch 库名称: transformers 数据集:
• garythung/trashnet

基于视觉变换器(ViT)的垃圾图像分类

本仓库包含使用Hugging Face预训练视觉变换器(ViT)模型实现的图像分类方案。该模型经过微调，可将图像分为六类：纸板、玻璃、金属、纸张、塑料和其他垃圾。

数据集

数据集包含来自garythung/trashnet的六类图像，具体分布如下：

• 纸板：806张
• 玻璃：1002张
• 金属：820张
• 纸张：1188张
• 塑料：964张
• 其他垃圾：274张

模型

我们采用Hugging Face的预训练视觉变换器模型google/vit-base-patch16-224-in21k进行图像分类。该模型在数据集上进行了微调以达到最佳性能。

训练完成的模型已发布于Hugging Face Hub：tribber93/my-trash-classification

使用方法

执行推理的步骤如下：

import torch
import requests
from PIL import Image
from transformers import AutoModelForImageClassification, AutoImageProcessor

url = 'https://cdn.grid.id/crop/0x0:0x0/700x465/photo/grid/original/127308_kaleng-bekas.jpg'
image = Image.open(requests.get(url, stream=True).raw)

model_name = "tribber93/my-trash-classification"
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)
inputs = processor(image, return_tensors="pt")

outputs = model(**inputs)
predictions = torch.argmax(outputs.logits, dim=-1)
print("预测类别:", model.config.id2label[predictions.item()])

训练结果

训练完成后模型表现如下：