triangulation_utils.py

import numpy as np
from itertools import combinations
import cv2

def project_points(keypoints, RT, einsum=None):
    homo = np.concatenate([keypoints[..., :3], np.ones_like(keypoints[..., :1])], axis=-1)
    if einsum is None:
        if len(homo.shape) == 2 and len(RT.shape) == 3:
            kpts2d = np.einsum('vab,kb->vka', RT, homo)
        elif len(homo.shape) == 2 and len(RT.shape) == 4:
            kpts2d = np.einsum('vkab,kb->vka', RT, homo)
        else:
            import ipdb; ipdb.set_trace()
    else:
        kpts2d = np.einsum(einsum, RT, homo)
    kpts2d[..., :2] /= kpts2d[..., 2:]
    return kpts2d

def project_and_distance(kpts3d, RT, kpts2d):
    kpts_proj = project_points(kpts3d, RT)
    # 1. distance between input and projection
    conf = (kpts3d[None, :, -1] > 0) * (kpts2d[:, :, -1] > 0)
    dist = np.linalg.norm(kpts_proj[..., :2] - kpts2d[..., :2], axis=-1) * conf
    return dist, conf

def batch_triangulate(keypoints_, Pall, min_view=2):
    """ triangulate the keypoints of whole body

    Args:
        keypoints_ (nViews, nJoints, 3): 2D detections
        Pall (nViews, 3, 4) | (nViews, nJoints, 3, 4): projection matrix of each view
        min_view (int, optional): min view for visible points. Defaults to 2.

    Returns:
        keypoints3d: (nJoints, 4)
    """
    # keypoints: (nViews, nJoints, 3)
    # Pall: (nViews, 3, 4)
    # A: (nJoints, nViewsx2, 4), x: (nJoints, 4, 1); b: (nJoints, nViewsx2, 1)
    v = (keypoints_[:, :, -1]>0).sum(axis=0)
    valid_joint = np.where(v >= min_view)[0]
    keypoints = keypoints_[:, valid_joint]
    conf3d = keypoints[:, :, -1].sum(axis=0)/v[valid_joint]
    # P2: P矩阵的最后一行：(1, nViews, 1, 4)
    if len(Pall.shape) == 3:
        P0 = Pall[None, :, 0, :]
        P1 = Pall[None, :, 1, :]
        P2 = Pall[None, :, 2, :]
    else:
        P0 = Pall[:, :, 0, :].swapaxes(0, 1)
        P1 = Pall[:, :, 1, :].swapaxes(0, 1)
        P2 = Pall[:, :, 2, :].swapaxes(0, 1)
    # uP2: x坐标乘上P2: (nJoints, nViews, 1, 4)
    uP2 = keypoints[:, :, 0].T[:, :, None] * P2
    vP2 = keypoints[:, :, 1].T[:, :, None] * P2
    conf = keypoints[:, :, 2].T[:, :, None]
    Au = conf * (uP2 - P0)
    Av = conf * (vP2 - P1)
    A = np.hstack([Au, Av])
    u, s, v = np.linalg.svd(A)
    X = v[:, -1, :]
    X = X / X[:, 3:]
    # out: (nJoints, 4) 
    result = np.zeros((keypoints_.shape[1], 4))
    result[valid_joint, :3] = X[:, :3]
    result[valid_joint, 3] = conf3d #* (conf[..., 0].sum(axis=-1)>min_view)
    return result

def iterative_triangulate(kpts2d, RT, previous=None,
    min_conf=0.1, min_view=3, min_joints=3, dist_max=0.05, dist_vel=0.05,
    thres_outlier_view=0.4, thres_outlier_joint=0.4, debug=False):
    kpts2d = kpts2d.copy()
    conf = kpts2d[..., -1]
    kpts2d[conf<min_conf] = 0.
    if debug:
        log('[triangulate] kpts2d: {}'.format(kpts2d.shape))
    # TODO: consider large motion
    if previous is not None:
        dist, conf = project_and_distance(previous, RT, kpts2d)
        nottrack = (dist > dist_vel) & conf
        if nottrack.sum() > 0:
            kpts2d[nottrack] = 0.
            if debug:
                log('[triangulate] Remove with track {}'.format(np.where(nottrack)))
    while True:
        # 0. triangulate and project
        kpts3d = batch_triangulate(kpts2d, RT, min_view=min_view)
        dist, conf = project_and_distance(kpts3d, RT, kpts2d)
        # 2. find the outlier
        vv, jj = np.where(dist > dist_max)
        if vv.shape[0] < 1:
            if debug:
                log('[triangulate] Not found outlier, break')
            break
        ratio_outlier_view = (dist>dist_max).sum(axis=1)/(1e-5 + conf.sum(axis=1))
        ratio_outlier_joint = (dist>dist_max).sum(axis=0)/(1e-5 + conf.sum(axis=0))
        # 3. find the totally wrong detections
        out_view = np.where(ratio_outlier_view > thres_outlier_view)[0]
        out_joint = np.where(ratio_outlier_joint > thres_outlier_joint)[0]
        if len(out_view) > 1:
            dist_view = dist.sum(axis=1)/(1e-5 + conf.sum(axis=1))
            out_view = out_view.tolist()
            out_view.sort(key=lambda x:-dist_view[x])
            if debug: mywarn('[triangulate] Remove outlier view: {}'.format(ratio_outlier_view))
        if remove_outview(kpts2d, out_view, debug): continue
        if remove_outjoint(kpts2d, RT, out_joint, dist_max, debug=debug): continue
        if debug:
            log('[triangulate] Directly remove {}, {}'.format(vv, jj))
        kpts2d[vv, jj, -1] = 0.
    if debug:
        log('[triangulate] finally {} valid points'.format((kpts3d[..., -1]>0).sum()))
    if (kpts3d[..., -1]>0).sum() < min_joints:
        kpts3d[..., -1] = 0.
        kpts2d[..., -1] = 0.
        return kpts3d, kpts2d
    return kpts3d, kpts2d

class Undistort:
    distortMap = {}
    @classmethod
    def image(cls, frame, K, dist, sub=None, interp=cv2.INTER_NEAREST):
        if sub is None:
            return cv2.undistort(frame, K, dist, None)
        else:
            if sub not in cls.distortMap.keys():
                h,  w = frame.shape[:2]
                mapx, mapy = cv2.initUndistortRectifyMap(K, dist, None, K, (w,h), 5)
                cls.distortMap[sub] = (mapx, mapy)
            mapx, mapy = cls.distortMap[sub]
            img = cv2.remap(frame, mapx, mapy, interp)
            return img

    @staticmethod
    def points(keypoints, K, dist):
        # keypoints: (N, 3)
        assert len(keypoints.shape) == 2, keypoints.shape
        kpts = keypoints[:, None, :2]
        kpts = np.ascontiguousarray(kpts)
        kpts = cv2.undistortPoints(kpts, K, dist, P=K)
        keypoints = np.hstack([kpts[:, 0], keypoints[:, 2:]])
        return keypoints
    
    @staticmethod
    def bbox(bbox, K, dist):
        keypoints = np.array([[bbox[0], bbox[1], 1], [bbox[2], bbox[3], 1]])
        kpts = Undistort.points(keypoints, K, dist)
        bbox = np.array([kpts[0, 0], kpts[0, 1], kpts[1, 0], kpts[1, 1], bbox[4]])
        return bbox

class SimpleTriangulate:
    def __init__(self, mode):
        self.mode = mode
    
    @staticmethod
    def undistort(points, cameras):
        nViews = len(points)
        pelvis_undis = []
        for nv in range(nViews):
            camera = {key:cameras[key][nv] for key in ['R', 'T', 'K', 'dist']}
            if points[nv].shape[0] > 0:
                pelvis = Undistort.points(points[nv], camera['K'], camera['dist'])
            else:
                pelvis = points[nv].copy()
            pelvis_undis.append(pelvis)
        return pelvis_undis

    def __call__(self, keypoints, cameras):
        '''
            keypoints: [nViews, nJoints, 3]
        
        output: 
            keypoints3d: (nJoints, 4)
        '''
        keypoints = self.undistort(keypoints, cameras)
        keypoints = np.stack(keypoints)
        if self.mode == 'naive':
            keypoints3d = batch_triangulate(keypoints, cameras['P'])
        else:
            keypoints3d, k2d = iterative_triangulate(keypoints, cameras['P'], dist_max=25)
        return {'keypoints3d': keypoints3d}