yolop 算子

yolop 从特定图像中识别车道和可驾驶区域。

更多信息在此: https://github.com/hustvl/YOLOP

您还可以选择使用环境变量在 GPU 中分配模型：

• PYTORCH_DEVICE: cuda # 或 cpu

输入

• 图像: 高 x 宽 x BGR array.

输出

• drivable_area: 驾驶面积做为轮廓数据点
• lanes: lanes(本变量)的60个点值表示车道

示例绘制 (车道是红色的，可驾驶区域是绿色的)

图描述

  - id: yolop
    operator: 
      outputs:
        - lanes
        - drivable_area
      inputs:
        image: webcam/image
      python: ../../operators/yolop_op.py

方法

__init__()

源码

    def __init__(self):
        self.model = torch.hub.load("hustvl/yolop", "yolop", pretrained=True)
        self.model.to(torch.device(DEVICE))
        self.model.eval()

.on_event(...)

源码

    def on_event(
        self,
        dora_event: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        if dora_event["type"] == "INPUT":
            return self.on_input(dora_event, send_output)
        return DoraStatus.CONTINUE

.on_input(...)

源码

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        # 推断
        frame = cv2.imdecode(
            np.frombuffer(
                dora_input["data"],
                np.uint8,
            ),
            -1,
        )

        frame = frame[:, :, :3]
        h0, w0, _ = frame.shape
        h, w = (640, 640)
        frame, _, (pad_w, pad_h) = letterbox_for_img(frame)
        ratio = w / w0
        pad_h, pad_w = (int(pad_h), int(pad_w))

        img = torch.unsqueeze(transform(frame), dim=0)
        half = False  # half precision only supported on CUDA
        img = img.half() if half else img.float()  # uint8 to fp16/32
        img = img.to(torch.device(DEVICE))
        det_out, da_seg_out, ll_seg_out = self.model(img)

        # det_out = [pred.reshape((1, -1, 6)) for pred in det_out]
        # inf_out = torch.cat(det_out, dim=1)

        # det_pred = non_max_suppression(
        # inf_out,
        # )
        # det = det_pred[0]

        da_predict = da_seg_out[:, :, pad_h : (h0 - pad_h), pad_w : (w0 - pad_w)]
        da_seg_mask = torch.nn.functional.interpolate(
            da_predict, scale_factor=1 / ratio, mode="bilinear"
        )
        _, da_seg_mask = torch.max(da_seg_mask, 1)
        da_seg_mask = da_seg_mask.int().squeeze().cpu().numpy()
        da_seg_mask = morphological_process(da_seg_mask, kernel_size=7)

        contours, _ = cv2.findContours(
            da_seg_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        if len(contours) != 0:
            contour = max(contours, key=cv2.contourArea)
            contour = contour.astype(np.int32)
            send_output("drivable_area", contour.tobytes(), dora_input["metadata"])
        else:
            send_output("drivable_area", np.array([]).tobytes(), dora_input["metadata"])

        ll_predict = ll_seg_out[:, :, pad_h : (h0 - pad_h), pad_w : (w0 - pad_w)]

        ll_seg_mask = torch.nn.functional.interpolate(
            ll_predict, scale_factor=1 / ratio, mode="bilinear"
        )

        _, ll_seg_mask = torch.max(ll_seg_mask, 1)
        ll_seg_mask = ll_seg_mask.int().squeeze().cpu().numpy()
        # Lane line post-processing
        ll_seg_mask = morphological_process(
            ll_seg_mask, kernel_size=7, func_type=cv2.MORPH_OPEN
        )
        ll_seg_points = np.array(connect_lane(ll_seg_mask), np.int32)
        send_output("lanes", ll_seg_points.tobytes(), dora_input["metadata"])
        return DoraStatus.CONTINUE