實例分割 模型訓練 InstanceSegmentation by detectron2

detectron2_repo_path = "" #detectron2 path 

img_path = r"coco/images/train"
json_path = r"coco/annotations/instances_train.json"

register_coco_instances("mydata", {}, json_path, img_path)
mydata_metadata = MetadataCatalog.get("mydata")
dataset_dicts = DatasetCatalog.get("mydata")

# test data
val_img_path = r"coco/images/val2017"
val_json_path = r"/coco/annotations/instances_val2017.json"
register_coco_instances("testdata", {}, val_json_path, val_img_path)
mydata_metadata = MetadataCatalog.get("testdata")
dataset_dicts = DatasetCatalog.get("testdata")

#註冊自己的數據集

cfg = get_cfg()
cfg.merge_from_file(
            os.path.join(detectron2_repo_path, "configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
            )
# github有提供不同骨幹：FPN,C4,DC5，
cfg.DATASETS.TRAIN = ("mydata",)
cfg.DATASETS.TEST = ("")  # no metrics implemented for this dataset
cfg.DATALOADER.NUM_WORKERS = 0
#cfg.MODEL.WEIGHTS = r"./model/model_final_f10217.pkl"  # 從 model zoo初始化
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.SOLVER.BASE_LR = 0.02
cfg.SOLVER.MAX_ITER = (300)  # 300 次迭代訓練

cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4 # 4 object type

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)

trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()  # 開始訓練 要打開!

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.8   # 設置此模型的測試閾值
cfg.DATASETS.TEST = ("data", )
predictor = DefaultPredictor(cfg)