基于YOLOv10+YOLOP+PYQT的可视化系统，实现多类别目标检测+可行驶区域分割+车道线分割【附代码】

---

前言

在往期博客中，我们详细介绍了如何搭建YOLOv10和YOLOP的环境。本期将结合这两个算法，实现多类别目标检测、可行驶区域分割和车道线分割等多种任务，并将其部署到PYQT界面中进行展示。

---

视频效果

基于YOLOv10+YOLOP+PYQT的可视化系统，实现多类别目标检测+可行驶区域分割+车道线分割多种任务【附代码】

---

必要环境

1. 配置yolov10环境 可参考往期博客
   地址：搭建YOLOv10环境 训练+推理+模型评估
   
2. 配置yolop环境 可参考往期博客
   地址：YOLOP 训练+测试+模型评估
   

---

一、代码结构

1、 训练参数解析

首先，我们利用 argparse 模块来设置命令行参数，以便灵活配置模型的权重路径、使用设备(cpu、gpu)等信息

# 解析命令行参数
parser.add_argument('--v10weights', default=r"yolov10s.pt", type=str, help='weights path')
parser.add_argument('--weights', default=r"weights/End-to-end.pth", type=str, help='weights path')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--imgsz', type=int, default=640, help='image size')
parser.add_argument('--merge_nms', default=False, action='store_true', help='merge class')
parser.add_argument('--conf_thre', type=float, default=0.3, help='conf_thre')
parser.add_argument('--iou_thre', type=float, default=0.2, help='iou_thre')
parser.add_argument('--augment', action='store_true', help='augmented inference')
opt = parser.parse_args()

关键参数详解：

1. –v10weights: 指定YOLOv10模型的权重文件路径。

2. –weights: 指定YOLOP模型的权重文件路径，这个模型包含了车道线分割和可行驶区域分割的任务

3. –device: 指定运行模型的设备，可以是单个GPU（如 0），或者是CPU（cpu）

4. –imgsz: 指定输入图像的尺寸，输入图像会被调整为这个尺寸，以适应模型的输入要求

5. –conf_thre: 设置初始置信度阈值，只有置信度高于这个阈值的检测框才会被保留

6. –iou_thre: 设置初始IOU阈值，在NMS过程中，只有IOU低于这个阈值的检测框才会被保留

2、 核心代码解析

此部分包含车道线分割、可行驶区域分割和目标检测等关键部分的代码

1.初始化Detector类

这段代码定义了一个名为Detector的类，该类初始化了两个模型：一个是用于End-to-end检测的YOLOP模型，另一个是用于目标检测的YOLOv10模型。通过加载权重文件、设置设备、调整图像大小以及配置模型参数，实现了对这两个模型的初始化和准备工作

class Detector:
    def __init__(self, v10weights, cfg, device, model_path=r'./best_dist_model.pt', imgsz=640, conf=0.5, iou=0.0625, merge_nms=False):
        self.device = device
        self.model = get_net(cfg)
        checkpoint = torch.load(model_path, map_location=device)
        self.model.load_state_dict(checkpoint['state_dict'])
        self.model = self.model.to(device)
        img_w = torch.zeros((1, 3, imgsz, imgsz), device=device)
        _ = self.model(img_w)
        self.model.eval()

        self.stride = int(self.model.stride.max())
        self.imgsz = check_img_size(imgsz, s=self.stride)
        self.merge_nms = merge_nms

        self.model_v10 = YOLOv10(v10weights)
        self.names = self.model_v10.names

2. @torch.no_grad()

这是一个装饰器，用于禁用梯度计算，可以减少内存消耗并加快推理速度，通常在推理时使用

@torch.no_grad()
def __call__(self, image: np.ndarray, conf, iou):

3. 复制输入图像并初始化计数器

复制输入图像以便在结果图像上进行操作，并初始化一个默认字典来记录每个类别的检测次数

img_vis = image.copy()
class_counts = defaultdict(int)

4. 调用YOLOv10模型进行目标检测

使用YOLOv10模型在输入图像上进行目标检测，返回检测结果

results = self.model_v10(image, verbose=True, conf=conf, iou=iou, device=self.device)

5. 提取检测结果信息

提取检测结果中的类别、置信度和边界框坐标

bboxes_cls = results[0].boxes.cls
bboxes_conf = results[0].boxes.conf
bboxes_xyxy = results[0].boxes.xyxy.cpu().numpy().astype('uint32')

6. 遍历检测结果并在图像上绘制边界框和标签

遍历所有检测到的目标，在图像上绘制边界框和标签，并记录每个类别的检测次数

for idx in range(len(bboxes_cls)):
    box_cls = int(bboxes_cls[idx])
    bbox_xyxy = bboxes_xyxy[idx]
    bbox_label = self.names[box_cls]
    class_counts[bbox_label] += 1
    box_conf = f"{bboxes_conf[idx]:.2f}"
    xmax, ymax, xmin, ymin = bbox_xyxy[2], bbox_xyxy[3], bbox_xyxy[0], bbox_xyxy[1]

    img_vis = cv2.rectangle(img_vis, (xmin, ymin), (xmax, ymax), get_color(box_cls + 2), 3)
    cv2.putText(img_vis, f'{str(bbox_label)}/{str(box_conf)}', (xmin, ymin - 10),
                cv2.FONT_HERSHEY_SIMPLEX, 1.0, get_color(box_cls + 2), 3)

7. 准备输入图像以适应End-to-end模型

对输入图像进行调整和预处理，以适应End-to-end模型的输入要求

img, ratio, pad = letterbox_for_img(image, new_shape=self.imgsz, auto=True)
pad_w, pad_h = pad
pad_w = int(pad_w)
pad_h = int(pad_h)
ratio = ratio[1]
img = np.ascontiguousarray(img)
img = transform(img).to(self.device)
im = img.float()
if im.ndimension() == 3:
    im = im.unsqueeze(0)

8. 使用YOLOP模型进行推理

在预处理后的图像上运行End-to-end模型，输出检测结果、车道线分割结果和可行驶区域分割结果

det_out, da_seg_out, ll_seg_out = self.model(im)

9. 处理可行驶区域分割结果

这段代码将对可行驶区域的分割结果进行后处理，首先从模型输出中裁剪出实际的分割结果，通过双线性插值恢复到原始图像尺寸，然后提取每个像素的类别索引，最终生成可行驶区域的分割掩码

_, _, height, width = im.shape
da_predict = da_seg_out[:, :, pad_h:(height - pad_h), pad_w:(width - pad_w)]
da_seg_mask = torch.nn.functional.interpolate(da_predict, scale_factor=int(1 / ratio), mode='bilinear')
_, da_seg_mask = torch.max(da_seg_mask, 1)
da_seg_mask = da_seg_mask.int().squeeze().cpu().numpy()

10. 处理车道线分割结果

这段代码将对车道线分割结果进行后处理，和处理可行驶区域分割结果同理，首先从模型输出中裁剪出实际的分割结果，并通过双线性插值恢复到原始图像尺寸，然后提取每个像素的类别索引，生成最终的分割掩码

ll_predict = ll_seg_out[:, :, pad_h:(height - pad_h), pad_w:(width - pad_w)]
ll_seg_mask = torch.nn.functional.interpolate(ll_predict, scale_factor=int(1 / ratio), mode='bilinear')
_, ll_seg_mask = torch.max(ll_seg_mask, 1)
ll_seg_mask = ll_seg_mask.int().squeeze().cpu().numpy()

二、效果展示

---

总结

本期博客就到这里啦，喜欢的小伙伴们可以点点关注，感谢！