python调用yolov5模型识别户型图

可以使用以下步骤调用 YOLOv5 模型来识别户型图：

1.安装 PyTorch 和 YOLOv5：

pip install torch torchvision
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip install -r requirements.txt

2.下载预训练模型：

可以从 https://github.com/ultralytics/yolov5/releases 下载预训练模型。

3.编写 Python 代码进行识别：

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

# 加载模型
weights_path = 'path/to/weights.pt'
device = select_device('cpu')
model = attempt_load(weights_path, map_location=device)

# 图像预处理
img_path = 'path/to/image.jpg'
img_size = 640
img = cv2.imread(img_path)
img = letterbox(img, new_shape=img_size)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(device).float()
img /= 255.0

# 模型推理
model.eval()
with torch.no_grad():
    pred = model(img.unsqueeze(0))

# 后处理
pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)
for i, det in enumerate(pred):
    if len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img.shape[1:]).round()
        for *xyxy, conf, cls in reversed(det):
            label = f'{model.names[int(cls)]} {conf:.2f}'
            plot_one_box(xyxy, img, label=label, color=(0, 255, 0))

# 显示结果
img = img.permute(1, 2, 0).cpu().numpy()
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imshow('result', img)
cv2.waitKey(0)

在上述代码中，我们首先加载了预训练模型，并使用 `letterbox` 函数将输入图像调整为模型输入大小。然后，我们将图像转换为 PyTorch 张量，并使用模型进行推理。最后，我们使用 `non_max_suppression` 函数过滤掉重叠的边界框，并使用 `plot_one_box` 函数将边界框绘制在图像上。