Module narya.tracker.homography_estimator

Classes

class HomographyEstimator (pretrained=True, weights_homo=None, weights_keypoints=None, shape_in=512.0, shape_out=320.0)

Class the homography estimator. Given an image, computes the homography.

Arguments

pretrained: Boolean, if the homography models should be pretrained with our weights or not. weights_homo: Path to weight for the homography model weights_keypoints: Path to weight for the keypoints model shape_in: Shape of the input image shape_out: Shape of the ouput image Call arguments: input_img: np.array

Expand source code

class HomographyEstimator:
    """Class the homography estimator. Given an image, computes the homography.

    Arguments:
        pretrained: Boolean, if the homography models should be pretrained with our weights or not.
        weights_homo: Path to weight for the homography model
        weights_keypoints: Path to weight for the keypoints model
        shape_in: Shape of the input image
        shape_out: Shape of the ouput image
    Call arguments:
        input_img: np.array
    """

    def __init__(
        self,
        pretrained=True,
        weights_homo=None,
        weights_keypoints=None,
        shape_in=512.0,
        shape_out=320.0,
    ):

        self.homo_model = DeepHomoModel()
        self.keypoints_model = KeypointDetectorModel(
            backbone="efficientnetb3", num_classes=29, input_shape=(320, 320),
        )

        if pretrained == True:

            checkpoints_homo = tf.keras.utils.get_file(
                HOMO_NAME, HOMO_PATH, HOMO_TOTAR,
            )
            checkpoints_keypoints = tf.keras.utils.get_file(
                KEYPOINTS_NAME, KEYPOINTS_PATH, KEYPOINTS_TOTAR,
            )

            self.homo_model.load_weights(checkpoints_homo)
            self.keypoints_model.load_weights(checkpoints_keypoints)

        elif weights_homo is not None and weights_keypoints is not None:

            self.homo_model.load_weights(weights_homo)
            self.keypoints_model.load_weights(weights_keypoints)

        self.shape_in = shape_in
        self.shape_out = shape_out
        self.ratio = shape_out / shape_in

    def __call__(self, input_img):
        pr_mask = self.keypoints_model(input_img)
        src, dst = _points_from_mask(pr_mask[0])
        if len(src) > 3:
            pred_homo = get_perspective_transform(dst, src)
            method = "cv"
        else:
            corners = self.homo_model(input_img)
            pred_homo = compute_homography(corners)[0]
            method = "torch"
        return pred_homo, method

    def get_field_coordinates(self, bbox, pred_homo, method):
        """Computes the warped coordinates of a bounding box
        Arguments:
            bbox: np.array of shape (4,).
            pred_homo: np.array of shape (3,3)
            method: string in {'cv','torch'}, to precise how the homography was predicted
        Returns:
            dst: np.array, the warped coordinates of the bouding box
        Raises:
        """
        x_1 = int(bbox[0])
        y_1 = int(bbox[1])
        x_2 = int(bbox[2])
        y_2 = int(bbox[3])
        x = (x_1 + x_2) / 2.0 * self.ratio
        y = max(y_1, y_2) * self.ratio
        if method == "cv":
            pts = np.array([float(x), float(y)])
            dst = warp_point(pts, np.linalg.inv(pred_homo), method="cv")
        else:
            pts = np.array([int(x), int(y)])
            dst = warp_point(pts, np.linalg.inv(pred_homo), method="torch")
        return dst

Methods

def get_field_coordinates(self, bbox, pred_homo, method)

Computes the warped coordinates of a bounding box

Arguments

bbox: np.array of shape (4,). pred_homo: np.array of shape (3,3) method: string in {'cv','torch'}, to precise how the homography was predicted

Returns

dst

np.array, the warped coordinates of the bouding box

Raises:

Expand source code

def get_field_coordinates(self, bbox, pred_homo, method):
    """Computes the warped coordinates of a bounding box
    Arguments:
        bbox: np.array of shape (4,).
        pred_homo: np.array of shape (3,3)
        method: string in {'cv','torch'}, to precise how the homography was predicted
    Returns:
        dst: np.array, the warped coordinates of the bouding box
    Raises:
    """
    x_1 = int(bbox[0])
    y_1 = int(bbox[1])
    x_2 = int(bbox[2])
    y_2 = int(bbox[3])
    x = (x_1 + x_2) / 2.0 * self.ratio
    y = max(y_1, y_2) * self.ratio
    if method == "cv":
        pts = np.array([float(x), float(y)])
        dst = warp_point(pts, np.linalg.inv(pred_homo), method="cv")
    else:
        pts = np.array([int(x), int(y)])
        dst = warp_point(pts, np.linalg.inv(pred_homo), method="torch")
    return dst