视觉检测大模型DRTR部署使用

import torch
import requests
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
import matplotlib.pyplot as plt

# 1. 加载模型和特征提取器
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
model.eval()

# 2. 读取图片
url = "http://images.cocodataset.org/val2017/000000039769.jpg"  # 示例图片
image = Image.open(requests.get(url, stream=True).raw)

# 3. 预处理并推理
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)

# 4. 后处理（DETR 无需 NMS，直接输出结果）
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]

# 5. 显示检测框
plt.figure(figsize=(12, 9))
plt.imshow(image)
ax = plt.gca()

for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
    box = [round(i, 2) for i in box.tolist()]
    # 绘制矩形框
    rect = plt.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1],
                         fill=False, color="red", linewidth=2)
    ax.add_patch(rect)
    # 添加标签
    ax.text(box[0], box[1], f"{model.config.id2label[label.item()]}: {score.item():.3f}",
            bbox=dict(facecolor="red", alpha=0.5), fontsize=10, color="white")

plt.axis('off')
plt.show()