pytorch检测代码

以下是一个使用 PyTorch 完成目标检测的简单代码示例：

import torch
import torchvision

# 加载模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)

# 设置模型为评估模式
model.eval()

# 加载图像
image = Image.open("test_image.jpg")

# 对图像进行预处理
transform = transforms.Compose([
    transforms.ToTensor()
])
image = transform(image)

# 将图像转换为批次
images = image.unsqueeze(0)

# 将图像传入模型进行预测
with torch.no_grad():
    outputs = model(images)

# 处理预测结果
boxes = outputs[0]['boxes']
labels = outputs[0]['labels']
scores = outputs[0]['scores']

# 输出预测结果
for i in range(len(boxes)):
    print("Box: ", boxes[i])
    print("Label: ", labels[i])
    print("Score: ", scores[i])

这个示例使用了 Faster R-CNN 网络，这是一个经典的目标检测模型。首先，我们通过 `torchvision.models.detection.fasterrcnn_resnet50_fpn` 加载了一个已经预训练好的 Faster R-CNN 模型。

之后，我们将图像加载并进行了简单的预处理，然后将其转换为一个 PyTorch 张量，并将其转换为一个大小为 `(1, C, H, W)` 的批次张量。

最后，我们将批次张量传入模型，并得到了一个包含预测结果的字典。我们从这个字典中提取出了预测框、标签和分数，并将其输出。