替换SlowFast中Detectron2为Yolov8

一 需求

复制代码
FaceBookReserch中SlowFast源码中检测框是用Detectron2进行目标检测,本文想实现用yolov8替换detectron2

二 实施方案

首先,yolov8 支持有自定义库ultralytics(仅支持yolov8),安装对应库

bash 复制代码
pip install ultralytics

源码中slowfast/visualization.py 43行中

python 复制代码
if cfg.DETECTION.ENABLE:
       self.object_detector = Detectron2Predictor(cfg, gpu_id=self.gpu_id)

根据ultralytics文档进行定义

创建对应YOLOPredictor类(加入了检测框及其标签,具体见前一篇文章)

python 复制代码
class YOLOPredictor:

    def __init__(self, cfg, gpu_id=None):
        # 加载预训练的 YOLOv8n 模型
        self.model = YOLO('/root/autodl-tmp/data/runs/detect/train/weights/best.pt')
        self.detect_names, _, _ = get_class_names(cfg.DEMO.Detect_File_Path, None, None)

    def __call__(self, task):
        """
        Return bounding boxes predictions as a tensor.
        Args:
            task (TaskInfo object): task object that contain
                the necessary information for action prediction. (e.g. frames)
        Returns:
            task (TaskInfo object): the same task info object but filled with
                prediction values (a tensor) and the corresponding boxes for
                action detection task.
        """
        # """得到预测置信度"""
        # scores = outputs["instances"].scores[mask].tolist()
        # """获取类别标签"""
        # pred_labels = outputs["instances"].pred_classes[mask]
        # pred_labels = pred_labels.tolist()
        # """进行标签匹配"""
        # for i in range(len(pred_labels)):
        #     pred_labels[i] = self.detect_names[pred_labels[i]]
        # preds = [
        #     "[{:.4f}] {}".format(s, labels) for s, labels in zip(scores, pred_labels)
        # ]
        # """加入预测标签"""
        # task.add_detect_preds(preds)
        # task.add_bboxes(pred_boxes)
        middle_frame = task.frames[len(task.frames) // 2]
        outputs = self.model(middle_frame)
        boxes = outputs[0].boxes
        mask = boxes.conf >= 0.5
        pred_boxes = boxes.xyxy[mask]
        scores = boxes.conf[mask].tolist()
        pred_labels = boxes.cls[mask].to(torch.int)
        pred_labels = pred_labels.tolist()
        for i in range(len(pred_labels)):
            pred_labels[i] = self.detect_names[pred_labels[i]]
        preds = [
            "[{:.4f}] {}".format(s, labels) for s, labels in zip(scores, pred_labels)
        ]
        """加入预测标签"""
        task.add_detect_preds(preds)
        task.add_bboxes(pred_boxes)

        return task
相关推荐
无风听海7 小时前
神经网络之窗口大小对词语义向量的影响
人工智能·深度学习·神经网络
sali-tec7 小时前
C# 基于halcon的视觉工作流-章52-生成标定板
开发语言·图像处理·人工智能·算法·计算机视觉
newxtc8 小时前
【江苏政务服务网-注册_登录安全分析报告】
人工智能·安全·yolo·政务·1024程序员节·安全爆破
Tiandaren8 小时前
自用提示词02 || Prompt Engineering || RAG数据切分 || 作用:通过LLM将文档切分成chunks
数据库·pytorch·深度学习·oracle·prompt·rag
阿水实证通8 小时前
面向社科研究者:用深度学习做因果推断(二)
深度学习·1024程序员节·因果推断·实证分析·科研创新
王哈哈^_^9 小时前
【数据集】【YOLO】【目标检测】口罩数据集,口罩佩戴识别数据集 1971 张,YOLO佩戴口罩检测算法实战训练教程。
人工智能·算法·yolo·目标检测·计算机视觉·ai·视觉检测
房开民9 小时前
OpenCV C++ 中,访问图像像素三种常用方法
c++·opencv·计算机视觉
文火冰糖的硅基工坊9 小时前
[人工智能-大模型-83]:模型层技术 - 前向预测:神经网络是如何产生涌现智能的?背后的本质是什么?
人工智能·深度学习·神经网络
子不语18010 小时前
深度学习——IDE之Jupyter
人工智能·深度学习·jupyter
格林威11 小时前
紫外工业相机入门介绍和工业检测核心场景
人工智能·数码相机·计算机视觉·目标跟踪·视觉检测