| |
| import logging |
|
|
| from detectron2.modeling import ROI_HEADS_REGISTRY, StandardROIHeads |
|
|
|
|
| @ROI_HEADS_REGISTRY.register() |
| class PointRendROIHeads(StandardROIHeads): |
| """ |
| Identical to StandardROIHeads, except for some weights conversion code to |
| handle old models. |
| """ |
|
|
| _version = 2 |
|
|
| def _load_from_state_dict( |
| self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs |
| ): |
| version = local_metadata.get("version", None) |
| if version is None or version < 2: |
| logger = logging.getLogger(__name__) |
| logger.warning( |
| "Weight format of PointRend models have changed! " |
| "Please upgrade your models. Applying automatic conversion now ..." |
| ) |
| for k in list(state_dict.keys()): |
| newk = k |
| if k.startswith(prefix + "mask_point_head"): |
| newk = k.replace(prefix + "mask_point_head", prefix + "mask_head.point_head") |
| if k.startswith(prefix + "mask_coarse_head"): |
| newk = k.replace(prefix + "mask_coarse_head", prefix + "mask_head.coarse_head") |
| if newk != k: |
| state_dict[newk] = state_dict[k] |
| del state_dict[k] |
|
|
| @classmethod |
| def _init_mask_head(cls, cfg, input_shape): |
| if cfg.MODEL.MASK_ON and cfg.MODEL.ROI_MASK_HEAD.NAME != "PointRendMaskHead": |
| logger = logging.getLogger(__name__) |
| logger.warning( |
| "Config of PointRend models have changed! " |
| "Please upgrade your models. Applying automatic conversion now ..." |
| ) |
| assert cfg.MODEL.ROI_MASK_HEAD.NAME == "CoarseMaskHead" |
| cfg.defrost() |
| cfg.MODEL.ROI_MASK_HEAD.NAME = "PointRendMaskHead" |
| cfg.MODEL.ROI_MASK_HEAD.POOLER_TYPE = "" |
| cfg.freeze() |
| return super()._init_mask_head(cfg, input_shape) |
|
|
