detectron2:使用API函数训练自己的coco格式的数据集
程序员文章站
2022-03-25 18:07:56
...
0. load 基本的依赖库
# 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
1. 注册自己coco格式的数据集:
按照 Detectron2 自定义数据集教程,将自己的数据集注册到 Detectron2。
from detectron2.data.datasets import register_coco_instances
register_coco_instances('self_coco_train', {},
'/home/workspace/datasets/cocodatasets_2017/annotations/instances_train2017.json',
'/home/workspace/datasets/cocodatasets_2017/train2017')
register_coco_instances('self_coco_val', {},
'/home/workspace/datasets/cocodatasets_2017/annotations/instances_val2017.json',
'/home/workspace/datasets/cocodatasets_2017/val2017')
获取数据的元数据
coco_val_metadata = MetadataCatalog.get("self_coco_val")
dataset_dicts = DatasetCatalog.get("self_coco_val")
coco_val_metadata
原始数据可视化
import random
for d in random.sample(dataset_dicts, 3):
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])
2.模型训练
加载基本的配置
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
from detectron2 import model_zoo
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_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-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 80 # coco datasets
cfg
训练与保存模型
import os
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
3.预测
加载已经训练好的模型
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set the testing threshold for this model
cfg.DATASETS.TEST = ("self_coco_val", )
predictor = DefaultPredictor(cfg)
预测结果可视化
from detectron2.utils.visualizer import ColorMode
for d in random.sample(dataset_dicts, 3):
im = cv2.imread(d["file_name"])
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=coco_val_metadata,
scale=0.8,
instance_mode=ColorMode.IMAGE_BW # remove the colors of unsegmented pixels
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(v.get_image()[:, :, ::-1])
上一篇: 浅谈Service中实现弹出对话框的坑
推荐阅读
-
在C#下使用TensorFlow.NET训练自己的数据集
-
C#使用TensorFlow.NET训练自己的数据集的方法
-
python 使用Yolact训练自己的数据集
-
在Linux下使用Pytorch运行yolov3训练自己的数据集初体验
-
C#使用TensorFlow.NET训练自己的数据集的方法
-
在C#下使用TensorFlow.NET训练自己的数据集
-
自己制作机器学习训练和测试使用的二进制数据集(C++)
-
使用keras-retinanet训练自己的数据集
-
PaddleDetection——使用自己制作的VOC数据集进行模型训练的避坑指南
-
detectron2:使用API函数训练自己的coco格式的数据集