使用 Detectron II 构建您的第一个物体检测器

写于：2021 年 5 月

大家好，我决定为想要使用detectron2 尝试他们的第一个对象检测的人们制作笔记本教程。

我是 PyTorch 的忠实粉丝，这已经不是什么秘密了，我喜欢这个框架，并且因为 fastai 之类的库而更加爱上了它。让我们跳过那部分，只谈谈detectron2。

所以detectron2是Facebook的人工智能研究（FAIR）平台，用于对象检测和分割。 Detectron2 实际上是 Detectron 或 Dectectron1 的重写，您可能喜欢这样称呼它。 Detectron1 是用 caffe2 编写的，但 Detectron2 完全是 pytorch。

Detectron2 可以执行的不仅仅是在检测到的对象上绘制边界框，它可以定位和检测人类的关键点，以及预测姿势和标记身体部位。 Detectron 2 还执行所谓的“全景分割”，屏幕上的每个像素都被标记。我想停在那里，随时查看 Detectron2 文档更多。

本笔记本将指导您了解如何使用自定义数据集构建您的第一个对象检测器。在这里，我使用了检测器来检测前置总线。

让我们潜入吧！

安装您需要的所有依赖项，这肯定涉及 PyTorch 和 torchvision。如果您更喜欢使用最新版本的 PyTorch，请确保确认它与 detectron2 完全兼容。

安装依赖项：（使用 cu101 因为 colab 有 CUDA 10.1）

!pip install -U torch1.5 torchvision0.6 -f https://download.pytorch.org/whl/cu101/torch_stable.html
!pip install cython pyyaml==5.1
!pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'
进口火炬，火炬视觉
打印（火炬.版本，火炬.cuda.is_available（））
!gcc --版本
# opencv 预装在 colab 上

让我们自己安装Detectron 2。

安装检测器2：

!pip install detectron2==0.1.3 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu101/torch1.5/index.html

将 Detectron 2 作为库导入 — — 非常简单！

记得在此之前重启你的运行时，让你的安装生效

设置detectron2记录器

进口检测器2
从detectron2.utils.logger 导入setup_logger
setup_logger()

导入一些常用库

将 numpy 导入为 np
导入简历2
随机导入
从 google.colab.patches 导入 cv2_imshow

导入一些常用的detectron2工具

从detectron2导入model_zoo
从detectron2.engine 导入DefaultPredictor
从detectron2.config 导入get_cfg
从detectron2.utils.visualizer 导入Visualizer
从detectron2.data 导入元数据目录
从detectron2.data.catalog 导入DatasetCatalog

添加我的谷歌驱动器。

在这里，我只是将我的谷歌驱动器导入 colab，因为那是我的数据集所在的位置。

挂载你的谷歌驱动器

从 google.colab 导入驱动器
drive.mount('/content/drive/') 安装在 /content/drive/

数据集

如前所述，我的数据集包含前置总线，这些图像具有 COCO 格式的注释。尽管我只有一个包含这些注释的文件，但我仍然想将其拆分为训练集和验证集，以便拥有训练集和验证集。为了拆分注释，我使用了一个名为 cocosplit 的项目中的一些技巧。寻找这里 cocosplit GitHub repo 的链接以查看源代码。

!pip install funcy #我们需要funcy 收藏功能
下载 https://files.pythonhosted.org/packages/44/52/5cf7401456a461e4b481650dfb8279bc000f31a011d0918904f86e755947/funcy-1.16-py2.py3-none-any.whl
安装收集的包：funcy
安装成功func-1.16 # 快速导入一些东西
导入json
导入参数解析
导入功能
从 sklearn.model_selection 导入 train_test_split

帮助执行拆分的主要功能。

def save_coco（文件、信息、许可证、图像、注释、类别）：
使用 open(file, 'wt', encoding='UTF-8') 作为 coco：
json.dump({ 'info': 信息, 'licenses': 许可证, 'images': 图像,
'annotations': 注释, 'categories': 类别}, coco, indent=2, sort_keys=True)

def filter_annotations（注释，图像）：
image_ids = funcy.lmap(lambda i: int(i['id']), images)
return funcy.lfilter(lambda a: int(a['image_id']) in image_ids, annotations)

定义 d_splitter():
使用 open(annotations_path, 'rt', encoding='UTF-8') 作为注释：
coco = json.load(注解)
信息 = 可可 ['信息']
许可证 = coco['许可证']
图像 = 可可 ['图像']
注释 = coco['注释']
类别=可可['类别']

number_of_images = len（图像）

images_with_annotations = funcy.lmap(lambda a: int(a['image_id']), annotations)

如果真实：
images = funcy.lremove(lambda i: i['id'] not in images_with_annotations, images)

x, y = train_test_split(图像, train_size=train_split_size)

save_coco(store_train_path, info, licenses, x, filter_annotations(annotations, x), categories)
save_coco(store_test_path, info, licenses, y, filter_annotations(annotations, y), categories)

print("在 {} 和 {} 中保存 {} 条目".format(len(x), store_train_path, len(y), store_test_path))

因此，为了结束这一切，我们必须设置这些参数。

annotations_path：主要注释文件所在的位置。
十进制的火车拆分大小。这是训练集的大小。
存储火车路径是您要在拆分后存储新火车集注释的位置。请注意，您必须指定要写入的文件而不是文件夹。
和上面一样，指定一个它可以创建的文件，这样它就可以在那里存储你的测试注释。

annotations_path = "/content/drive/MyDrive/machine_learning_stories/Omdena/Bus_face_LED_panel_174_images/annotations_bus_face_only/instances_default.json"
train_split_size = 0.9
store_train_path = "/content/train.json"
store_test_path = "/content/test.json"

立即拆分！

d_splitter() 在 /content/train.json 中保存了 156 个条目，在 /content/test.json 中保存了 18 个条目

现在我们有了训练集注释和验证注释，我们可以继续将数据集注册到detectron2。在这里，我注册了两个数据集，火车和验证。 register_coco_instances 函数接受您要用于数据集的实例名称、一个空字典、注释文件的路径，最后是图像所在的目录。

从detectron2.data.datasets 导入register_coco_instances
register_coco_instances("my_dataset_train", {},"/content/train.json", "/content/drive/MyDrive/machine_learning_stories/Omdena/Bus_face_LED_panel_174_images/images")
register_coco_instances("my_dataset_val", {}, "/content/test.json", "/content/drive/MyDrive/machine_learning_stories/Omdena/Bus_face_LED_panel_174_images/images")

呼！！我们已经成功加载了数据集，我们可以先看一下数据集。让我们快速可视化它。

可视化训练数据

my_dataset_train_metadata = MetadataCatalog.get("my_dataset_train")
dataset_dicts = DatasetCatalog.get("my_dataset_train")

对于 random.sample(dataset_dicts, 3) 中的 d：
img = cv2.imread(d["file_name"])
可视化器 = 可视化器（img[:, :, ::-1]，元数据=my_dataset_train_metadata，规模=0.5）
vis = Visualizer.draw_dataset_dict(d)
cv2_imshow(vis.get_image()[:, :, ::-1])

好的！是时候训练模型了。

所以，是的，detectron2 有许多很棒的最先进的预训练模型，您可以根据自己的任务进行选择。在本教程中，我使用 COCO 预训练的 faster_rcnn 来训练我的模型，我还训练了大约 1500 次迭代，大约需要 35 分钟

from .detectron2.tools.train_net 导入培训师

from detectron2.engine 导入 DefaultTrainer

在此处从 modelzoo 中选择：https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md#coco-object-detection-baselines

从detectron2.config 导入get_cfg
#from detectron2.evaluation.coco_evaluation import COCOEvaluator
进口我们
从detectron2.engine 导入DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ("my_dataset_val",)

cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml") # 让训练从模型动物园初始化
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.001

cfg.SOLVER.WARMUP_ITERS = 1000
cfg.SOLVER.MAX_ITER = 1500 #如果val mAP还在上升就上调，如果过拟合就下调
cfg.SOLVER.STEPS = (1000, 1500)
cfg.SOLVER.GAMMA = 0.05

cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 #你的班级数量 + 1

os.makedirs（cfg.OUTPUT_DIR，exist_ok=True）
培训师 = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
培训师.train()

让我们看看使用张量板工具的训练曲线和其他细节

查看 tensorboard 中的训练曲线：

%load_ext 张量板
%tensorboard --logdir 输出张量板扩展已加载。要重新加载它，请使用：
%reload_ext 张量板

在端口 6006 (pid 723) 上重用 TensorBoard，于 0:54:06 前开始。（使用 '!kill 723' 杀死它。）

<IPython.core.display.Javascript object>

推理

使用经过训练的模型进行推理非常容易，您必须选择一个阈值，因此如果您对这种模型置信度没问题，我们会在它们周围绘制边界框。

推理应该使用带有训练中使用的参数的配置

cfg 现在已经包含了我们之前设置的所有内容。我们对其进行了一些更改以进行推理：

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") # 我们刚刚训练的模型的路径
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8 # 设置自定义测试阈值
预测器 = DefaultPredictor(cfg)

查看预测

和哟！我们完成了，我决定对验证数据集进行推断，检查预测，结果很好。

从detectron2.utils.visualizer 导入ColorMode
my_dataset_val_metadata = MetadataCatalog.get("my_dataset_val")
dataset_dicts = DatasetCatalog.get("my_dataset_val")
对于 random.sample(dataset_dicts, 3) 中的 d：
im = cv2.imread(d["file_name"])
output = predictor(im) # 格式记录在 https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
v = Visualizer(im[:, :, ::-1],
元数据=my_dataset_val_metadata，
比例=0.5)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])