Script: homography estimation.py

#!/usr/bin/env python3

"""  Copyright (C) 2020 RidgeRun, LLC (http://www.ridgerun.com)
 All Rights Reserved.

 The contents of this software are proprietary and confidential to RidgeRun,
 LLC.  No part of this program may be photocopied, reproduced or translated
 into another programming language without prior written consent of
 RidgeRun, LLC.  The user is free to modify the source code after obtaining
 a software license from RidgeRun.  All source code changes must be provided
 back to RidgeRun without any encumbrance. """

""" Tool for estimating the homography matrix """


import argparse
import json
import sys

import cv2
import numpy as np

HOMOGRAPHY_DIMENSION = 3
MIN_MATCHES = 4
KERNEL_SIZE = 5
CHANNELS = 3
DEFAULT_REPROJ_ERROR = 4.0
DEFAULT_RATIO = 0.75
DEFAULT_SIGMA = 0
DEFAULT_FOV = 70
DEFAULT_OVERLAP = 15
DEFAULT_CROP = 0
DEFAULT_UNDISTORT = True


def drawMatches(imageA, imageB, keypointsA, keypointsB, matches, status):
    '''
   Returns an image with the matches found.

           Parameters:
                   imageA (np.array): Image A
                   imageB (np.array): Image B
                   keypointsA (np.array): Keypoints for image A
                   keypointsB (np.array): Keypoints for image B
                   matches (np.array): List of found matches between image
                    A and image B
                   status (int): status of homography estimation

           Returns:
                   Image (np.array)
   '''

    # initialize the output visualization image
    (heightA, widthA) = imageA.shape[:2]
    (heightB, widthB) = imageB.shape[:2]
    vis = np.zeros(
        (max(
            heightA,
            heightB),
            widthA +
            widthB,
            CHANNELS),
        dtype="uint8")
    vis[0:heightA, 0:widthA] = imageA
    vis[0:heightB, widthA:] = imageB

    for ((trainIdx, queryIdx), s) in zip(matches, status):
        # only process the match if the keypoint was successfully
        # matched
        if s == 1:
            # draw the match
            ptA = (int(keypointsA[queryIdx][0]), int(keypointsA[queryIdx][1]))
            ptB = (int(keypointsB[trainIdx][0]) +
                   widthA, int(keypointsB[trainIdx][1]))
            cv2.line(vis, ptA, ptB, (0, 255, 0), 1)

    return vis


def detectAndDescribe(gray, mask):
    '''
    Returns the keyponts and the correspoinding descriptors for the
    gray image.

            Parameters:
                    gray (np.array): Input image in grayscale.
                    mask (np.array): Mask to apply before the feature
                        extraction.

            Returns:
                    Tuple with the keypoins and their descriptors.
    '''

    # detect and extract features from the image
    detector = cv2.xfeatures2d.SIFT_create()
    (kps, features) = detector.detectAndCompute(gray, mask)
    kps = np.float32([kp.pt for kp in kps])

    return (kps, features)


def matchKeypoints(keypointsA, keypointsB, featuresA, featuresB,
                   ratio, reprojThresh):
    '''
    Function that searches the correspondiencies for the
    given descriptors (featuresA and featuresB) and with those
    correspondencies, the homography matrix is calculated.

            Parameters:
                    keypointsA (np.array): Keypoints of image A.
                    keypointsB (np.array): Keypoints of image B.
                    featuresA (np.array): Descriptors of image A
                    featuresB (np.array): Descriptors of image B
                    ratio (float): Max distance between a possible
                        correspondence.
                    reprojThresh (float): Reprojection error of the
                        homography estimation

            Returns:
                    Tuple with the resulting matches, the homography H
                    and the estimation status.
    '''

    # compute the raw matches and initialize the list of actual matches
    matcher = cv2.DescriptorMatcher_create("BruteForce")
    rawMatches = matcher.knnMatch(featuresA, featuresB, 2)
    matches = []

    # loop over the raw matches
    for m in rawMatches:
        # ensure the distance is within a certain ratio of each
        # other (i.e. Lowe's ratio test)
        if len(m) == 2 and m[0].distance < m[1].distance * ratio:
            matches.append((m[0].trainIdx, m[0].queryIdx))

    print("matches: " + str(len(matches)))
    if len(matches) >= MIN_MATCHES:
        # construct the two sets of points
        ptsA = np.float32([keypointsA[i] for (_, i) in matches])
        ptsB = np.float32([keypointsB[i] for (i, _) in matches])

        # compute the homography between the two sets of points
        (H, status) = cv2.findHomography(ptsA, ptsB, cv2.RANSAC,
                                         reprojThresh)

        # return the matches along with the homograpy matrix
        # and status of each matched point
        return (matches, H, status)

    # otherwise, no homograpy could be computed
    return None


def imageEnhancement(img, sigma):
    '''
    Returns an image with some enhancement applied. Remove the
    noise of the input image img using a Gaussian filter and
    then convert it to grayscale.

            Parameters:
                    img (np.array): Input image img
                    sigma (float): Sigma value of the Gaussian filter.

            Returns:
                    Image
    '''

    # Remove noise
    img = cv2.GaussianBlur(img, (KERNEL_SIZE, KERNEL_SIZE), sigma)

    # Convert to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    return gray


def removeDistortion(img, cameraMatrix, distortionParameters):
    '''
    Returns an image with the distortion provoked by the camera lens removed.

            Parameters:
                    img (np.array): Image img.
                    cameraMatrix (np.array): Camera matrix with a dimension
                        of (3x3).
                    distortionParameters (np.array): Vector with the
                        distortion coefficients.
            Returns:
                    Undistorted image.
    '''

    # Read an example image and acquire its size
    h, w = img.shape[:2]

    # Generate new camera matrix from parameters
    newCameraMatrix, roi = cv2.getOptimalNewCameraMatrix(
        cameraMatrix, distortionParameters, (w, h), 0)

    # Generate look-up tables for remapping the camera image
    mapX, mapY = cv2.initUndistortRectifyMap(
        cameraMatrix, distortionParameters, None, newCameraMatrix, (w, h), 5)

    # Remap the original image to a new image
    newImg = cv2.remap(img, mapX, mapY, cv2.INTER_LINEAR)
    return newImg


def parseJSON(filename):
    '''
    Returns the algorithm variables read from a JSON configuration file.

            Parameters:
                    filename (string): Path to the JSON configuration file.

            Returns:
                    Tuple the read variables: (K, D, reprojError,
                    matchRatio, sigma, overlap, crop, fov, undistort)
    '''

    with open(filename) as json_file:
        data = json.load(json_file)

        # load config variables
        if data['reprojError'] is not None:
            reprojError = float(data['reprojError'])
        else:
            reprojError = DEFAULT_REPROJ_ERROR
            print("Key reprojError not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_REPROJ_ERROR))

        if data['matchRatio'] is not None:
            matchRatio = float(data['matchRatio'])
        else:
            matchRatio = DEFAULT_RATIO
            print("Key matchRatio not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_RATIO))

        if data['sigma'] is not None:
            sigma = float(data['sigma'])
        else:
            sigma = DEFAULT_SIGMA
            print("Key sigma not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_SIGMA))

        if data['fov'] is not None:
            fov = float(data['fov'])
        else:
            fov = DEFAULT_FOV
            print("Key fov not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_FOV))

        if data['overlap'] is not None:
            overlap = float(data['overlap'])
        else:
            overlap = DEFAULT_OVERLAP
            print("Key overlap not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_OVERLAP))

        if data['crop'] is not None:
            crop = float(data['crop'])
        else:
            crop = DEFAULT_CROP
            print("Key crop not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_CROP))

        if data['undistort'] is not None:
            undistort = bool(data['undistort'])
        else:
            undistort = DEFAULT_UNDISTORT
            print("Key undistort not found in JSON configuration file. "
                  "Setting {} value".format(DEFAULT_UNDISTORT))

        # load camera matrix and distortion coeficients
        K = None
        D = None
        if undistort:
            if data['cameraMatrix'] is None:
                return (
                    K,
                    D,
                    reprojError,
                    matchRatio,
                    sigma,
                    overlap,
                    crop,
                    fov,
                    False)

            K = np.zeros(HOMOGRAPHY_DIMENSION * HOMOGRAPHY_DIMENSION)
            for i, k in enumerate(data['cameraMatrix']):
                K[i] = k
            K = K.reshape(HOMOGRAPHY_DIMENSION, HOMOGRAPHY_DIMENSION)

            if data['distortionParameters'] is None:
                return (
                    K,
                    D,
                    reprojError,
                    matchRatio,
                    sigma,
                    overlap,
                    crop,
                    fov,
                    False)

            D = np.zeros([1, len(data['distortionParameters'])])
            for i, d in enumerate(data['distortionParameters']):
                D[0, i] = d

        return (K, D, reprojError, matchRatio,
                sigma, overlap, crop, fov, undistort)


def getHomography(left, right, config):
    '''
    Returns the homography matrix along with the processed left and right
    images.

            Parameters:
                    left (np.array): Left input image.
                    right (np.array): Right input image.
                    config (string): Path to the JSON configuration file.

            Returns:
                    Tuple with: (left, right, homography H)
    '''

    (cameraMatrix, distortionParameters, reprojThresh, ratio, sigma, overlap,
     crop, fov, undistort) = parseJSON(config)
    height, width = left.shape[:2]

    # crop
    cropImg = int((crop * left.shape[1]) / fov)
    if cropImg > 0:
        left = left[:, :-1 * cropImg]
        left = cv2.resize(left, (width, height),
                          interpolation=cv2.INTER_CUBIC)
        right = right[:, cropImg:]
        right = cv2.resize(right, (width, height),
                           interpolation=cv2.INTER_CUBIC)

    # remove undistort
    if undistort:
        left = removeDistortion(left, cameraMatrix, distortionParameters)
        right = removeDistortion(right, cameraMatrix, distortionParameters)

    # Enhance images
    enhancedLeft = imageEnhancement(left, sigma)
    enhancedRight = imageEnhancement(right, sigma)

    # Mask
    overlapImg = int((overlap * left.shape[1]) / fov)
    leftMask = np.zeros(enhancedLeft.shape, dtype=np.uint8)
    leftMask[0:, (-1 * overlapImg):] = 1
    rightMask = np.zeros(enhancedRight.shape, dtype=np.uint8)
    rightMask[0:, 0:overlapImg] = 1

    # Extract keypoints
    (kpsA, featuresA) = detectAndDescribe(enhancedRight, rightMask)
    (kpsB, featuresB) = detectAndDescribe(enhancedLeft, leftMask)

    # Match features
    (matches, H, status) = matchKeypoints(
        kpsA, kpsB, featuresA, featuresB, ratio, reprojThresh)

    vis = drawMatches(right, left, kpsA, kpsB, matches, status)
    cv2.imwrite('vis.jpg', vis)
    print("Created vis.jpg image")

    return (left, right, H)


def left_fixed(config, targetImage, originalImage, homographyScale):
    '''
    Performs the homography estimation betwesn two images, leaving the left one
    fixed and transforming the right one to align them.

            Parameters:
                    config (string): Path to the JSON configuration file.
                    targetImage (string): Path to the target iamge.
                    originalImage (string): Path to the original image.
                    homographyScale (float): Scale factor for the generated
                        homography.

            Returns:
                    No return value
    '''

    # Load images
    target = cv2.imread(targetImage)
    original = cv2.imread(originalImage)

    (left, right, H) = getHomography(target, original, config)

    # Apply homography
    result = cv2.warpPerspective(
        right, H, (left.shape[1] + right.shape[1], right.shape[0]))
    result[0:left.shape[0], 0:left.shape[1]] = left
    cv2.imwrite('result.jpg', result)
    print("Created result.jpg image")

    # Scale homography
    if homographyScale > 0:
        Hmap = cv2.reg_MapProjec(H)
        Hmap.scale(homographyScale)
        H = Hmap.getProjTr()

    print(
        """RC_HOMOGRAPHY=\"{{\\"h00\\":{},\\"h01\\":{}, \\"h02\\":{}, \\"h10\\":{}"""
        """, \\"h11\\":{}, \\"h12\\":{}, \\"h20\\":{}, \\"h21\\":{}, \\"h22\\":{}}}\"""".format(
            H[0][0],
            H[0][1],
            H[0][2],
            H[1][0],
            H[1][1],
            H[1][2],
            H[2][0],
            H[2][1],
            H[2][2]))


def right_fixed(config, targetImage, originalImage, homographyScale):
    '''
    Performs the homography estimation between two images, leaving the right one
    fixed and transforming the left one to align them.

            Parameters:
                    config (string): Path to the JSON configuration file.
                    targetImage (string): Path to the target iamge.
                    originalImage (string): Path to the original image.
                    homographyScale (float): Scale factor for the generated
                        homography.

            Returns:
                    No return value
    '''

    # Load images
    target = cv2.imread(targetImage)
    original = cv2.imread(originalImage)

    (left, right, H) = getHomography(original, target, config)

    # X translation shift
    H = np.linalg.inv(H)
    H[0][-1] += original.shape[1]

    # Apply homography
    result = cv2.warpPerspective(
        left, H, (right.shape[1] + left.shape[1], left.shape[0]))
    result[0:right.shape[0], right.shape[1]:] = right
    cv2.imwrite('result.jpg', result)
    print("Created result.jpg image")

    # Scale homography
    if homographyScale > 0:
        Hmap = cv2.reg_MapProjec(H)
        H[0][-1] -= original.shape[1]
        Hmap.scale(homographyScale)
        H = Hmap.getProjTr()

    print(
        """LC_HOMOGRAPHY=\"{{\\"h00\\":{},\\"h01\\":{}, \\"h02\\":{}, \\"h10\\":{}"""
        """, \\"h11\\":{}, \\"h12\\":{}, \\"h20\\":{}, \\"h21\\":{}, \\"h22\\":{}}}\"""".format(
            H[0][0],
            H[0][1],
            H[0][2],
            H[1][0],
            H[1][1],
            H[1][2],
            H[2][0],
            H[2][1],
            H[2][2]))


def cmdline(argv):
    '''
    Function that parses the command line options before executing the algorithm

            Parameters:
                    config (list): List of command line arguments.

            Returns:
                    No return value
    '''

    prog = argv[0]
    parser = argparse.ArgumentParser(
        prog=prog,
        description='Tool for estimating the homography between two images.',
        epilog='Type "%s <command> -h" for more information.' %
        prog)

    subparsers = parser.add_subparsers(dest='command')
    subparsers.required = True

    def add_command(cmd, desc, example=None):
        epilog = 'Example: %s %s' % (
            prog, example) if example is not None else None
        return subparsers.add_parser(
            cmd, description=desc, help=desc, epilog=epilog)

    p = add_command(
        'left_fixed',
        'Estimation of homography between two images, with the left one fixed.')

    p.add_argument(
        '--config',
        help='Path of the configuration file',
        default='')
    p.add_argument(
        '--targetImage',
        help='Path of the target image',
        default='')
    p.add_argument(
        '--originalImage',
        help='Path of the original image',
        default='')
    p.add_argument(
        '--homographyScale',
        help='Scale factor of the homography. For example if you go from 1920x1080 '
        'in the estimation to 640x360 in the processing the scale factor should be 1/3',
        type=float,
        default=0)

    p = add_command(
        'right_fixed',
        'Estimation of homography between two images, with the right one fixed.')

    p.add_argument(
        '--config',
        help='Path of the configuration file',
        default='')
    p.add_argument(
        '--targetImage',
        help='Path of the target image',
        default='')
    p.add_argument(
        '--originalImage',
        help='Path of the original image',
        default='')
    p.add_argument(
        '--homographyScale',
        help='Scale factor of the homography. For example if you go from 1920x1080 '
        'in the estimation to 640x360 in the processing the scale factor should be 1/3',
        type=float,
        default=0)

    args = parser.parse_args(argv[1:] if len(argv) > 1 else ['-h'])
    func = globals()[args.command]
    del args.command
    func(**vars(args))


if __name__ == "__main__":
    cmdline(sys.argv)