detectron2 训练自己定义好的coco数据集

本文主要介绍如何在colab上训练自己的coco数据集实现目标检测的功能！

使用coco数据集主要的原因时 代码中 有现成的注册数据集等格式，无需自己去定义一个新的函数

目录

1 coco数据集的存放

2 查看detectron2 的目录

3 基于colab的detectron2训练

3.1 查看设备切换路径、安装等一系列问题

3.2 安装好环境，尝试导包测试等

3.3 先加载数据，对数据进行注册

3.4 查看读取的图像的标注框

3.5 训练数据集

3.6 进行测试

3.6.1 读取训练的模型生产一个预测器

3.6.2 预测一张图片后的效果

3.6.3 做整体测试后的结果

3.7 测试推断时间

---

1 coco数据集的存放

annotations 主要是存放对应的json文件

train2017主要存放 训练时的图像

val2017主要存放  测试时的图像

2 查看detectron2 的目录

3 基于colab的detectron2训练

3.1 查看设备切换路径、安装等一系列问题

 此处无需再安装pytorch，默认使用最新版的1.9版本就行

!nvidia-smi
%cd /content/drive/MyDrive/

##将 detectron2 装载到自己的云盘
!git clone https://github.com/facebookresearch/detectron2.git  
    
##搭建好环境
%cd /content/drive/MyDrive/
!python -m pip install -e detectron2  

3.2 安装好环境，尝试导包测试等

尝试导入不同的包，若没有成功导入，则尝试使用exit()后再重新使用GPU就行了

此时可能需要重新执行以下语句 !python -m pip install -e detectron2

然后再尝试导“包”就能成功

# You may need to restart your runtime prior to this, to let your installation take effect
# Some basic setup
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import matplotlib.pyplot as plt
import numpy as np
import cv2
from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

3.3 先加载数据，对数据进行注册

from detectron2.data import MetadataCatalog
from detectron2.data.datasets import register_coco_instances
##第一个参数为自己注册的数据名称，第二个参数无需管，第三个参数为数据的地址
register_coco_instances('self_coco_train', {}, 
                        '/content/drive/MyDrive/coco/annotations/instances_train2017.json', 
                       '/content/drive/MyDrive/coco/train2017')
register_coco_instances('self_coco_val', {}, 
                        '/content/drive/MyDrive/coco/annotations/instances_val2017.json', 
                       '/content/drive/MyDrive/coco/val2017')

#查看一下自己的元数据
##运行后可以看到thing_classes=["自己的数据类别"]
coco_val_metadata = MetadataCatalog.get("self_coco_val")
dataset_dicts = DatasetCatalog.get("self_coco_val")
coco_val_metadata

结果：（直接使用coco数据集进行注册可以直接看到数据集的“类别”等具体信息）

3.4 查看读取的图像的标注框

import random

for d in random.sample(dataset_dicts, 1):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=coco_val_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    cv2_imshow(vis.get_image()[:, :, ::-1])

 结果：

3.5 训练数据集

from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2 import model_zoo
import os

cfg = get_cfg()
cfg.merge_from_file("/content/drive/MyDrive/detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  #配置的地方
cfg.DATASETS.TRAIN = ("self_coco_train",)                              #设置好训练的数据集，也就是刚刚自己注册的数据集
cfg.DATASETS.TEST = ("self_coco_val", )

cfg.DATALOADER.NUM_WORKERS = 2        
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")  #让训练初试化从对应的配置开始，这里也可以先下载
cfg.SOLVER.IMS_PER_BATCH = 2                                    #Batchsize
cfg.SOLVER.BASE_LR = 0.00025                                    #单个GPU应该设置的学习率
cfg.SOLVER.MAX_ITER = 300    # 迭代的次数   图片的数量/Batchsize等于一个epoch迭代的次数
cfg.TEST.EVAL_PERIOD= 100   # 每迭代几次就进行一次评估，但结果未显现？？？表示很疑问
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2  # 自己数据集的类别

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)    #传入output文件夹 存放训练好的权重等
trainer = DefaultTrainer(cfg)
#DefaultTrainer.test(evaluators='coco')
trainer.resume_or_load(resume=False)
trainer.train()

结果：

3.6 进行测试

3.6.1 读取训练的模型生产一个预测器

import os
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")  # 我们进行训练的 权值文件  存放处
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # 设置一个阈值
predictor = DefaultPredictor(cfg)

3.6.2 预测一张图片后的效果

from detectron2.utils.visualizer import ColorMode
for d in random.sample(dataset_dicts, 1):
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im,
                   metadata=coco_val_metadata,       
                   scale=0.9,
                   instance_mode =  ColorMode.IMAGE
        )
                   # remove the colors of unsegmented pixels
    print(outputs['instances'].pred_classes)
    print(outputs["instances"].pred_boxes)

    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    cv2_imshow(v.get_image())

因为我这里 训练的次数没达标 所以出来的效果不是很好

3.6.3 做整体测试后的结果

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader

evaluator = COCOEvaluator("self_coco_val", cfg, False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "self_coco_val")
inference_on_dataset(trainer.model, val_loader, evaluator)

结果：（因为这里训练的次数 没达标 所以效果 很差）

3.7 测试推断时间

import time

def speedTest(predictor,img):
    N = 20
    startT = time.time()
    for i in range(N):
        predictor(img)
    endT = time.time()
    print("run time:{:.3}s/per image,imageShape:{}".format((endT-startT)/N,img.shape))
speedTest(predictor,im)

 结果：