输入一张动物的图片进行目标检测和分类
!pip install yolov5
import torch
from PIL import Image
from torchvision import transforms
from yolov5.models.experimental import attempt_load
from yolov5.utils.general import non_max_suppression
# 加载YOLOv5模型
device = torch.device('cpu') # 或者使用torch.device('cuda')以使用GPU
model = attempt_load("yolov5s.pt", device=device)
# 设置模型为评估模式
model.eval()
# 加载图像
image_path = "images/example_image.jpg"
image = Image.open(image_path)
# 调整图像大小为模型期望的输入大小（416x416）
resize_transform = transforms.Resize((416, 416))
image_resized = resize_transform(image)
# 定义图像转换
transform = transforms.Compose([
transforms.ToTensor(),
])
# 对图像进行转换
input_tensor = transform(image_resized).unsqueeze(0).to(device)
# 进行目标检测
with torch.no_grad():
results = model(input_tensor)
# 对结果进行非极大值抑制
results = non_max_suppression(results, conf_thres=0.3)
# 如果检测到目标，打印检测到的目标信息
if results is not None and len(results[0]) > 0:
for detection in results[0]:
print(f"类别: {detection[-1]}, 置信度: {detection[4]}, 边界框: {detection[:4]}")
else:
print("未检测到任何目标。")
返回结果： Fusing layers... YOLOv5s summary: 270 layers, 7235389 parameters, 0 gradients, 16.6 GFLOPs
类别: 17.0, 置信度: 0.8820466995239258, 边界框: tensor([ 95.17632, 88.61916, 327.81903, 387.15582])
图片 baidu
