许可协议：Apache-2.0
标签：

• 目标检测
• 视觉

数据集：

• sku110k

微件示例：

• 图片链接：https://github.com/Isalia20/DETR-finetune/blob/main/IMG_3507.jpg?raw=true
  示例标题：商店示例（非SKU110K数据集）

基于ResNet-50骨干网络并在SKU110K数据集上训练的DETR（端到端目标检测）模型（查询数为400）

该DEtection TRansformer（DETR）模型在SKU110K目标检测数据集（包含8,000张标注图像）上进行了端到端训练。与原始模型的主要区别在于其查询数设置为400，并且是在SKU110K数据集上进行预训练的。

使用方法

以下是使用该模型的示例代码：

from transformers import DetrImageProcessor, DetrForObjectDetection  
import torch  
from PIL import Image, ImageOps  
import requests  

url = "https://github.com/Isalia20/DETR-finetune/blob/main/IMG_3507.jpg?raw=true"  
image = Image.open(requests.get(url, stream=True).raw)  
image = ImageOps.exif_transpose(image)  

# 如果不希望依赖timm库，可以指定修订标签  
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", revision="no_timm")  
model = DetrForObjectDetection.from_pretrained("isalia99/detr-resnet-50-sku110k")  
model = model.eval()  
inputs = processor(images=image, return_tensors="pt")  
outputs = model(**inputs)  

# 将输出（边界框和类别分数）转换为COCO API格式  
# 仅保留置信度分数大于0.8的检测结果  
target_sizes = torch.tensor([image.size[::-1]])  
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.8)[0]  

for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):  
        box = [round(i, 2) for i in box.tolist()]  
        print(  
                f"检测到 {model.config.id2label[label.item()]}，置信度为 "  
                f"{round(score.item(), 3)}，位置为 {box}"  
        )  

预期输出示例：

检测到 LABEL_1，置信度为 0.983，位置为 [665.49, 480.05, 708.15, 650.11]  
检测到 LABEL_1，置信度为 0.938，位置为 [204.99, 1405.9, 239.9, 1546.5]  
...  
检测到 LABEL_1，置信度为 0.998，位置为 [772.85, 169.49, 829.67, 372.18]  
检测到 LABEL_1，置信度为 0.999，位置为 [828.28, 1475.16, 874.37, 1593.43]  

目前，特征提取器和模型均支持PyTorch。

训练数据

DETR模型在SKU110K数据集上训练，该数据集包含8,219/588/2,936张分别用于训练/验证/测试的标注图像。

训练过程

训练

模型分两阶段训练：

1. 仅微调解码器：在1块RTX 4060 Ti GPU上训练140个周期，批次大小为8。
2. 微调整个网络：训练70个周期，批次大小为3，并累积梯度3步。

评估结果

该模型在SKU110K验证集上达到了58.9的mAP（平均精度均值）。结果使用torchmetrics的MeanAveragePrecision类计算。

训练代码

代码已发布在此仓库。请注意，代码尚未完全定稿或充分测试，但核心功能已实现。