基于yolov5写一个对图片目标检测的代码
时间: 2024-05-12 09:17:12 浏览: 7
好的,我可以帮你回答,以下是基于 yolov5 的图片目标检测代码实现:
```
import torch
import cv2
from PIL import Image
from torch.backends import cudnn
from models.experimental import attempt_load
from utils.general import check_img_size, non_max_suppression, apply_classifier, scale_coords
from utils.torch_utils import select_device, time_sync
def detect(yolo, img, device):
img_tensor = torch.from_numpy(img).to(device)
pred = yolo(img_tensor)
pred = non_max_suppression(pred, conf_thres=0.4, iou_thres=0.5, classes=None, agnostic=False)
if len(pred[0]):
det = pred[0].cpu().detach().numpy()
return det
else:
return []
if __name__ == '__main__':
device = select_device('')
half = device.type != 'cpu'
yolov5_model = attempt_load('yolov5s.pt', map_location=device)
img = cv2.imread('test.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_size = check_img_size(img, s=yolov5_model.stride.max())
img = cv2.resize(img, img_size)
if half:
yolov5_model.half()
cudnn.benchmark = True
yolov5_model.eval()
with torch.no_grad():
det = detect(yolov5_model, img, device)
if len(det)>0:
print('检测到的目标数量:{}'.format(len(det)))
for obj in det:
x1, y1, x2, y2, conf, cls = obj[0], obj[1], obj[2], obj[3], obj[4], obj[5]
print('置信度:{:0.2f}, 类别: {}'.format(conf, cls))
cv2.rectangle(img, (int(x1), int(y1)), (int(x2), int(y2)), (255, 0, 0), 2)
cv2.putText(img, 'conf:{:0.2f}'.format(conf), (int(x1), int(y1)-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
cv2.imshow('result', img)
cv2.waitKey(0)
```
希望这能回答您的问题,如果您有其他问题,可以继续问我哦。