EasyMocap/easymocap/mytools/triangulator.py

import numpy as np
import cv2
from easymocap.datasets.base import crop_image

from easymocap.estimator.wrapper_base import bbox_from_keypoints
from easymocap.mytools.vis_base import merge, plot_keypoints_auto
from .debug_utils import log, mywarn, myerror

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): 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)
    P0 = Pall[None, :, 0, :]
    P1 = Pall[None, :, 1, :]
    P2 = Pall[None, :, 2, :]
    # 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 project_points(keypoints, RT, einsum='vab,kb->vka'):
    homo = np.concatenate([keypoints[..., :3], np.ones_like(keypoints[..., :1])], axis=-1)
    kpts2d = np.einsum(einsum, RT, homo)
    kpts2d[..., :2] /= kpts2d[..., 2:]
    return kpts2d

def make_Cnk(n, k):
    import itertools
    res = {}
    for n_ in range(3, n+1):
        n_0 = [i for i in range(n_)]
        for k_ in range(2, k+1):
            res[(n_, k_)] = list(map(list, itertools.combinations(n_0, k_)))
    return res

MAX_VIEWS = 30
Cnk = make_Cnk(MAX_VIEWS, 3)

def robust_triangulate_point(kpts2d, Pall, dist_max, min_v = 3):
    nV = kpts2d.shape[0]
    if len(kpts2d) < min_v:# 重建失败
        return [], None
    # min_v = max(2, nV//2)
    # 1. choose the combination of min_v
    index_ = Cnk[(len(kpts2d), min(min_v, len(kpts2d)))]
    # 2. proposals: store the reconstruction points of each proposal
    proposals = np.zeros((len(index_), 4))
    weight_self = np.zeros((nV, len(index_)))
    for i, index in enumerate(index_):
        weight_self[index, i] = 100.
        point = batch_triangulate(kpts2d[index, :], Pall[index], min_view=min_v)
        proposals[i] = point
    # 3. project the proposals to each view
    #    and calculate the reprojection error
    # (nViews, nProposals, 4)
    kpts_repro = project_points(proposals, Pall)
    conf = (proposals[None, :, -1] > 0) * (kpts2d[..., -1] > 0)
    # err: (nViews, nProposals)
    err = np.linalg.norm(kpts_repro[..., :2] - kpts2d[..., :2], axis=-1) * conf
    valid = 1. - err/dist_max
    valid[valid<0] = 0
    # consider the weight of different view
    # TODO:naive weight:
    conf = kpts2d[..., -1]
    weight = conf
    # (valid > 0)*weight_self 一项用于强制要求使用到的两个视角都需要被用到
    # 增加一项使用的视角数的加成
    weight_sum = (weight * valid).sum(axis=0) + ((valid > 0)*weight_self).sum(axis=0) - min_v * 100
    if weight_sum.max() < 0:# 重建失败
        return [], None
    best = weight_sum.argmax()
    if (err[index_[best], best] > dist_max).any():
        return [], None
    # 对于选出来的proposal，寻找其大于0的其他视角
    point = proposals[best]
    best_add = np.where(valid[:, best])[0].tolist()
    index = list(index_[best])
    best_add.sort(key=lambda x:-weight[x])
    for add in best_add:
        if add in index:
            continue
        index.append(add)
        point = batch_triangulate(kpts2d[index, :], Pall[index], min_view=min_v)
        kpts_repro = project_points(point, Pall[index])
        err = np.linalg.norm(kpts_repro[..., :2] - kpts2d[index, ..., :2], axis=-1)
        if (err > dist_max).any():
            index.remove(add)
            break
    return index, point

def remove_outview(kpts2d, out_view, debug):
    if len(out_view) == 0:
        return False
    outv = out_view[0]
    if debug:
        mywarn('[triangulate] remove outview: {} from {}'.format(outv, out_view))
    kpts2d[outv] = 0.
    return True

def remove_outjoint(kpts2d, Pall, out_joint, dist_max, min_view=3, debug=False):
    if len(out_joint) == 0:
        return False
    if debug:
        mywarn('[triangulate] remove outjoint: {}'.format(out_joint))
    for nj in out_joint:
        valid = np.where(kpts2d[:, nj, -1] > 0)[0]
        if len(valid) < min_view:
            # if less than 3 visible view, set these unvisible
            kpts2d[:, nj, -1] = 0
            continue
        if len(valid) > MAX_VIEWS:
            # only select max points
            conf = -kpts2d[:, nj, -1]
            valid = conf.argsort()[:MAX_VIEWS]
        index_j, point = robust_triangulate_point(kpts2d[valid, nj:nj+1], Pall[valid], dist_max=dist_max, min_v=3)
        index_j = valid[index_j]
        # print('select {} for joint {}'.format(index_j, nj))
        set0 = np.zeros(kpts2d.shape[0])
        set0[index_j] = 1.
        kpts2d[:, nj, -1] *= set0
    return True

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 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 BaseTriangulator:
    def __init__(self, config, debug, keys) -> None:
        self.config = config
        self.debug = debug
        self.keys = keys

    def project_and_check(self, kpts3d, kpts2d, RT):
        kpts_proj = project_points(kpts3d, RT)
        conf = (kpts3d[None, :, -1] > 0) * (kpts2d[:, :, -1] > 0)
        dist = np.linalg.norm(kpts_proj[..., :2] - kpts2d[..., :2], axis=-1) * conf
        return conf, dist
    
    def triangulate_with_results(self, pid, data, results):
        new = {'id': pid}
        for key in self.keys:
            key3d = key.replace('2d', '3d')
            if len(results) == 0:
                kpts3d, kpts2d = iterative_triangulate(data[key + '_unproj'], data['RT'],
                    debug=self.debug, **self.config[key])
            else:
                if len(results) == 1:
                    previous = results[-1][key3d] # TODO: mean previous frame
                elif len(results) >= 2:
                    # TODO: mean previous velocity
                    previous0 = results[-2][key3d] # TODO: mean previous frame
                    previous1 = results[-1][key3d] # TODO: mean previous frame
                    vel = (previous1[:, :3] - previous0[:, :3])*((previous0[:, -1:]>0)&(previous0[:, -1:]>0))
                    previous = previous1.copy()
                    previous[:, :3] += vel
                kpts3d, kpts2d = iterative_triangulate(data[key + '_unproj'], data['RT'],
                    debug=self.debug, previous=previous, **self.config[key])
                vel = np.linalg.norm(kpts3d[:, :3] - previous[:, :3], axis=-1)
            new[key] = np.concatenate([data[key+'_distort'][..., :-1], kpts2d[..., -1:]], axis=-1)
            new[key3d] = kpts3d
        return new

class SimpleTriangulator(BaseTriangulator):
    def __init__(self, keys, debug, config,
        pid=0) -> None:
        super().__init__(config, debug, keys)
        self.results = []
        self.infos = []
        self.dim_name = ['_joints', '_views']
        self.pid = pid

    def __call__(self, data, results=None):
        info = {}
        if results is None:
            results = self.results
        new = {'id': self.pid}
        for key in self.keys:
            if key not in data.keys(): continue
            key3d = key.replace('2d', '3d')
            if self.debug:
                log('[triangulate] {}'.format(key))
            if len(results) == 0:
                kpts3d, kpts2d = iterative_triangulate(data[key + '_unproj'], data['RT'],
                    debug=self.debug, **self.config[key])
            else:
                if len(results) == 1:
                    previous = results[-1][key3d] # TODO: mean previous frame
                elif len(results) >= 2:
                    # TODO: mean previous velocity
                    previous0 = results[-2][key3d] # TODO: mean previous frame
                    previous1 = results[-1][key3d] # TODO: mean previous frame
                    vel = (previous1[:, :3] - previous0[:, :3])*((previous0[:, -1:]>0)&(previous0[:, -1:]>0))
                    previous = previous1.copy()
                    previous[:, :3] += vel
                kpts3d, kpts2d = iterative_triangulate(data[key + '_unproj'], data['RT'],
                    debug=self.debug, previous=previous, **self.config[key])
                vel = np.linalg.norm(kpts3d[:, :3] - previous[:, :3], axis=-1)
            new[key] = np.concatenate([data[key+'_distort'][..., :-1], kpts2d[..., -1:]], axis=-1)
            new[key.replace('2d', '3d')] = kpts3d
            if self.debug:
                conf, dist = self.project_and_check(kpts3d, kpts2d, data['RT'])
                for dim in [0, 1]:
                    info_dim = {
                        'valid': conf.sum(axis=dim),
                        'dist': 10000*dist.sum(axis=dim)/(1e-5 + conf.sum(axis=dim)),
                    }
                    info[key+self.dim_name[dim]] = info_dim
                info[key+'_joints']['valid3d'] = kpts3d[:, -1] >0
        results.append(new)
        self.infos.append(info)
        return [new]

    def report(self):
        if not self.debug:
            return 0
        from .debug_utils import print_table
        for key in self.infos[0].keys():
            metrics = list(self.infos[0][key].keys())
            values = [np.mean(np.stack([info[key][metric] for info in self.infos]), axis=0) for metric in metrics]
            metrics = [key] + metrics
            values = [[i for i in range(values[0].shape[0])]] + values
            print_table(metrics, values)

class SimpleTriangulatorMulti(SimpleTriangulator):
    def __init__(self, pids, **cfg) -> None:
        super().__init__(**cfg)
        self.results = {}
    
    def __call__(self, data, results=None):
        res_now = []
        for ipid, pid in enumerate(data['pid']):
            if pid not in self.results.keys():
                self.results[pid] = []
            data_ = {'RT': data['RT']}
            for key in self.keys:
                data_[key+'_distort'] = data[key+'_distort'][:, ipid]
                data_[key+'_unproj'] = data[key+'_unproj'][:, ipid]
                data_[key] = data[key][:, ipid]
            res = self.triangulate_with_results(pid, data_, self.results[pid])
            self.results[pid].append(res)
            res_now.append(res)
        return res_now

def skew_op(x):
    skew_op = lambda x: np.array([[0, -x[2], x[1]], [x[2], 0, -x[0]], [-x[1], x[0], 0]])
    res = np.zeros((3, 3), dtype=x.dtype)
    # 0, -z, y
    res[0, 1] = -x[2, 0]
    res[0, 2] =  x[1, 0]
    # z, 0, -x
    res[1, 0] =  x[2, 0]
    res[1, 2] = -x[0, 0]
    # -y, x, 0
    res[2, 0] = -x[1, 0]
    res[2, 1] =  x[0, 0]
    return res

def fundamental_op(K0, K1, R_0, T_0, R_1, T_1):
    invK0 = np.linalg.inv(K0)
    return invK0.T @ (R_0 @ R_1.T) @ K1.T @ skew_op(K1 @ R_1 @ R_0.T @ (T_0 - R_0 @ R_1.T @ T_1))

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 = img1.shape[:2]
    for r,pt1,pt2 in zip(lines,pts1,pts2):
        pt1 = list(map(lambda x:int(x+0.5), pt1[:2].tolist()))
        pt2 = list(map(lambda x:int(x+0.5), pt2[:2].tolist()))
        if pt1[0] < 0 or pt1[1] < 0:
            continue
        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] ])
        img1 = cv2.line(img1, (x0,y0), (x1,y1), color,1)
        img1 = cv2.circle(img1,tuple(pt1),5,color,-1)
        img2 = cv2.circle(img2,tuple(pt2),5,color,-1)
    return img1,img2

def SimpleConstrain(dimGroups):
    constrain = np.ones((dimGroups[-1], dimGroups[-1]))
    for i in range(len(dimGroups)-1):
        start, end = dimGroups[i], dimGroups[i+1]
        constrain[start:end, start:end] = 0
    N = constrain.shape[0]
    constrain[range(N), range(N)] = 1
    return constrain

def check_cluster(affinity, row, views, dimGroups, indices, p2dAssigned, visited):
    affinity_row = affinity[row].copy()
    # given affinity and row, select the combine of all possible set
    cluster = np.where((affinity[row]>0)&(p2dAssigned==-1)&(visited==0))[0].tolist()
    cluster.sort(key=lambda x:-affinity[row, x])
    views_ = views[cluster]
    view_count = np.bincount(views[cluster])
    indices_all = [indices]
    for col in cluster:
        v = views[col]
        nOld = len(indices_all)
        if indices[v] != -1: # already assigned, copy and make new 
            for i in range(nOld):
                ind = indices_all[i].copy()
                ind[v] = col
                indices_all.append(ind)
        else: # not assigned, assign
            for i in range(nOld):
                indices_all[i][v] = col
    return indices_all

def views_from_dimGroups(dimGroups):
    views = np.zeros(dimGroups[-1], dtype=int)
    for nv in range(len(dimGroups) - 1):
        views[dimGroups[nv]:dimGroups[nv+1]] = nv
    return views

class SimpleMatchAndTriangulator(SimpleTriangulator):
    def __init__(self, num_joints, min_views, min_joints, cfg_svt, cfg_track, **cfg) -> None:
        super().__init__(**cfg)
        self.nJoints = num_joints
        self.cfg_svt = cfg_svt
        self.cfg_track = cfg_track
        self.min_views = min_views
        self.min_joints = min_joints
        self.time = -1
        self.max_id = 0
        self.tracks = {}
        self.loglevel_dict = {
            'info': 0,
            'warn': 1,
            'error': 2,
        }
        self.loglevel = self.loglevel_dict['info'] # ['info', 'warn', 'error']
        self.debug = False
        self.data = None
        self.people = None
    
    def log(self, text):
        if self.loglevel > 0:
            return 0
        log(text)
    
    def warn(self, text):
        if self.loglevel > 1:
            return 0
        mywarn(text)
    

    @staticmethod
    def distance_by_epipolar(pts0, pts1, K0, K1, R0, T0, R1, T1):
        F = fundamental_op(K0, K1, R0, T0, R1, T1)
        # Find epilines corresponding to points in left image (first image) and
        # drawing its lines on right image
        lines0 = cv2.computeCorrespondEpilines(pts0[..., :2].reshape (-1,1,2), 2, F)
        # Find epilines corresponding to points in right image (second image) and
        # drawing its lines on left image
        lines1 = cv2.computeCorrespondEpilines(pts1[..., :2].reshape(-1,1,2), 1, F)
        if False:
            H, W = 1080, 1920
            img0 = np.zeros((H, W, 3), dtype=np.uint8) +255
            img4, img3 = drawlines(img0.copy(), img0.copy(), lines0.reshape(-1, 3), pts1.reshape(-1, 3), pts0.reshape(-1,3))
            img5,img6 = drawlines(img0.copy(), img0.copy(), lines1.reshape(-1, 3), pts0.reshape(-1,3), pts1.reshape(-1,3))
            import matplotlib.pyplot as plt
            plt.subplot(121)
            plt.imshow(img5)
            plt.subplot(122)
            plt.imshow(img4)
            plt.show()        
        lines0 = lines0.reshape(pts0.shape)
        lines1 = lines1.reshape(pts1.shape)
        # dist: (D_v0, D_v1, nJoints)
        dist01 = np.abs(np.sum(lines0[:, None, :, :2] * pts1[None, :, :, :2], axis=-1) + lines0[:, None, :, 2])
        conf = pts0[:, None, :, 2] * pts1[None, :, :, 2]
        dist10 = np.abs(np.sum(lines1[:, None, :, :2] * pts0[None, :, :, :2], axis=-1) + lines1[:, None, :, 2])

        dist = np.sum(dist01 * conf + dist10.transpose(1, 0, 2) * conf, axis=-1)/(conf.sum(axis=-1) + 1e-5)/2
        return dist

    def _simple_associate2d_triangulate(self, data, affinity, dimGroups, prev_id):
        # sum1 = affinity.sum(axis=1)
        # 注意：这里的排序应该是对每个视角，挑选最大的一个
        sum1 = np.zeros((affinity.shape[0]))
        for i in range(len(dimGroups)-1):
            start, end = dimGroups[i], dimGroups[i+1]
            if end == start:continue
            sum1 += affinity[:, start:end].max(axis=-1)
        n2d = affinity.shape[0]
        nViews = len(dimGroups) - 1
        idx_zero = np.zeros(nViews, dtype=int) - 1
        views = views_from_dimGroups(dimGroups)
        # the assigned results of each person
        p2dAssigned = np.zeros(n2d, dtype=int) - 1
        visited = np.zeros(n2d, dtype=int)
        sortidx = np.argsort(-sum1)
        pid = 0
        k3dresults = []
        for idx in sortidx:
            if p2dAssigned[idx] != -1:
                continue
            if prev_id[idx] != -1:
                results = [self.people[prev_id[idx]]]
            else:
                results = []
            proposals = check_cluster(affinity, row=idx, views=views, 
                dimGroups=dimGroups, indices=idx_zero.copy(), p2dAssigned=p2dAssigned, visited=visited)
            for indices in proposals:
                if (indices > -1).sum() < self.min_views - (len(results)):
                    continue
                # set keypoints2d
                info = {'RT': data['RT']}
                for name in ['keypoints2d', 'keypoints2d_unproj', 'keypoints2d_distort']:
                    info[name] = np.zeros((nViews, self.nJoints, 3), dtype=np.float32)
                for nv in range(nViews):
                    if indices[nv] == -1: continue
                    for name in ['keypoints2d', 'keypoints2d_unproj', 'keypoints2d_distort']:
                        info[name][nv] = data[name][nv][indices[nv]-dimGroups[nv]]

                res = super().__call__(info, results=results)[0]

                k2d = res['keypoints2d']
                valid_view = (k2d[..., 2] > 0).sum(axis=-1) > self.min_joints
                # if valid_view.sum() < self.min_views - len(results): # 这里如果是有前一帧的话，len(results)会是2；不知道之前为啥有这个条件使用
                if valid_view.sum() < self.min_views:
                    self.log('[associate] Skip proposal {}->{} with not enough valid view {}'.format(idx, indices, (k2d[..., 2] > 0).sum(axis=-1)))
                    continue
                valid_joint = res['keypoints3d'][:, -1] > 0.1
                if valid_joint.sum() < self.min_joints:
                    self.log('[associate] Skip proposal {}->{} as not enough joints'.format(idx, indices))
                    continue
                indices[~valid_view] = -1
                if (indices < 0).all():
                    import ipdb; ipdb.set_trace()
                self.log('[associate] Add indices {}, valid {}'.format(indices, (k2d[..., 2] > 0).sum(axis=-1)))
                res['id'] = pid
                res['indices'] = indices
                res['valid_view'] = valid_view
                res['valid_joints'] = res['keypoints3d'][:, -1] > 0.1
                k3dresults.append(res)
                for nv in range(nViews):
                    if valid_view[nv] and indices[nv] != -1:
                        p2dAssigned[indices[nv]] = pid
                        visited[indices[nv]] = 1
                pid += 1
                break
            visited[idx] = 1
        self.log('[associate] {} points not visited, {} not assigned'.format(visited.shape[0] - visited.sum(), (p2dAssigned==-1).sum()))
        k3dresults.sort(key=lambda x: -x['keypoints2d'][..., -1].sum())
        return k3dresults

    def _calculate_affinity_MxM(self, dims, dimGroups, data, key):
        M = dimGroups[-1]
        distance = np.zeros((M, M), dtype=np.float32)
        nViews = len(dims)
        for v0 in range(nViews-1):
            for v1 in range(1, nViews):
                # calculate distance between (v0, v1)
                if v0 >= v1:
                    continue
                if dims[v0] == 0 or dims[v1] == 0:
                    continue
                if True:
                    pts0 = data[key][v0]
                    pts1 = data[key][v1]
                    K0, K1 = data['K'][v0], data['K'][v1]
                    R0, T0 = data['Rc'][v0], data['Tc'][v0]
                    R1, T1 = data['Rc'][v1], data['Tc'][v1]
                    dist = self.distance_by_epipolar(pts0, pts1, K0, K1, R0, T0, R1, T1)
                    dist /= (K0[0, 0] + K1[0, 0])/2
                else:
                    dist = self.distance_by_ray(pts0, pts1, R0, T0, R1, T1)
                distance[dimGroups[v0]:dimGroups[v0+1], dimGroups[v1]:dimGroups[v1+1]] = dist
                distance[dimGroups[v1]:dimGroups[v1+1], dimGroups[v0]:dimGroups[v0+1]] = dist.T
        DIST_MAX = self.cfg_track.track_dist_max
        for nv in range(nViews):
            distance[dimGroups[nv]:dimGroups[nv+1], dimGroups[nv]:dimGroups[nv+1]] = DIST_MAX
        distance -= np.eye(M) * DIST_MAX
        aff = (DIST_MAX - distance)/DIST_MAX
        aff = np.clip(aff, 0, 1)
        return aff

    def _calculate_affinity_MxN(self, dims, dimGroups, data, key, results):
        M = dimGroups[-1]
        N = len(results)
        distance = np.zeros((M, N), dtype=np.float32)
        nViews = len(dims)
        k3d = np.stack([r['keypoints3d'] for r in results])
        kpts_proj = project_points(k3d, data['KRT'], einsum='vab,pkb->vpka')
        depth = kpts_proj[..., -1]
        kpts_proj[depth<0] = -10000
        for v in range(nViews):
            if dims[v] == 0:
                continue
            focal = data['K'][v][0, 0]
            pts2d = data[key][v][:, None]
            pts_repro = kpts_proj[v][None]
            conf = np.sqrt(pts2d[..., -1]*k3d[None, ..., -1])
            diff = np.linalg.norm(pts2d[..., :2] - pts_repro[..., :2], axis=-1)
            diff = np.sum(diff*conf, axis=-1)/(1e-5 + np.sum(conf, axis=-1))
            dist = diff / focal
            distance[dimGroups[v]:dimGroups[v+1], :] = dist
        DIST_MAX = self.cfg_track.track_repro_max
        aff = (DIST_MAX - distance)/DIST_MAX
        aff = np.clip(aff, 0, 1)
        return aff

    def _svt_optimize_affinity(self, affinity, dimGroups):
        # match SVT
        import pymatchlr
        observe = np.ones_like(affinity)
        aff_svt = pymatchlr.matchSVT(affinity, dimGroups, SimpleConstrain(dimGroups), observe, self.cfg_svt)
        aff_svt[aff_svt<self.cfg_svt.aff_min] = 0.
        if False:
            import matplotlib.pyplot as plt
            M = affinity.shape[0]
            plt.subplot(121)
            plt.imshow(affinity)
            plt.hlines([i-0.5 for i in dimGroups[1:]], -0.5, M-0.5, 'w')
            plt.vlines([i-0.5 for i in dimGroups[1:]], -0.5, M-0.5, 'w')
            plt.subplot(122)
            sum_row = aff_svt.sum(axis=1, keepdims=True)/(len(dimGroups) - 1)
            plt.imshow(np.hstack([aff_svt, sum_row]))
            plt.hlines([i-0.5 for i in dimGroups[1:]], -0.5, M-0.5, 'w')
            plt.vlines([i-0.5 for i in dimGroups[1:]], -0.5, M-0.5, 'w')
            plt.ioff()
            plt.show()
            import ipdb;ipdb.set_trace()
        return aff_svt

    def _track_add(self, res):
        pid = res['id']
        if pid == -1:
            pid = self.max_id
            res['id'] = pid
            self.max_id += 1
            self.log('[{:06d}] Create track {} <- {}'.format(self.time, pid, res['indices']))
            if False:
                crops = []
                data = self.data
                kpts = np.vstack(data['keypoints2d'])
                for nv in range(len(data['imgname'])):
                    img = cv2.imread(data['imgname'][nv])
                    if res['indices'][nv] == -1: continue
                    _kpts = kpts[res['indices'][nv]]
                    bbox = bbox_from_keypoints(_kpts)
                    plot_keypoints_auto(img, _kpts, pid)
                    crop = crop_image(img, bbox, crop_square=True)
                    crops.append(crop)
                debug = merge(crops)
                cv2.imwrite('debug/{:06d}.jpg'.format(pid), debug)
        else:
            self.max_id = max(self.max_id, pid+1)
            self.log('[{:06d}] Initialize track {}, valid joints={}'.format(self.time, pid, (res['keypoints3d'][:, -1]>0.01).sum()))
        self.tracks[pid] = {
            'start_time': self.time,
            'end_time': self.time+1,
            'missing_frame': [],
            'infos': [res]
        }
    
    def _track_update(self, res, pid):
        res['id'] = pid
        info = self.tracks[pid]
        self.log('[{:06d}] Update track {} [{}->{}], valid joints={}'.format(self.time, pid, info['start_time'], info['end_time'], (res['keypoints3d'][:, -1]>0.1).sum()))
        self.tracks[pid]['end_time'] = self.time + 1
        self.tracks[pid]['infos'].append(res)
    
    def _track_merge(self, res, pid):
        res['id'] = -1
        # TODO: merge

    def _track_and_update(self, data, results):
        cfg = self.cfg_track
        self.time += 1
        if self.time == 0:
            # initialize the tracks
            for res in results:
                self._track_add(res)
            return results
        # filter the missing frames
        for pid in list(self.tracks.keys()):
            if self.time - self.tracks[pid]['end_time'] > cfg.max_missing_frame:
                self.warn('[{:06d}] Remove track {}'.format(self.time, pid))
                self.tracks.pop(pid)
        # track the results with greedy matching
        for idx_match, res in enumerate(results):
            res['id'] = -1
            # compute the distance
            k3d = res['keypoints3d'][None]
            pids_free = [pid for pid in self.tracks.keys() if self.tracks[pid]['end_time'] != self.time+1]
            pids_used = [pid for pid in self.tracks.keys() if self.tracks[pid]['end_time'] == self.time+1]
            def check_dist(k3d_check):
                dist = np.linalg.norm(k3d[..., :3] - k3d_check[..., :3], axis=-1)
                conf = np.sqrt(k3d[..., 3] * k3d_check[..., 3])
                dist_mean = ((conf>0.1).sum(axis=-1) < self.min_joints)*cfg.track_dist_max + np.sum(dist * conf, axis=-1)/(1e-5 + np.sum(conf, axis=-1))
                argmin = dist_mean.argmin()
                dist_min = dist_mean[argmin]
                return dist_mean, argmin, dist_min
            # check free
            NOT_VISITED = -2
            NOT_FOUND = -1
            flag_tracked, flag_current = NOT_VISITED, NOT_VISITED
            if len(pids_free) > 0:
                k3d_check = np.stack([self.tracks[pid]['infos'][-1]['keypoints3d'] for pid in pids_free])
                dist_track, best, best_dist_track = check_dist(k3d_check)
                if best_dist_track < cfg.track_dist_max:
                    flag_tracked = best
                else:
                    flag_tracked = NOT_FOUND
            # check used
            if len(pids_used) > 0:
                k3d_check = np.stack([self.tracks[pid]['infos'][-1]['keypoints3d'] for pid in pids_used])
                dist_cur, best, best_dist_curr = check_dist(k3d_check)
                if best_dist_curr < cfg.track_dist_max:
                    flag_current = best
                else:
                    flag_current = NOT_FOUND
            if flag_tracked >= 0 and (flag_current == NOT_VISITED or flag_current == NOT_FOUND):
                self._track_update(res, pids_free[flag_tracked])
            elif (flag_tracked == NOT_FOUND or flag_tracked==NOT_VISITED) and flag_current >= 0:
                # 没有跟踪到，但是有当前帧的3D的，合并
                self.log('[{:06d}] Merge track {} to {}'.format(self.time, idx_match, pids_used[flag_current]))
                self._track_merge(res, pids_used[flag_current])
            elif flag_tracked == NOT_FOUND and flag_current == NOT_FOUND:
                # create a new track
                self._track_add(res)
            else:
                # 丢弃
                self.log('[{:06d}] Remove track {}. No close points'.format(self.time, idx_match))

        for pid in list(self.tracks.keys()):
            if self.tracks[pid]['end_time'] != self.time + 1:
                self.warn('[{:06d}] Tracking {} missing'.format(self.time, pid))
        results = [r for r in results if r['id']!=-1]
        return results

    def __call__(self, data):
        # match the data
        self.data = data
        key = 'keypoints2d'
        dims = [d.shape[0] for d in data[key]]
        dimGroups = np.cumsum([0] + dims)
        # 1. compute affinity
        affinity = self._calculate_affinity_MxM(dims, dimGroups, data, key)
        N2D = affinity.shape[0]
        if self.people is not None and len(self.people) > 0:
            # add 3d affinity
            _affinity = affinity
            affinity_3d = self._calculate_affinity_MxN(dims, dimGroups, data, key, self.people)
            affinity = np.concatenate([affinity, affinity_3d], axis=1)
            eye3d = np.eye(affinity_3d.shape[1])
            affinity = np.concatenate([affinity, np.hstack((affinity_3d.T, eye3d))], axis=0)
            dimGroups = dimGroups.tolist()
            dimGroups.append(dimGroups[-1]+affinity_3d.shape[1])
            affinity = self._svt_optimize_affinity(affinity, dimGroups)
            # affinity = self._svt_optimize_affinity(_affinity, dimGroups[:-1])
            # recover
            affinity_3d = np.hstack([np.ones((N2D, 1))*0.5, affinity[:N2D, N2D:]])
            prev_id = affinity_3d.argmax(axis=-1) - 1
            affinity = affinity[:N2D, :N2D]
            dimGroups = np.array(dimGroups[:-1])
        else:
            affinity = self._svt_optimize_affinity(affinity, dimGroups)
            prev_id = np.zeros(N2D) - 1
        # 2. associate and triangulate
        results = self._simple_associate2d_triangulate(data, affinity, dimGroups, prev_id)
        # 3. track, filter and return
        results = self._track_and_update(data, results)
        results.sort(key=lambda x:x['id'])
        self.people = results
        return results