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

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

importtorchimportrequests from PILimportImagefrom transformersimportDetrImageProcessor,DetrForObjectDetectionimportmatplotlib.pyplotasplt #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")withtorch.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()forscore,label,boxinzip(results["scores"],results["labels"],results["boxes"]):box=[round(i,2)foriinbox.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()