42 lines
910 B
Python
42 lines
910 B
Python
import torch
|
|
from torchvision.models.detection import (
|
|
fasterrcnn_resnet50_fpn,
|
|
FasterRCNN_ResNet50_FPN_Weights
|
|
)
|
|
from PIL import Image
|
|
import torchvision.transforms as T
|
|
|
|
# Load pretrained model
|
|
weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT
|
|
model = fasterrcnn_resnet50_fpn(weights=weights)
|
|
|
|
model.eval()
|
|
|
|
# Load image
|
|
image = Image.open("image.jpg").convert("RGB")
|
|
|
|
# Convert image to tensor
|
|
transform = T.ToTensor()
|
|
image_tensor = transform(image)
|
|
|
|
# Run detection
|
|
with torch.no_grad():
|
|
predictions = model([image_tensor])
|
|
|
|
prediction = predictions[0]
|
|
|
|
# COCO class names
|
|
categories = weights.meta["categories"]
|
|
|
|
# Print detected objects
|
|
for box, label, score in zip(
|
|
prediction["boxes"],
|
|
prediction["labels"],
|
|
prediction["scores"]
|
|
):
|
|
if score > 0.5:
|
|
print(
|
|
f"{categories[label]}: "
|
|
f"{score:.2f}, "
|
|
f"box={box.tolist()}"
|
|
) |