Ejemplo n.º 1
0
    def get_predictor(cls):
        ''' load trained model'''

        with cls.lock:
            # check if model is already loaded
            if cls.predictor:
                return cls.predictor

            os.environ['TENSORPACK_FP16'] = 'true'
        
            # create a mask r-cnn model
            mask_rcnn_model = ResNetFPNModel(True)

            try:
                model_dir = os.environ['SM_MODEL_DIR']
            except KeyError:
                model_dir = '/opt/ml/model'

            try:
                cls.pretrained_model = os.environ['PRETRAINED_MODEL']
            except KeyError:
                pass

            # file path to previoulsy trained mask r-cnn model
            latest_trained_model = ""
            model_search_path = os.path.join(model_dir, "model-*.index" )
            for model_file in glob.glob(model_search_path):
                if model_file > latest_trained_model:
                    latest_trained_model = model_file

            trained_model = latest_trained_model
            print(f'Using model: {trained_model}')

            # fixed resnet50 backbone weights
            cfg.BACKBONE.WEIGHTS = os.path.join(cls.pretrained_model)
            cfg.MODE_FPN = True
            cfg.MODE_MASK = True

            # calling detection dataset gets the number of coco categories 
            # and saves in the configuration
            DetectionDataset()
            finalize_configs(is_training=False)

            # Create an inference model
            # PredictConfig takes a model, input tensors and output tensors
            cls.predictor = OfflinePredictor(PredictConfig(
                model=mask_rcnn_model,
                session_init=get_model_loader(trained_model),
                input_names=['images', 'orig_image_dims'],
                output_names=[
                    'generate_{}_proposals_topk_per_image/boxes'.format('fpn' if cfg.MODE_FPN else 'rpn'),
                    'generate_{}_proposals_topk_per_image/scores'.format('fpn' if cfg.MODE_FPN else 'rpn'),
                    'fastrcnn_all_scores',
                    'output/boxes',
                    'output/scores',
                    'output/labels',
                    'output/masks'
                ]))
            return cls.predictor
Ejemplo n.º 2
0
    def get_predictor(cls):
        """load trained model"""

        with cls.lock:
            # check if model is already loaded
            if cls.predictor:
                return cls.predictor

            os.environ["TENSORPACK_FP16"] = "true"

            # create a mask r-cnn model
            mask_rcnn_model = ResNetFPNModel(True)

            try:
                model_dir = os.environ["SM_MODEL_DIR"]
            except KeyError:
                model_dir = "/opt/ml/model"
            try:
                resnet_arch = os.environ["RESNET_ARCH"]
            except KeyError:
                resnet_arch = "resnet50"

            # file path to previoulsy trained mask r-cnn model
            latest_trained_model = ""
            model_search_path = os.path.join(model_dir, "model-*.index")
            for model_file in glob.glob(model_search_path):
                if model_file > latest_trained_model:
                    latest_trained_model = model_file

            trained_model = latest_trained_model
            print(f"Using model: {trained_model}")

            # fixed resnet50 backbone weights
            cfg.MODE_FPN = True
            cfg.MODE_MASK = True
            if resnet_arch == "resnet101":
                cfg.BACKBONE.RESNET_NUM_BLOCKS = [3, 4, 23, 3]
            else:
                cfg.BACKBONE.RESNET_NUM_BLOCKS = [3, 4, 6, 3]

            cfg_prefix = "CONFIG__"
            for key, value in dict(os.environ).items():
                if key.startswith(cfg_prefix):
                    attr_name = key[len(cfg_prefix):]
                    attr_name = attr_name.replace("__", ".")
                    value = eval(value)
                    print(f"update config: {attr_name}={value}")
                    nested_var = cfg
                    attr_list = attr_name.split(".")
                    for attr in attr_list[0:-1]:
                        nested_var = getattr(nested_var, attr)
                    setattr(nested_var, attr_list[-1], value)

            # calling detection dataset gets the number of coco categories
            # and saves in the configuration
            DetectionDataset()
            finalize_configs(is_training=False)

            # Create an inference model
            # PredictConfig takes a model, input tensors and output tensors
            cls.predictor = OfflinePredictor(
                PredictConfig(
                    model=mask_rcnn_model,
                    session_init=get_model_loader(trained_model),
                    input_names=["images", "orig_image_dims"],
                    output_names=[
                        "generate_{}_proposals_topk_per_image/boxes".format(
                            "fpn" if cfg.MODE_FPN else "rpn"),
                        "generate_{}_proposals_topk_per_image/scores".format(
                            "fpn" if cfg.MODE_FPN else "rpn"),
                        "fastrcnn_all_scores",
                        "output/boxes",
                        "output/scores",
                        "output/labels",
                        "output/masks",
                    ],
                ))
            return cls.predictor
Ejemplo n.º 3
0
                        help="Log the full git diff",
                        action="store_false")

    #################################################################################################################

    if get_tf_version_tuple() < (1, 6):
        # https://github.com/tensorflow/tensorflow/issues/14657
        logger.warn(
            "TF<1.6 has a bug which may lead to crash in FasterRCNN if you're unlucky."
        )

    args = parser.parse_args()
    if args.config:
        cfg.update_args(args.config)

    MODEL = ResNetFPNModel(args.fp16)
    DetectionDataset(
    )  # initialize the config with information from our dataset

    if args.visualize or args.evaluate or args.predict:
        assert tf.test.is_gpu_available()
        assert args.load
        finalize_configs(is_training=False)

        if args.predict or args.visualize:
            cfg.TEST.RESULT_SCORE_THRESH = cfg.TEST.RESULT_SCORE_THRESH_VIS

        if args.visualize:
            do_visualize(MODEL, args.load)
        else:
            predcfg = PredictConfig(
    def get_predictor(cls):
        ''' load trained model'''

        with cls.lock:
            # check if model is already loaded
            if cls.predictor:
                return cls.predictor

            os.environ['TENSORPACK_FP16'] = 'true'

            # create a mask r-cnn model
            mask_rcnn_model = ResNetFPNModel(True)

            try:
                model_dir = os.environ['SM_MODEL_DIR']
            except KeyError:
                model_dir = '/opt/ml/model'

            try:
                cls.pretrained_model = os.environ['PRETRAINED_MODEL']
            except KeyError:
                pass

            try:
                div = int(eval(os.environ['divisor']))
            except KeyError:
                div = 1
                pass

            rpn_anchor_stride = int(16 / div)
            rpn_anchor_sizes = (int(32 / div), int(64 / div), int(128 / div),
                                int(256 / div), int(512 / div))

            try:
                rpn_anchor_stride = int(eval(os.environ['rpnanchor_stride']))
            except KeyError:
                pass

            try:
                nms_topk = int(eval(os.environ['NMS_TOPK']))
            except KeyError:
                nms_topk = 2
                pass

            try:
                nms_thresh = eval(os.environ['NMS_THRESH'])
            except KeyError:
                nms_thresh = 0.7
                pass

            try:
                results_per_img = eval(os.environ['res_perimg'])
            except KeyError:
                results_per_img = 400
                pass

            # file path to previoulsy trained mask r-cnn model
            latest_trained_model = ""
            model_search_path = os.path.join(model_dir, "model-*.index")
            for model_file in glob.glob(model_search_path):
                if model_file > latest_trained_model:
                    latest_trained_model = model_file

            trained_model = latest_trained_model
            print(f'Using model: {trained_model}')

            # fixed resnet50 backbone weights
            cfg.BACKBONE.WEIGHTS = os.path.join(cls.pretrained_model)
            cfg.MODE_FPN = True
            cfg.MODE_MASK = True
            cfg.RPN.ANCHOR_STRIDE = rpn_anchor_stride
            cfg.RPN.ANCHOR_SIZES = rpn_anchor_sizes
            cfg.RPN.TEST_PRE_NMS_TOPK = int(6000 * nms_topk)
            cfg.RPN.TEST_POST_NMS_TOPK = int(1000 * nms_topk)
            cfg.RPN.TEST_PER_LEVEL_NMS_TOPK = int(1000 * nms_topk)
            # testing -----------------------
            cfg.TEST.FRCNN_NMS_THRESH = nms_thresh
            cfg.TEST.RESULT_SCORE_THRESH = 0.05
            cfg.TEST.RESULT_SCORE_THRESH_VIS = 0.2  # only visualize confident results
            cfg.TEST.RESULTS_PER_IM = results_per_img

            # calling detection dataset gets the number of coco categories
            # and saves in the configuration
            DetectionDataset()
            finalize_configs(is_training=False)

            # Create an inference model
            # PredictConfig takes a model, input tensors and output tensors
            cls.predictor = OfflinePredictor(
                PredictConfig(
                    model=mask_rcnn_model,
                    session_init=get_model_loader(trained_model),
                    input_names=['images', 'orig_image_dims'],
                    output_names=[
                        'generate_{}_proposals_topk_per_image/boxes'.format(
                            'fpn' if cfg.MODE_FPN else 'rpn'),
                        'generate_{}_proposals_topk_per_image/scores'.format(
                            'fpn' if cfg.MODE_FPN else 'rpn'),
                        'fastrcnn_all_scores', 'output/boxes', 'output/scores',
                        'output/labels', 'output/masks'
                    ]))
            return cls.predictor