def predict()

in miscellaneous/distributed_tensorflow_mask_rcnn/container-serving-optimized/resources/predict.py [0:0]

• 30 lines of code
• 4 McCabe index (conditional complexity)

    def predict(cls, img=None, img_id=None, rpn=False, score_threshold=0.8, mask_threshold=0.5):

        (
            rpn_boxes,
            rpn_scores,
            all_scores,
            final_boxes,
            final_scores,
            final_labels,
            masks,
        ) = cls.predictor(np.expand_dims(img, axis=0), np.expand_dims(np.array(img.shape), axis=0))

        predictions = {"img_id": str(img_id)}

        annotations = []

        img_shape = (img.shape[0], img.shape[1])
        for box, mask, score, category_id in zip(final_boxes, masks, final_scores, final_labels):
            a = {}
            b = box.tolist()
            a["bbox"] = [int(b[0]), int(b[1]), int(b[2] - b[0]), int(b[3] - b[1])]

            if round(score, 1) >= score_threshold:
                a["category_id"] = int(category_id)
                a["category_name"] = cfg.DATA.CLASS_NAMES[int(category_id)]
                b_mask = cls.get_binary_mask(img_shape, box, mask, threshold=mask_threshold)
                rle = cls.binary_mask_to_rle(b_mask)
                a["segmentation"] = rle
                annotations.append(a)

        predictions["annotations"] = annotations
        if rpn:
            predictions["rpn_boxes"] = rpn_boxes.tolist()
            predictions["rpn_scores"] = rpn_scores.tolist()
            predictions["all_scores"] = all_scores.tolist()

        return predictions