加载预训练的 Faster R-CNN 模型进行物体检查

import torchvision.transforms as T
from PIL import Image
import matplotlib.pyplot as plt

import torchvision
import torch

# 加载预训练的 Faster R-CNN 模型
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
model = model.eval()

# 移除分类器
# model.roi_heads.box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(1024, num_classes=1)
# Replace 'your_image_path.jpg' with the path to your actual image
image_path = 'I:/GTOT/BlackCar/v/00002v.png'

# Load the image
image = Image.open(image_path).convert("RGB")
image = image.resize((256, 256))


# Define the transformation
transform = T.Compose([T.ToTensor()])

# Apply the transformation to the image
input_image = transform(image).unsqueeze(0)

# Run the model
with torch.no_grad():
    prediction = model(input_image)

# Display the original image with predicted boxes
boxes = prediction[0]['boxes'].numpy()
scores = prediction[0]['scores'].numpy()

# Set a threshold for displaying boxes (adjust as needed)
threshold =0.9
selected_boxes = boxes[scores > threshold]

# 创建 ax 对象
fig, ax = plt.subplots(1, figsize=(12, 9))

# Display the image with boxes
ax.imshow(image)

for box in selected_boxes:
    # Draw a rectangle on the image
    rect = plt.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], linewidth=2, edgecolor='r', facecolor='none')
    ax.add_patch(rect)

plt.axis('off')
plt.show()
print(prediction)
print(input_image.shape)
print("Predicted Boxes:", boxes)
print("Confidence Scores:", scores)