yolov5 特征提取 代码

以下是使用PyTorch实现的YoloV5模型特征提取的代码示例：

import torch
import yaml
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device

def load_yaml(path):
    with open(path, 'r') as f:
        return yaml.load(f, Loader=yaml.FullLoader)

# 加载模型配置文件和权重文件
cfg = load_yaml('yolov5s.yaml')
weights = 'yolov5s.pt'

# 设置使用的设备
device = select_device('')

# 初始化模型
model = attempt_load(weights, map_location=device)

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

# 获取输入图片大小
img_size = cfg['input_size']

# 加载图片
img_path = 'test.jpg'
img = torch.from_numpy(cv2.imread(img_path)).float()
img /= 255.0
img = img.permute(2, 0, 1).unsqueeze(0)

# 将图片送入模型进行特征提取
with torch.no_grad():
    output = model(img.to(device))

# 进行非极大值抑制处理，得到预测框
pred = non_max_suppression(output, cfg['conf_thres'], cfg['iou_thres'], classes=None, agnostic=False)

# 将预测框坐标进行缩放，得到在原图中的坐标
if pred[0] is not None:
    pred[0][:, :4] = scale_coords(img.shape[2:], pred[0][:, :4], img.shape[2:]).round()

print(pred)

需要注意的是，上述代码中的`yolov5s.yaml`文件是自己根据YoloV5官方提供的`yolov5s.yaml`文件修改得到的，其中包括了模型的各种参数配置，可以根据需要进行调整。此外，还需要安装YoloV5所依赖的Python包，可以通过以下命令进行安装：

pip install -r requirements.txt