Detectron2解读全部文章链接：
Detectron新手教程
官方 Colab Notebook 上 Beginner Tutorial 部分阅读：
1. 使用预训练的Detectron2模型
下载一张图片，我们需要创建detectron2 config，随后根据config创建一个Default Predictor去进行单张图片推理。
cfg = get_cfg() # 获取 Default Config
# 根据 mask_rcnn_R_50_FPN_3x.yaml 的配置文件更新 config
cfg.merge_from_file(model_zoo.get_config_file("COCOInstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
# 推理的显示阈值设置为 0.5
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
# 根据配置文件选择对应的训练好的模型
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
# 创建 DefaultPredictor
predictor = DefaultPredictor(cfg)
# 推理
outputs = predictor(im)
获取网络输出结果：
print(outputs['instances'].pred_classes)
print(outputs['instances'].pred_boxes)
我们可以使用 Visualizer 类去可视化输出结果：
v = Visualizer(im[:,:,::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to('cpu'))
cv2_imshow(out.get_image()[:,:,::-1])
2. 在自定义数据集上训练
准备数据集
在这一部分，我们学习如何在自定义数据集上训练。我们以气球分割数据集为例，这个数据集只有一类：气球。我们使用在COCO数据集上预训练的模型。
准备好数据集后，我们尽量把数据集转换成COCO格式。对于COCO格式的数据集，我们可以直接用如下方法注册数据集：
from detectron2.data.datasets import register_coco_instances
register_coco_instances(
"my_dataset_train",
{},
"json_annotation_train.json",
"path/to/image/dir"
)
register_coco_instances(
"my_dataset_val",
{},
"json_annotation_train.json",
"path/to/image/dir"
)
如果我们的数据集是自己的格式，我们需要自定义一个函数来把自己格式的数据集输入detectron2。首先我们需要定义 get_balloon_dicts(img_dir) 方法来从数据集文件夹获取 dataset_dicts，之后通过DatasetCatalog.register方法把dataset_dicts注册到detectron2中。训练时使用 MetadataCatalog.get 来获取 Metadata。
from detectron2.structures import BoxMode
def get_balloon_dicts(img_dir):
json_file = os.path.join(img_dir, "via_region_data.json")
with open(json_file) as f:
imgs_anns = json.load(f)
dataset_dicts = []
for idx, v in enumerate(imgs_anns.values()):
record = {}
filename = os.path.join(img_dir, v["filename"])
height, width = cv2.imread(filename).shape[:2]
"""
对于每张图片，我们需要记录：
1.图片的文件名，2.图片的编号，3.图片的高，4.图片的宽
"""
record["file_name"] = filename
record["image_id"] = idx
record["height"] = height
record["width"] = width
annos = v["regions"]
objs = []
# 对于单张图片的每个标注（annotation）
for _, anno in annos.items():
assert not anno["region_attributes"]
anno = anno["shape_attributes"]
px = anno["all_points_x"]
py = anno["all_points_y"]
poly = [(x+0.5, y+0.5) for x, y in zip(px, py)]
poly = [p for x in poly for p in x]
# 创建物体的字典
obj = {
"bbox":[np.min(px), np.min(py), np.max(px), np.max(py)], # 物体轮廓同时可以转换为一个框
"bbox_mode": BoxMode.XYXY_ABS,
"segmentation": [poly],
"category_id": 0,
}
objs.append(obj)
# 把转换完形式的字典放进 record 中
record["annotations"] = objs
# 把单张图片的信息加入 dataset_dict
dataset_dicts.append(record)
return dataset_dicts
for d in ["train", "val"]:
DatasetCatalog.register("balloon_"+d, lambda d=d:get_ballon_dicts("balloon/"+d))
MetadataCatalog.get("balloon_"+d).set(things_classes=["balloon"])
ballon_metadata = MetadataCatalog.get("balloon_train")
这些工作做完之后，我们可以测试数据集是否注册正确：
dataset_dicts = get_ballon_dicts("balloon/train")
for d in random.sample(dataset_dicts, 3):
img = cv2.imread(d["filename"])
visualizer = Visualizer(img[:,:,::-1], metadata=ballon_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(d)
cv2_imshow(out.get_image()[:,:,::-1])
#cv2_imshow方法不是opencv中的，可以使用cv2.imshow来显示图片。
训练
准备好数据集之后，我们即可开始训练，作为例子，我们使用在COCO数据集上与训练的 R50-FPN Mask R-CNN 模型。
from detectron2.engine import DefaultTrainer
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
# "balloon_train"数据集已经被注册进去了，这里指定训练数据集
cfg.DATASETS.TRAIN = ("balloon_train",)
# 加载预训练模型
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
# 训练的一些基础参数
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 300
cfg.SOLVER.STEPS = []
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
#调用 trainer.train() 开始训练
trainer.train()
训练曲线可以在tensorboard上面可视化：
%load_ext tensorboard
%tensorboard --logdir output
推理和Evaluation
现在模型已经训练好，我们需要对没模型进行评测和单张图片推理，我们需要创建一个 predictor。推理时使用的config必须和训练时一致，我们只需要对config进行一些小改动：
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth") # 加载训练好的模型
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # 评分在0.7以上的框才可视化
predictor = DefaultPredictor(cfg)
之后我们随机从测试集选取几张图片：
from detectron2.utils.visualizer import ColorMode
dataset_dicts = get_balloon_dicts("balloon/val")
for d in random.sample(dataset_dicts, 3):
im = cv2.imread(d["file_name"])
outputs = predictor(im)
# 查看outputs的格式请浏览：https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
v = Visualizer(im[:, :, ::-1],
metadata=balloon_metadata,
scale=0.5,
instance_mode=ColorMode.IMAGE_BW
)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])
我们可以通过 COCOEvaluator 来评测数据集的 AP 等表现：
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
evaluator = COCOEvaluator("balloon_val", ("bbox", "segm"), False, output_dir="./output/")
val_loader = build_detection_test_loader(cfg, "balloon_val")
print(inference_on_dataset(trainer.model, val_loader, evaluator))
其它模型的训练同理。
扫一扫
1168
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
