文章目录
- 前言
- 视频效果
- 必要环境
- 一、代码结构
- 1、 训练参数解析
- 2、 核心代码解析
- 1.初始化Detector类
- 2. @torch.no_grad()
- 3. 复制输入图像并初始化计数器
- 4. 调用YOLOv10模型进行目标检测
- 5. 提取检测结果信息
- 6. 遍历检测结果并在图像上绘制边界框和标签
- 7. 准备输入图像以适应End-to-end模型
- 8. 使用YOLOP模型进行推理
- 9. 处理可行驶区域分割结果
- 10. 处理车道线分割结果
- 二、效果展示
- 三、完整代码获取
- 总结
前言
在往期博客中,我们详细介绍了如何搭建YOLOv10和YOLOP的环境。本期将结合这两个算法,实现多类别目标检测、可行驶区域分割和车道线分割等多种任务,并将其部署到PYQT界面中进行展示。
视频效果
b站链接:基于YOLOv10+YOLOP+PYQT的可视化系统,实现多类别目标检测+可行驶区域分割+车道线分割多种任务
必要环境
- 配置yolov10环境 可参考往期博客
地址:搭建YOLOv10环境 训练+推理+模型评估
- 配置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()
关键参数详解:
-
–v10weights: 指定YOLOv10模型的权重文件路径。
-
–weights: 指定YOLOP模型的权重文件路径,这个模型包含了车道线分割和可行驶区域分割的任务
-
–device: 指定运行模型的设备,可以是单个GPU(如 0),或者是CPU(cpu)
-
–imgsz: 指定输入图像的尺寸,输入图像会被调整为这个尺寸,以适应模型的输入要求
-
–conf_thre: 设置初始置信度阈值,只有置信度高于这个阈值的检测框才会被保留
-
–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()
二、效果展示
三、完整代码获取
链接:基于YOLOv10+YOLOP+PYQT的可视化系统,实现多类别目标检测+可行驶区域分割+车道线分割
总结
本期博客就到这里啦,喜欢的小伙伴们可以点点关注,感谢!
最近经常在b站上更新一些有关目标检测的视频,大家感兴趣可以来看看 https://b23.tv/1upjbcG
学习交流群:995760755