基于detectron2的目标检测API

1 运行环境

可以直接使用服务器上的sharePytorch环境

  1. source activate sharePytorch

自己配置需要根据https://github.com/facebookresearch/detectron2的安装步骤安装相关环境

2 下载

服务器上的代码路径为/home/chbls/PycharmProjects/detectron2Test，也直接从语雀下载object_detection.py

3 使用

调用方式

  1. detector = Detectron2()
  2. result = detector(path="street.jpg", is_folder=False, visualization=True)

调用时有以下三个参数

  1. path
    • 指定单个图片或文件夹的路径
    • 如果指定的是文件夹路径，则模型会识别文件夹下所有图片（后缀为png、jpg、jpeg）
  2. is_folder
    • 默认值为False
    • 当path指定的是文件夹路径时，需要指定该参数为True
  3. visualization
    • 默认值为False
    • 该参数指定是否可视化目标检测的结果

4 输出结果

目标检测接口会返回一个InstancesData列表，其中InstancesData封装了单张图片的目标检测结果

  • InstancesData定义

    1. class InstancesData:
    2.   """
    3.   封装了单张图片的目标检测结果
    4.   """
    6.   def __init__(self, boxes, classes, scores, features):
    7.       self.boxes = boxes  # 物体bbx
    8.       self.classes = classes  # 物体类别
    9.       self.scores = scores  # 置信度
    10.       self.features = features  # 物体的feature map
    12.   # 图片中物体的数量
    13.   @property
    14.   def instances_count(self):
    15.       return len(self.boxes)

5 测试

测试代码
分别测试检测单张图片和检测一个文件夹内的所有图片

  1. from object_detection import Detectron2
  3. object_detector = Detectron2()
  5. print("----test one image----")
  6. result = object_detector("street.jpg", is_folder=False, visualization=True)
  7. print(result[0])
  9. print("----test folder----")
  10. object_detector("/new_dev/kingston_a400/STTran/STTran/ag/frames/0A8CF.mp4",
  11.                 is_folder=True,
  12.                 visualization=True)

测试结果

  • 单张图片

-可视化
1.png
-数据

  1. ----test one image----
  2. instancesData
  3.  instances_count: 9
  4.  boxes: Boxes(tensor([[373.9778, 587.7932, 535.8943, 829.8896],
  5.         [196.1853, 489.5289, 352.8455, 776.3320],
  6.         [ 51.1151, 558.6667, 192.0185, 801.9966],
  7.         [165.3411, 565.9411, 302.1731, 788.3586],
  8.         [542.0828, 579.9363, 638.2467, 780.9273],
  9.         [448.1255, 540.5609, 640.0000, 797.8308],
  10.         [213.0738, 242.0820, 338.7947, 347.1292],
  11.         [437.5472, 898.0849, 459.5511, 912.9505],
  12.         [199.7362, 524.1202, 227.8205, 545.3105]], device='cuda:0')) 
  13.  classes: tensor([ 0,  0,  3,  3,  3,  3,  9, 14, 46], device='cuda:0')
  14.  scores: tensor([0.9970, 0.9832, 0.9694, 0.9493, 0.9482, 0.7045, 0.6820, 0.6048, 0.5721],
  15.        device='cuda:0')
  16.  features: tensor([[0.6648, 0.2306, 0.5742,  ..., 0.5689, 0.0000, 0.0000],
  17.         [0.8326, 0.1943, 0.8564,  ..., 0.0000, 0.0000, 0.0000],
  18.         [0.9263, 0.0272, 1.9297,  ..., 0.0000, 0.0000, 0.0000],
  19.         ...,
  20.         [0.3757, 0.0000, 0.0000,  ..., 0.0000, 0.0000, 0.0000],
  21.         [1.7139, 0.1994, 0.0000,  ..., 0.0000, 0.0000, 0.5677],
  22.         [0.0166, 0.0000, 0.0000,  ..., 0.0000, 0.1855, 0.0000]],
  23.        device='cuda:0')
  • 图片文件夹

-可视化2.png