utils.py

# Copyright 2020 Google LLC, University of Victoria, Czech Technical University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# The majority of this code is taken from https://github.com/vcg-uvic/image-matching-benchmark

import os
from copy import deepcopy

import cv2
import h5py
import matplotlib.pyplot as plt
import numpy as np


def load_h5(filename):
    """Loads dictionary from hdf5 file"""
    dict_to_load = {}
    try:
        with h5py.File(filename, "r") as f:
            keys = [key for key in f.keys()]
            for key in keys:
                dict_to_load[key] = f[key][()]
    except:
        print("Cannot find file {}".format(filename))
    return dict_to_load


def save_h5(dict_to_save, filename):
    """Saves dictionary to HDF5 file"""

    with h5py.File(filename, "w") as f:
        for key in dict_to_save:
            # if 'img' in key:
            #    continue
            # print (key)
            f.create_dataset(key, data=dict_to_save[key])


def normalize_keypoints(keypoints, K):
    """Normalize keypoints using the calibration data."""

    C_x = K[0, 2]
    C_y = K[1, 2]
    f_x = K[0, 0]
    f_y = K[1, 1]
    keypoints = (keypoints - np.array([[C_x, C_y]])) / np.array([[f_x, f_y]])

    return keypoints


def get_E_from_F(F, K1, K2):
    return np.matmul(np.matmul(K2.T, F), K1)


def quaternion_from_matrix(matrix, isprecise=False):
    """Return quaternion from rotation matrix.
    If isprecise is True, the input matrix is assumed to be a precise rotation
    matrix and a faster algorithm is used.
    >>> q = quaternion_from_matrix(numpy.identity(4), True)
    >>> numpy.allclose(q, [1, 0, 0, 0])
    True
    >>> q = quaternion_from_matrix(numpy.diag([1, -1, -1, 1]))
    >>> numpy.allclose(q, [0, 1, 0, 0]) or numpy.allclose(q, [0, -1, 0, 0])
    True
    >>> R = rotation_matrix(0.123, (1, 2, 3))
    >>> q = quaternion_from_matrix(R, True)
    >>> numpy.allclose(q, [0.9981095, 0.0164262, 0.0328524, 0.0492786])
    True
    >>> R = [[-0.545, 0.797, 0.260, 0], [0.733, 0.603, -0.313, 0],
    ...      [-0.407, 0.021, -0.913, 0], [0, 0, 0, 1]]
    >>> q = quaternion_from_matrix(R)
    >>> numpy.allclose(q, [0.19069, 0.43736, 0.87485, -0.083611])
    True
    >>> R = [[0.395, 0.362, 0.843, 0], [-0.626, 0.796, -0.056, 0],
    ...      [-0.677, -0.498, 0.529, 0], [0, 0, 0, 1]]
    >>> q = quaternion_from_matrix(R)
    >>> numpy.allclose(q, [0.82336615, -0.13610694, 0.46344705, -0.29792603])
    True
    >>> R = random_rotation_matrix()
    >>> q = quaternion_from_matrix(R)
    >>> is_same_transform(R, quaternion_matrix(q))
    True
    >>> R = euler_matrix(0.0, 0.0, numpy.pi/2.0)
    >>> numpy.allclose(quaternion_from_matrix(R, isprecise=False),
    ...                quaternion_from_matrix(R, isprecise=True))
    True
    """

    M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]
    if isprecise:
        q = np.empty((4,))
        t = np.trace(M)
        if t > M[3, 3]:
            q[0] = t
            q[3] = M[1, 0] - M[0, 1]
            q[2] = M[0, 2] - M[2, 0]
            q[1] = M[2, 1] - M[1, 2]
        else:
            i, j, k = 1, 2, 3
            if M[1, 1] > M[0, 0]:
                i, j, k = 2, 3, 1
            if M[2, 2] > M[i, i]:
                i, j, k = 3, 1, 2
            t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3]
            q[i] = t
            q[j] = M[i, j] + M[j, i]
            q[k] = M[k, i] + M[i, k]
            q[3] = M[k, j] - M[j, k]
        q *= 0.5 / math.sqrt(t * M[3, 3])
    else:
        m00 = M[0, 0]
        m01 = M[0, 1]
        m02 = M[0, 2]
        m10 = M[1, 0]
        m11 = M[1, 1]
        m12 = M[1, 2]
        m20 = M[2, 0]
        m21 = M[2, 1]
        m22 = M[2, 2]

        # symmetric matrix K
        K = np.array(
            [
                [m00 - m11 - m22, 0.0, 0.0, 0.0],
                [m01 + m10, m11 - m00 - m22, 0.0, 0.0],
                [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],
                [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22],
            ]
        )
        K /= 3.0

        # quaternion is eigenvector of K that corresponds to largest eigenvalue
        w, V = np.linalg.eigh(K)
        q = V[[3, 0, 1, 2], np.argmax(w)]

    if q[0] < 0.0:
        np.negative(q, q)

    return q


def evaluate_R_t(R_gt, t_gt, R, t, q_gt=None):
    t = t.flatten()
    t_gt = t_gt.flatten()

    eps = 1e-15

    if q_gt is None:
        q_gt = quaternion_from_matrix(R_gt)
    q = quaternion_from_matrix(R)
    q = q / (np.linalg.norm(q) + eps)
    q_gt = q_gt / (np.linalg.norm(q_gt) + eps)
    loss_q = np.maximum(eps, (1.0 - np.sum(q * q_gt) ** 2))
    err_q = np.arccos(1 - 2 * loss_q)

    t = t / (np.linalg.norm(t) + eps)
    t_gt = t_gt / (np.linalg.norm(t_gt) + eps)
    loss_t = np.maximum(eps, (1.0 - np.sum(t * t_gt) ** 2))
    err_t = np.arccos(np.sqrt(1 - loss_t))

    if np.sum(np.isnan(err_q)) or np.sum(np.isnan(err_t)):
        # This should never happen! Debug here
        print(R_gt, t_gt, R, t, q_gt)
        import IPython

        IPython.embed()

    return err_q, err_t


def eval_essential_matrix(p1n, p2n, E, dR, dt):
    if len(p1n) != len(p2n):
        raise RuntimeError("Size mismatch in the keypoint lists")

    if p1n.shape[0] < 5:
        return np.pi, np.pi / 2

    if E.size > 0:
        _, R, t, _ = cv2.recoverPose(E, p1n, p2n)
        try:
            err_q, err_t = evaluate_R_t(dR, dt, R, t)
        except:
            err_q = np.pi
            err_t = np.pi / 2

    else:
        err_q = np.pi
        err_t = np.pi / 2

    return err_q, err_t


def decolorize(img):
    return cv2.cvtColor(cv2.cvtColor(img, cv2.COLOR_RGB2GRAY), cv2.COLOR_GRAY2RGB)


def drawlines(img1, img2, lines, pts1, pts2):
    """img1 - image on which we draw the epilines for the points in img2
    lines - corresponding epilines"""
    r, c, ch = img1.shape
    img1o = deepcopy(img1)
    img2o = deepcopy(img2)
    for r, pt1, pt2 in zip(lines, pts1, pts2):
        color = tuple(np.random.randint(0, 255, 3).tolist())
        x0, y0 = map(int, [0, -r[2] / r[1]])
        x1, y1 = map(int, [c, -(r[2] + r[0] * c) / r[1]])
        img1o = cv2.line(img1o, (x0, y0), (x1, y1), color, 1)
        img1o = cv2.circle(img1o, tuple(pt1.squeeze().astype(np.int32)), 5, color, -1)
        img2o = cv2.circle(img2o, tuple(pt2.squeeze().astype(np.int32)), 5, color, -1)
    return img1o, img2o


def draw_everything(img1, img2, pts1_good, pts2_good, F_gt):
    lines1gt = cv2.computeCorrespondEpilines(pts2_good.reshape(-1, 1, 2), 2, F_gt)
    lines1gt = lines1gt.reshape(-1, 3)
    img5gt, img6gt = drawlines(img1, img2, lines1gt, pts1_good, pts2_good)
    # Find epilines corresponding to points in left image (first image) and
    # drawing its lines on right image
    lines2gt = cv2.computeCorrespondEpilines(pts1_good.reshape(-1, 1, 2), 1, F_gt)
    lines2gt = lines2gt.reshape(-1, 3)
    img3gt, img4gt = drawlines(img2, img1, lines2gt, pts2_good, pts1_good)
    plt.figure(figsize=(12, 12))
    plt.subplot(121)
    plt.imshow(img5gt)
    plt.subplot(122)
    plt.imshow(img3gt)
    return


def get_h_imgpair(key, dataset, split="val"):
    DIR = "homography"
    if dataset == "EVD":
        img1_fname = f"{DIR}/{dataset}/{split}/imgs/1/" + key.split("-")[0] + ".png"
        img2_fname = f"{DIR}/{dataset}/{split}/imgs/2/" + key.split("-")[0] + ".png"
    elif dataset == "HPatchesSeq":
        img1_fname = f"{DIR}/{dataset}/{split}/imgs/{key[:-4]}/1.ppm"
        img2_fname = f"{DIR}/{dataset}/{split}/imgs/{key[:-4]}/{key[-1]}.ppm"
    else:
        raise ValueError("Unknown dataset, try EVD or HPatchesSeq")
    img1 = cv2.cvtColor(cv2.imread(img1_fname), cv2.COLOR_BGR2RGB)
    img2 = cv2.cvtColor(cv2.imread(img2_fname), cv2.COLOR_BGR2RGB)
    return img1, img2


def get_h_imgpair2(key, DIR):
    if "EVD" in DIR:
        img1_fname = f"{DIR}/imgs/1/" + key.split("-")[0] + ".png"
        img2_fname = f"{DIR}/imgs/2/" + key.split("-")[0] + ".png"
    elif "HPatchesSeq" in DIR:
        img1_fname = f"{DIR}/imgs/{key[:-4]}/1.ppm"
        img2_fname = f"{DIR}/imgs/{key[:-4]}/{key[-1]}.ppm"
    else:
        raise ValueError("Unknown dataset, try EVD or HPatchesSeq")
    img1 = cv2.cvtColor(cv2.imread(img1_fname), cv2.COLOR_BGR2RGB)
    img2 = cv2.cvtColor(cv2.imread(img2_fname), cv2.COLOR_BGR2RGB)
    return img1, img2


def get_output_dir(problem: str, split: str, method: str, params: dict):
    problem = problem.lower()
    if problem not in ["e", "f", "h", "pnp"]:
        raise ValueError(f"{problem} is unknown problem. Try e, f, h, or pnp")
    param_string = ""
    sorted_keys = sorted([str(x) for x in params.keys()])
    for k in sorted_keys:
        param_string += f"_{k}-{str(params[k])}"
    return os.path.join("results", split, problem, method, param_string)


# from ACNe code
def compute_T_with_imagesize(w, h, f=None, ratio=1.0):
    cx = (w - 1.0) * 0.5
    cy = (h - 1.0) * 0.5
    mean = np.array([cx, cy])
    if f is not None:
        f = f
    else:
        f = max(w - 1.0, h - 1.0) * ratio

    scale = 1.0 / f
    T = np.zeros(
        (
            3,
            3,
        )
    )
    T[0, 0], T[1, 1], T[2, 2] = scale, scale, 1
    T[0, 2], T[1, 2] = -scale * mean[0], -scale * mean[1]

    return T.copy()