Yolov5 算子

Yolov5 对象检测算子 在检测到对象的图像上生成边界框。

更多信息在此: https://github.com/ultralytics/yolov5

Yolov5 尚未在模拟中进行微调，而是直接从 Pytorch Hub 导入权重信息。

如果您想在没有互联网的情况下运行yolov5，你可以 克隆 https://github.com/ultralytics/yolov5 并 从 releases 页面 下载您要使用的权重文件，然后在 yaml 图中指定两个环境变量：

• YOLOV5_PATH: 您的本地Yolov5路径
• YOLOV5_WEIGHT_PATH: 您的本地 Yolov5 权重文件路径

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

• PYTORCH_DEVICE: cuda # 或 cpu

输入

• 图像: 高 x 宽 x BGR array.

输出

• bbox: N_BBOX, X_MIN, X_MAX, Y_MIN, Y_MAX, CONDIDENCE, CLASS, array

示例图像

图描述

  - id: yolov5
    operator: 
      outputs:
        - bbox
      inputs:
        image: webcam/image
      python: ../../operators/yolov5_op.py

图可视化

方法

__init__()

源码

    def __init__(self):
        if YOLOV5_PATH is None:
            # With internet
            self.model = torch.hub.load(
                "ultralytics/yolov5",
                "yolov5n",
            )
        else:
            # 没有互联网
            #
            # 至安装:
            # cd $DORA_HOME_DEP/dependecies # 可选
            # git clone https://github.com/ultralytics/yolov5.git
            # rm yolov5/.git -rf
            # 在您的 YAML 图 添加 YOLOV5_PATH 与 YOLOV5_WEIGHT_PATH 

            self.model = torch.hub.load(
                YOLOV5_PATH,
                "custom",
                path=YOLOV5_WEIGHT_PATH,
                source="local",
            )

        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(...)

图像句柄 参数: dora_input["id"] (str): 在 yaml 配置中的输入已声明的 Id dora_input["value"] (arrow.array (UInt8)): 字节形式的输入消息 send_output (Callable[[str, bytes]]): 函数启用将输出发送回 dora。

源码

    def on_input(
        self,
        dora_input: dict,
        send_output: Callable[[str, bytes], None],
    ) -> DoraStatus:
        """
        Handle image
        Args:
            dora_input["id"] (str): 在 yaml 配置中的输入已声明的Id 
            dora_input["value"] (arrow.array (UInt8)): 字节形式的输入消息
            send_output (Callable[[str, bytes]]): 函数启用将输出发送回 dora。
        """
        if dora_input["id"] == "image":
            frame = (
                dora_input["value"].to_numpy().reshape((IMAGE_HEIGHT, IMAGE_WIDTH, 4))
            )
            frame = frame[:, :, :3]

            results = self.model(frame)  # 包含 NMS (non max suppression，非极大值抑制)
            arrays = np.array(results.xyxy[0].cpu())[
                :, [0, 2, 1, 3, 4, 5]
            ]  # xyxy -> xxyy
            arrays[:, 4] *= 100
            arrays = arrays.astype(np.int32)
            arrays = pa.array(arrays.ravel())
            send_output("bbox", arrays, dora_input["metadata"])
            return DoraStatus.CONTINUE