一 需求

FaceBookReserch中SlowFast源碼中檢測框是用Detectron2進行目標檢測，本文想實現(xiàn)用yolov8替換detectron2

二 實施方案

首先，yolov8 支持有自定義庫ultralytics（僅支持yolov8），安裝對應(yīng)庫

pip install ultralytics

源碼中slowfast/visualization.py 43行中

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

根據(jù)ultralytics文檔進行定義
創(chuàng)建對應(yīng)YOLOPredictor類(加入了檢測框及其標簽，具體見前一篇文章)

class YOLOPredictor:def __init__(self, cfg, gpu_id=None):# 加載預(yù)訓(xùn)練的 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."""# """得到預(yù)測置信度"""# 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)# ]# """加入預(yù)測標簽"""# 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)]"""加入預(yù)測標簽"""task.add_detect_preds(preds)task.add_bboxes(pred_boxes)return task