yolov5检测代码

### YOLOv5目标检测代码实现示例

对于YOLOv5的目标检测，其核心在于模型加载、图像预处理、推理以及结果解析几个方面。下面提供了一个简单的Python脚本用于展示如何利用YOLOv5执行目标检测。

from models.experimental import attempt_load
import torch
from utils.general import non_max_suppression, scale_coords
from utils.datasets import letterbox
import cv2
import numpy as np


def detect_objects(image_path='data/images/bus.jpg', weights='yolov5s.pt'):
    """
    使用YOLOv5进行对象检测并返回标注后的图片.
    
    参数:
        image_path (str): 输入待检测的图片路径，默认为'data/images/bus.jpg'.
        weights (str): 预训练权重文件的位置，默认为'yolov5s.pt'.

    返回:
        img_with_boxes (numpy.ndarray): 绘制边框和标签的结果图.
    """

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = attempt_load(weights=weights).to(device)

    # 设置模型为评估模式
    model.eval()

    # 加载原始图片
    original_img = cv2.imread(image_path)
    h0, w0 = original_img.shape[:2]

    # 图像预处理：调整大小到640x640，并应用letterboxing保持纵横比不变
    processed_image, ratio, pad = letterbox(original_img.copy(), new_shape=(640, 640), auto=True)[0], None, None
    
    # 转换颜色通道顺序(BGR->RGB)，转置维度(HWC->CHW)，转换数据类型(uint8->float32)
    processed_image = processed_image[:, :, ::-1].transpose(2, 0, 1)  
    processed_image = np.ascontiguousarray(processed_image.astype(np.float32))

    with torch.no_grad():
        input_tensor = torch.from_numpy(processed_image).unsqueeze(dim=0).to(device=device)/255.

        predictions = model(input_tensor)[0]
        
        prediction_results = non_max_suppression(predictions=predictions,
                                                 conf_thres=0.25,
                                                 iou_thres=0.45,
                                                 classes=None,
                                                 agnostic=False)

        results = []
        for det in prediction_results:  # 每张图片的结果
            if len(det):
                boxes = det[:, :4]
                
                # 将边界框坐标缩放回原图尺寸
                boxes = scale_coords(img1_shape=input_tensor.shape[2:], 
                                     coords=boxes, 
                                     img0_shape=original_img.shape).round()
            
                labels = det[:, -1].int().tolist()
                scores = det[:, 4].tolist()

                result_dict = {'boxes': boxes.tolist(),
                               'labels': labels,
                               'scores': scores}
                results.append(result_dict)


    # 可视化预测结果
    img_with_boxes = original_img.copy()
    colors = [(np.random.randint(0, 255), np.random.randint(0, 255),
               np.random.randint(0, 255)) for _ in range(len(results))]
    
    for idx, res in enumerate(results):
        box_color = colors[idx % len(colors)]
        for bbox, label_id, score in zip(res['boxes'], res['labels'], res['scores']):
            p1, p2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3]))
            text = f'{label_id} {score:.2f}'
            font_face = cv2.FONT_HERSHEY_SIMPLEX
            
            cv2.rectangle(img_with_boxes, pt1=p1, pt2=p2, color=box_color, thickness=2)
            cv2.putText(img_with_boxes, org=(p1[0]+5, p1[1]-5), text=text, 
                        fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.7, 
                        color=box_color, thickness=2)

    return img_with_boxes

此段代码展示了如何通过调用`detect_objects()`函数来完成一次完整的物体识别过程[^1]。该函数接收一张输入图片及其对应的YOLOv5预训练模型参数作为输入，经过一系列操作之后输出带有标记框的对象检测效果图。