Arulkumar03/Wheat_HEAD_Detection_Counting_ComputerVision_Model

1# Copyright (c) Facebook, Inc. and its affiliates.2import logging3 4from detectron2.utils.file_io import PathHandler, PathManager5 6 7class ModelCatalog(object):8    """9    Store mappings from names to third-party models.10    """11 12    S3_C2_DETECTRON_PREFIX = "https://dl.fbaipublicfiles.com/detectron"13 14    # MSRA models have STRIDE_IN_1X1=True. False otherwise.15    # NOTE: all BN models here have fused BN into an affine layer.16    # As a result, you should only load them to a model with "FrozenBN".17    # Loading them to a model with regular BN or SyncBN is wrong.18    # Even when loaded to FrozenBN, it is still different from affine by an epsilon,19    # which should be negligible for training.20    # NOTE: all models here uses PIXEL_STD=[1,1,1]21    # NOTE: Most of the BN models here are no longer used. We use the22    # re-converted pre-trained models under detectron2 model zoo instead.23    C2_IMAGENET_MODELS = {24        "MSRA/R-50": "ImageNetPretrained/MSRA/R-50.pkl",25        "MSRA/R-101": "ImageNetPretrained/MSRA/R-101.pkl",26        "FAIR/R-50-GN": "ImageNetPretrained/47261647/R-50-GN.pkl",27        "FAIR/R-101-GN": "ImageNetPretrained/47592356/R-101-GN.pkl",28        "FAIR/X-101-32x8d": "ImageNetPretrained/20171220/X-101-32x8d.pkl",29        "FAIR/X-101-64x4d": "ImageNetPretrained/FBResNeXt/X-101-64x4d.pkl",30        "FAIR/X-152-32x8d-IN5k": "ImageNetPretrained/25093814/X-152-32x8d-IN5k.pkl",31    }32 33    C2_DETECTRON_PATH_FORMAT = (34        "{prefix}/{url}/output/train/{dataset}/{type}/model_final.pkl"  # noqa B95035    )36 37    C2_DATASET_COCO = "coco_2014_train%3Acoco_2014_valminusminival"38    C2_DATASET_COCO_KEYPOINTS = "keypoints_coco_2014_train%3Akeypoints_coco_2014_valminusminival"39 40    # format: {model_name} -> part of the url41    C2_DETECTRON_MODELS = {42        "35857197/e2e_faster_rcnn_R-50-C4_1x": "35857197/12_2017_baselines/e2e_faster_rcnn_R-50-C4_1x.yaml.01_33_49.iAX0mXvW",  # noqa B95043        "35857345/e2e_faster_rcnn_R-50-FPN_1x": "35857345/12_2017_baselines/e2e_faster_rcnn_R-50-FPN_1x.yaml.01_36_30.cUF7QR7I",  # noqa B95044        "35857890/e2e_faster_rcnn_R-101-FPN_1x": "35857890/12_2017_baselines/e2e_faster_rcnn_R-101-FPN_1x.yaml.01_38_50.sNxI7sX7",  # noqa B95045        "36761737/e2e_faster_rcnn_X-101-32x8d-FPN_1x": "36761737/12_2017_baselines/e2e_faster_rcnn_X-101-32x8d-FPN_1x.yaml.06_31_39.5MIHi1fZ",  # noqa B95046        "35858791/e2e_mask_rcnn_R-50-C4_1x": "35858791/12_2017_baselines/e2e_mask_rcnn_R-50-C4_1x.yaml.01_45_57.ZgkA7hPB",  # noqa B95047        "35858933/e2e_mask_rcnn_R-50-FPN_1x": "35858933/12_2017_baselines/e2e_mask_rcnn_R-50-FPN_1x.yaml.01_48_14.DzEQe4wC",  # noqa B95048        "35861795/e2e_mask_rcnn_R-101-FPN_1x": "35861795/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_1x.yaml.02_31_37.KqyEK4tT",  # noqa B95049        "36761843/e2e_mask_rcnn_X-101-32x8d-FPN_1x": "36761843/12_2017_baselines/e2e_mask_rcnn_X-101-32x8d-FPN_1x.yaml.06_35_59.RZotkLKI",  # noqa B95050        "48616381/e2e_mask_rcnn_R-50-FPN_2x_gn": "GN/48616381/04_2018_gn_baselines/e2e_mask_rcnn_R-50-FPN_2x_gn_0416.13_23_38.bTlTI97Q",  # noqa B95051        "37697547/e2e_keypoint_rcnn_R-50-FPN_1x": "37697547/12_2017_baselines/e2e_keypoint_rcnn_R-50-FPN_1x.yaml.08_42_54.kdzV35ao",  # noqa B95052        "35998355/rpn_R-50-C4_1x": "35998355/12_2017_baselines/rpn_R-50-C4_1x.yaml.08_00_43.njH5oD9L",  # noqa B95053        "35998814/rpn_R-50-FPN_1x": "35998814/12_2017_baselines/rpn_R-50-FPN_1x.yaml.08_06_03.Axg0r179",  # noqa B95054        "36225147/fast_R-50-FPN_1x": "36225147/12_2017_baselines/fast_rcnn_R-50-FPN_1x.yaml.08_39_09.L3obSdQ2",  # noqa B95055    }56 57    @staticmethod58    def get(name):59        if name.startswith("Caffe2Detectron/COCO"):60            return ModelCatalog._get_c2_detectron_baseline(name)61        if name.startswith("ImageNetPretrained/"):62            return ModelCatalog._get_c2_imagenet_pretrained(name)63        raise RuntimeError("model not present in the catalog: {}".format(name))64 65    @staticmethod66    def _get_c2_imagenet_pretrained(name):67        prefix = ModelCatalog.S3_C2_DETECTRON_PREFIX68        name = name[len("ImageNetPretrained/") :]69        name = ModelCatalog.C2_IMAGENET_MODELS[name]70        url = "/".join([prefix, name])71        return url72 73    @staticmethod74    def _get_c2_detectron_baseline(name):75        name = name[len("Caffe2Detectron/COCO/") :]76        url = ModelCatalog.C2_DETECTRON_MODELS[name]77        if "keypoint_rcnn" in name:78            dataset = ModelCatalog.C2_DATASET_COCO_KEYPOINTS79        else:80            dataset = ModelCatalog.C2_DATASET_COCO81 82        if "35998355/rpn_R-50-C4_1x" in name:83            # this one model is somehow different from others ..84            type = "rpn"85        else:86            type = "generalized_rcnn"87 88        # Detectron C2 models are stored in the structure defined in `C2_DETECTRON_PATH_FORMAT`.89        url = ModelCatalog.C2_DETECTRON_PATH_FORMAT.format(90            prefix=ModelCatalog.S3_C2_DETECTRON_PREFIX, url=url, type=type, dataset=dataset91        )92        return url93 94 95class ModelCatalogHandler(PathHandler):96    """97    Resolve URL like catalog://.98    """99 100    PREFIX = "catalog://"101 102    def _get_supported_prefixes(self):103        return [self.PREFIX]104 105    def _get_local_path(self, path, **kwargs):106        logger = logging.getLogger(__name__)107        catalog_path = ModelCatalog.get(path[len(self.PREFIX) :])108        logger.info("Catalog entry {} points to {}".format(path, catalog_path))109        return PathManager.get_local_path(catalog_path, **kwargs)110 111    def _open(self, path, mode="r", **kwargs):112        return PathManager.open(self._get_local_path(path), mode, **kwargs)113 114 115PathManager.register_handler(ModelCatalogHandler())116