给我一段可执行的python代码,用yolov5库检测图片中的目标物体
时间: 2024-03-17 14:41:27 浏览: 14
好的,以下是一段可执行的Python代码,使用yolov5库来检测图片中的目标物体:
```python
import torch
import cv2
import argparse
import numpy as np
from PIL import Image
# 定义命令行参数
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='yolov5s.pt', help='model.pt path')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[width, height]')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
parser.add_argument('source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class')
args = parser.parse_args()
# 加载模型
device = torch.device(args.device)
model = torch.hub.load('ultralytics/yolov5', 'custom', path_or_model=args.weights).to(device)
# 加载图片
img0 = Image.open(args.source).convert('RGB')
# 图片预处理
img = np.array(img0)
if img.shape[0] != args.img_size[1] or img.shape[1] != args.img_size[0]:
img = cv2.resize(img, tuple(args.img_size[::-1]))
img = img.transpose(2, 0, 1) # HWC to CHW
img = img[np.newaxis, ...] # 增加batch维度
img = torch.from_numpy(img).to(device).float() # numpy to tensor
# 进行目标检测
results = model(img, conf_thres=args.conf_thres, iou_thres=args.iou_thres)
# 打印检测结果
print(results.pandas().xyxy[0])
# 显示图片
if args.view_img:
results.render()
cv2.imshow('results', results.imgs[0])
cv2.waitKey(0)
# 保存结果到txt
if args.save_txt:
results.save()
```
这段代码使用了yolov5库进行目标检测,可以在命令行中指定要检测的图片路径、模型权重文件路径、以及一些参数,例如目标置信度阈值、NMS阈值等。程序会输出检测结果并将结果可视化展示出来,也可以选择将结果保存到txt文件中。