当前位置：首页 > news >正文

替换SlowFast中Detectron2为Yolov8

news 2025/8/6 5:58:38

一需求

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

二实施方案

首先，yolov8 支持有自定义库ultralytics（仅支持yolov8），安装对应库

pip install ultralytics

源码中slowfast/visualization.py 43行中

if cfg.DETECTION.ENABLE:self.object_detector = Detectron2Predictor(cfg, gpu_id=self.gpu_id)

根据ultralytics文档进行定义
创建对应YOLOPredictor类(加入了检测框及其标签，具体见前一篇文章)

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 containthe necessary information for action prediction. (e.g. frames)Returns:task (TaskInfo object): the same task info object but filled withprediction values (a tensor) and the corresponding boxes foraction 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].boxesmask = boxes.conf >= 0.5pred_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