def get_mapping_ind(bwd_flows, bwd_occs, imgs, scale=1.0):
    """
    FLATTEN: Optical fLow-guided attention (Temoporal-guided attention)
    Find pixel correspondence between every consecutive frames in a batch

    [input]
    bwd_flow: (N-1)*2*H*W
    bwd_occ: (N-1)*H*W
    imgs: N*3*H*W

    [output]
    fwd_mappings: N*1*HW
    bwd_mappings: N*1*HW
    flattn_mask: HW*1*N*N
    i.e., imgs[i,:,fwd_mappings[i]] corresponds to imgs[0]
    i.e., imgs[i,:,fwd_mappings[i]][:,bwd_mappings[i]] restore the original imgs[i]
    """
    N, H, W = imgs.shape[0], int(imgs.shape[2] // scale), int(imgs.shape[3] // scale)
    iterattn_mask = torch.ones(H * W, N, N, dtype=torch.bool).to(imgs.device)
    for i in range(len(imgs) - 1):
        one_mask = torch.ones(N, N, dtype=torch.bool).to(imgs.device)
        one_mask[: i + 1, i + 1 :] = False
        one_mask[i + 1 :, : i + 1] = False
        mapping_ind, unlinkedmask = get_single_mapping_ind(
            bwd_flows[i : i + 1], bwd_occs[i : i + 1], imgs[i : i + 2], scale
        )
        if i == 0:
            fwd_mapping = [torch.arange(len(mapping_ind)).to(mapping_ind.device)]
            bwd_mapping = [torch.arange(len(mapping_ind)).to(mapping_ind.device)]
        iterattn_mask[unlinkedmask[fwd_mapping[-1]]] = torch.logical_and(
            iterattn_mask[unlinkedmask[fwd_mapping[-1]]], one_mask
        )
        fwd_mapping += [mapping_ind[fwd_mapping[-1]]]
        bwd_mapping += [torch.sort(fwd_mapping[-1])[1]]
    fwd_mappings = torch.stack(fwd_mapping, dim=0).unsqueeze(1)
    bwd_mappings = torch.stack(bwd_mapping, dim=0).unsqueeze(1)
    return fwd_mappings, bwd_mappings, iterattn_mask.unsqueeze(1)

    FLATTEN: Optical fLow-guided attention (Temoporal-guided attention)
    Find pixel correspondence between every consecutive frames in a batch

    [input]
    bwd_flow: (N-1)*2*H*W
    bwd_occ: (N-1)*H*W
    imgs: N*3*H*W

    [output]
    fwd_mappings: N*1*HW
    bwd_mappings: N*1*HW
    flattn_mask: HW*1*N*N
    i.e., imgs[i,:,fwd_mappings[i]] corresponds to imgs[0]
    i.e., imgs[i,:,fwd_mappings[i]][:,bwd_mappings[i]] restore the original imgs[i]
    
examples/community/fresco_v2v.py