Yolo的工具包Supervision 目标检测（图片，视频）

Supervision 目标检测（图片，视频）

1. 安装

pip install supervision   
pip install ultralytics   # Yolov8 提供的工具包

2. 代码实现

2.1 目标检测 （图片）

采用Supervision 目标检测，工具包集成 Yolov8 模型

1. 导入需要的包

   import cv2
   import supervision as sv
   
   from ultralytics import YOLO
   import matplotlib.pyplot as plt
   

2. 加载模型，得到模型训练的结果

   # yolov8s.pt  yolov8n.pt 
   model = YOLO("yolov8x.pt")
   image = cv2.imread("F:/ai/supervision/material/pic/微信图片_20240605122212.jpg")
   results = model(image, verbose=False)[0]
   # 从YOLOv8中加载数据结果
   detections = sv.Detections.from_ultralytics(results)
   # 查看输出结果
   detections
   

   其中 detections，是模型预测的结果信息

   Detections(
       xyxy=array([[     1.1242,      230.52,      804.28,      741.87],
          [     668.53,      392.16,      809.58,      879.81],
          [     50.536,      398.19,      247.58,       900.9],
          [     221.96,      406.15,      343.99,      860.09],
          [    0.21836,      550.76,      78.427,      872.18],
          [    0.23819,      550.13,      79.402,      1064.2],
          [     666.25,      15.416,      748.06,      90.126]], dtype=float32), 
       mask=None, 
       confidence=array([    0.96796,     0.93112,     0.91831,     0.90192,     0.69789,     0.53844,     0.43446], dtype=float32), 
       class_id=array([5, 0, 0, 0, 0, 0, 1]), 
       tracker_id=None, 
       data={'class_name': array(['bus', 'person', 'person', 'person', 'person', 'person', 'bicycle'], dtype='<U7')}
   )
   

   输出参数

   xyxy : 模型最大和最小的坐标点；

   confidence : 模型预测最大概率值；

   data ： 分类名称 class_name

3. 可视化结果

   我这里采用的是supervsion 自带的可视化的方法，当然你也可以根据detections 所提供的信息采用别的可视化的方法（opencv）

   # 可视化识别结果
    
   # 确定可视化参数
   bounding_box_annotator = sv.BoundingBoxAnnotator()  # 初始化外框
   label_annotator = sv.LabelAnnotator()  # 初始化文字框
   
   labels = [
       f"{class_name} {confidence:.2f}"
       for class_name, confidence
       in zip(detections['class_name'], detections.confidence)
   ]
   #绘制外框
   annotated_image = bounding_box_annotator.annotate(
       scene=image, detections=detections)
   #绘制文字
   annotated_image = label_annotator.annotate(
       scene=annotated_image, detections=detections, labels=labels)
   sv.plot_image(annotated_image)
   

4. 结果

2.2 目标识别 （视频）

1. 需要安装的包

   inference==0.9.17
   supervision==0.19.0
   tqdm
   ultralytics
   

2. 代码实现

   该代码需要上传4个参数 ：

   • --source_weights_path: Yolo模型参数权重文件
   • --source_video_path: 源文件所在的文件位置
   • --target_video_path: 输出的文件所在位置，用于保存识别后的视频文件
   • --confidence_threshold (optional): 设置模型的信任值，默认是0.3。越高的值对模型具有多选择性，较低的参数模型有可能识别成为多种可能
   • --iou_threshold (optional): IOU 模型的阈值，默认是0.7。这个参数有助于区分不同的物体，尤其是在拥挤的场景中。

   def process_video(
       source_weights_path: str,
       source_video_path: str,
       target_video_path: str,
       confidence_threshold: float = 0.3,
       iou_threshold: float = 0.7,
   ) -> None:
       model = YOLO(source_weights_path)
   	# 追踪者，用于跟随每一个检测到的物体并标记
       tracker = sv.ByteTrack()
       box_annotator = sv.BoundingBoxAnnotator()
       label_annotator = sv.LabelAnnotator()
       frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
       video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
   
       with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
           # 遍历每一帧的数据
           for frame in tqdm(frame_generator, total=video_info.total_frames):
               # 加载模型
               results = model(
                   frame, verbose=False, conf=confidence_threshold, iou=iou_threshold
               )[0]
   			# 从YOLOv8中加载数据结果
               detections = sv.Detections.from_ultralytics(results)
               # 更新被检测的物体的信息
               detections = tracker.update_with_detections(detections)
   
               annotated_frame = box_annotator.annotate(
                   scene=frame.copy(), detections=detections
               )
   			
               annotated_labeled_frame = label_annotator.annotate(
                   scene=annotated_frame, detections=detections
               )
   			# 重新绘制每一帧
               sink.write_frame(frame=annotated_labeled_frame)
   
   

3. 完成代码

   import argparse
   
   from tqdm import tqdm
   from ultralytics import YOLO
   
   import supervision as sv
   import os
   os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
   
   
   def process_video(
       source_weights_path: str,
       source_video_path: str,
       target_video_path: str,
       confidence_threshold: float = 0.3,
       iou_threshold: float = 0.7,
   ) -> None:
       model = YOLO(source_weights_path)
   
       tracker = sv.ByteTrack()
       box_annotator = sv.BoundingBoxAnnotator()
       label_annotator = sv.LabelAnnotator()
       frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
       video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
   
       with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
           for frame in tqdm(frame_generator, total=video_info.total_frames):
               results = model(
                   frame, verbose=False, conf=confidence_threshold, iou=iou_threshold
               )[0]
               detections = sv.Detections.from_ultralytics(results)
               detections = tracker.update_with_detections(detections)
   
               annotated_frame = box_annotator.annotate(
                   scene=frame.copy(), detections=detections
               )
   
               annotated_labeled_frame = label_annotator.annotate(
                   scene=annotated_frame, detections=detections
               )
   
               sink.write_frame(frame=annotated_labeled_frame)
   
   
   if __name__ == "__main__":
       parser = argparse.ArgumentParser(
           description="Video Processing with YOLO and ByteTrack"
       )
       parser.add_argument(
           "--source_weights_path",
           required=True,
           help="Path to the source weights file",
           type=str,
       )
       parser.add_argument(
           "--source_video_path",
           required=True,
           help="Path to the source video file",
           type=str,
       )
       parser.add_argument(
           "--target_video_path",
           required=True,
           help="Path to the target video file (output)",
           type=str,
       )
       parser.add_argument(
           "--confidence_threshold",
           default=0.3,
           help="Confidence threshold for the model",
           type=float,
       )
       parser.add_argument(
           "--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
       )
   
       args = parser.parse_args()
   
       process_video(
           source_weights_path=args.source_weights_path,
           source_video_path=args.source_video_path,
           target_video_path=args.target_video_path,
           confidence_threshold=args.confidence_threshold,
           iou_threshold=args.iou_threshold,
       )
   
   

4. 执行命令

   python ultralytics_example.py --source_weights_path yolov8s.pt --source_video_path input.mp4 --target_video_path tracking_result.mp4
   

3【参考】

https://github.com/roboflow/supervision?tab=readme-ov-file

后续更新，目标分割，速度检测，颜色识别，以及imutils工具包的使用