本范例演示使用非常有名的目标检测框架detectron2 🤗🤗
在自己的数据集(balloon数据)上训练实例分割模型MaskRCNN的方法。
detectron2框架的设计有以下一些优点：
1，强大：提供了包括目标检测、实例分割、全景分割等非常广泛的视觉任务模型库。
2，灵活：可以通过注册机制自定义模块或模型结构，从而进行扩展和改进。
3，易用：通过list of dict格式定义自己的数据集, 简单好用。
公众号算法美食屋后台回复关键词: torchkeras，获取本文源代码和balloon数据集下载链接。
我们首先需要安装并导入detectron库~
!pip install 'git+https://github.com/facebookresearch/detectron2.git'
!pip install torchkerasimport numpy as np
import os, json, cv2, random
from PIL import Image
import torch
import detectron2
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
#from detectron2.utils.logger import setup_logger
#setup_logger()
def cv2_show(arr):
img = Image.fromarray(cv2.cvtColor(arr, cv2.COLOR_BGR2RGB))
return img0，预训练模型
from torchkeras import data
#下载测试图片
img = data.get_example_image('park.jpg')
img.save('park.jpg')cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
im = cv2.imread("park.jpg")
outputs = predictor(im)print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_show(out.get_image()[:, :, ::-1])
一，准备数据
detectron准备数据集，需要先注册。
如果是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_val.json","path/to/image/dir")非coco格式的数据，可以用以下步骤进行注册：
1，先将数据集整理成字典组成的列表形式
2，使用DatasetCatalog注册数据集
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]
record["file_name"] = filename
record["image_id"] = idx
record["height"] = height
record["width"] = width
annos = v["regions"]
objs = []
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["annotations"] = objs
dataset_dicts.append(record)
return dataset_dicts
try:
#DatasetCatalog.remove('balloon_train')
#DatasetCatalog.remove('balloon_val')
DatasetCatalog.register("balloon_train", lambda : get_balloon_dicts("./data/balloon/train"))
MetadataCatalog.get("balloon_train" ).set(thing_classes=["balloon"])
DatasetCatalog.register("balloon_val", lambda : get_balloon_dicts("./data/balloon/val"))
MetadataCatalog.get("balloon_val" ).set(thing_classes=["balloon"])
except Exception as err:
pass
balloon_metadata = MetadataCatalog.get("balloon_train")我们来可视化一下数据，看看是否正确。
dicts_train = DatasetCatalog.get('balloon_train') #get_balloon_dicts("./data/balloon/train")
dicts_val = DatasetCatalog.get('balloon_val') #get_balloon_dicts("./data/balloon/val")dic = dicts_train[3]
img = cv2.imread(dic["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(dic)
cv2_show(out.get_image()[:, :, ::-1])
二，定义模型
detectron2通过配置文件定义模型。可以查看 detectron2目录下的configs路径，有各种各样功能的模型配置文件可以使用。
包括：Detection(检测), InstanceSegmentation(实例分割), Keypoints(关键点检测), Panoptic(全景分割) 等各种类型
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # The "RoIHead batch size". 128 is faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon).
#model = detectron2.modeling.build_model(cfg) # This way do not load pretrained weights
predictor = DefaultPredictor(cfg)
model = predictor.model三，训练模型
以下代码使用detectron2原生的DefaultTrainer进行训练，比较简单。
但是这个DefaultTrainer灵活性一般，当你想在训练循环中加入自己想要的功能时比较麻烦，并且日志输出不够直观。
此外也没有earlystopping，不能够保存验证集上最优的权重。
from detectron2.engine import DefaultTrainer
cfg.DATASETS.TRAIN = ("balloon_train",)
cfg.DATASETS.TEST = ("balloon_val",)
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 4
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2 # This is the real "batch size" commonly known to deep learning people
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 600
cfg.SOLVER.STEPS = [] # do not decay learning rate
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()
下面使用我们的 梦中情炉 ~ torchkeras ~ 实现最优雅的训练循环~ 😋😋
ds_train = detectron2.data.DatasetFromList(dicts_train)
ds_val = detectron2.data.DatasetFromList(dicts_val)
mp = detectron2.data.DatasetMapper(cfg,is_train=True)batch_size = 16
dl_train = detectron2.data.build_detection_train_loader(ds_train,
mapper=mp,total_batch_size=batch_size,num_workers=2)
dl_train.size = len(ds_train)//batch_size
dl_val = detectron2.data.build_detection_train_loader(ds_val,
mapper=mp,total_batch_size=1,num_workers=2)
dl_val.size = len(dicts_val)for batch in dl_val:
breakfrom torchkeras import KerasModel
from tqdm import tqdm
from detectron2.utils.events import EventStorage
class StepRunner:
def __init__(self, net, loss_fn, accelerator,
stage = "train", metrics_dict = None,
optimizer = None, lr_scheduler = None
):
self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage
self.optimizer,self.lr_scheduler = optimizer,lr_scheduler
self.accelerator = accelerator
if self.stage=='train':
self.net.train()
else:
self.net.train()
def __call__(self, batch):
#loss
with EventStorage() as event_storage:
loss_dict = self.net(batch)
loss = sum(loss_dict.values())
#backward()
if self.optimizer is not None and self.stage=="train":
self.accelerator.backward(loss)
self.optimizer.step()
if self.lr_scheduler is not None:
self.lr_scheduler.step()
self.optimizer.zero_grad()
all_loss = self.accelerator.gather(loss).sum()
#losses
step_losses = {self.stage+"_loss":all_loss.item()}
#metrics
step_metrics = {}
if self.stage=="train":
if self.optimizer is not None:
step_metrics['lr'] = self.optimizer.state_dict()['param_groups'][0]['lr']
else:
step_metrics['lr'] = 0.0
return step_losses,step_metrics
class EpochRunner:
def __init__(self,steprunner,quiet=False):
self.steprunner = steprunner
self.stage = steprunner.stage
self.accelerator = self.steprunner.accelerator
self.quiet = quiet
def __call__(self,dataloader):
try:
n = len(dataloader)
except Exception as err:
n = dataloader.size
loop = tqdm(enumerate(dataloader,start=1),
total =n,
file=sys.stdout,
disable=not self.accelerator.is_local_main_process or self.quiet,
ncols = 100
)
epoch_losses = {}
for step, batch in loop:
step_losses,step_metrics = self.steprunner(batch)
step_log = dict(step_losses,**step_metrics)
for k,v in step_losses.items():
epoch_losses[k] = epoch_losses.get(k,0.0)+v
if step<n:
loop.set_postfix(**step_log)
elif step==n:
epoch_metrics = step_metrics
epoch_metrics.update({self.stage+"_"+name:metric_fn.compute().item()
for name,metric_fn in self.steprunner.metrics_dict.items()})
epoch_losses = {k:v/step for k,v in epoch_losses.items()}
epoch_log = dict(epoch_losses,**epoch_metrics)
loop.set_postfix(**epoch_log)
for name,metric_fn in self.steprunner.metrics_dict.items():
metric_fn.reset()
else:
break
return epoch_log
KerasModel.StepRunner = StepRunner
KerasModel.EpochRunner = EpochRunnerparams = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.AdamW(params, lr=1e-4)
keras_model = KerasModel(model,
loss_fn = None,
metrics_dict=None,
optimizer= optimizer
)ckpt_path = 'checkpoint.pt'
keras_model.fit(train_data=dl_train,val_data=dl_val,
epochs=30,patience=10,
monitor='val_loss',
mode='min',
ckpt_path =ckpt_path,
plot=True
)
四，评估模型
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
evaluator = COCOEvaluator("balloon_val", output_dir="./output")
dl_val = build_detection_test_loader(cfg, "balloon_val")
print(inference_on_dataset(model, dl_val, evaluator))五，使用模型
from detectron2.engine import DefaultPredictor
#cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.WEIGHTS = ckpt_path
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7 # set a custom testing threshold
predictor = DefaultPredictor(cfg)from detectron2.utils.visualizer import ColorMode
im = cv2.imread(dicts_val[10]['file_name'])outputs = predictor(im) # format is documented at https://detectron2.readthedocs.io/tutorials/models.html#model-output-format
vis = Visualizer(im[:, :, ::-1],
metadata=balloon_metadata,
scale=0.5,
instance_mode=ColorMode.IMAGE_BW # remove the colors of unsegmented pixels. This option is only available for segmentation models
)out = vis.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_show(out.get_image()[:, :, ::-1])
公众号算法美食屋后台回复关键词: torchkeras，获取本文源代码和balloon数据集下载链接。
万水千山总是情，点个赞赞行不行？😋😋
