Add 'None' option for point_reduction in chamfer_distance function

[nematoli]

🚀 Feature

Currently point_reduction entry in the chamfer_distance function should either be "mean" or "sum". That would be great to have the option of accessing output without any reduction.

Motivation

As the chamfer distance uses nearest neighbor for data association it is usually quite noisy. In practice it would be helpful to also enforce smoothing constraints (such as edge-aware smoothness) on the computed distance transform to reduce the noise. That's why having access to the not-reduced form of distance transform can be helpful.

Pitch

if point_reduction is not None:
    # Apply point reduction
    cham_x = cham_x.sum(1)  # (N,)
    cham_y = cham_y.sum(1)  # (N,)
    if return_normals:
        cham_norm_x = cham_norm_x.sum(1)  # (N,)
        cham_norm_y = cham_norm_y.sum(1)  # (N,)
    if point_reduction == "mean":
        cham_x /= x_lengths
        cham_y /= y_lengths
        if return_normals:
            cham_norm_x /= x_lengths
            cham_norm_y /= y_lengths
    if batch_reduction is not None:
        # batch_reduction == "sum"
        cham_x = cham_x.sum()
        cham_y = cham_y.sum()
        if return_normals:
            cham_norm_x = cham_norm_x.sum()
            cham_norm_y = cham_norm_y.sum()
        if batch_reduction == "mean":
            div = weights.sum() if weights is not None else N
            cham_x /= div
            cham_y /= div
            if return_normals:
                cham_norm_x /= div
                cham_norm_y /= div
cham_dist = cham_x + cham_y
cham_normals = cham_norm_x + cham_norm_y if return_normals else None
return cham_dist, cham_normals

[gkioxari]

You can submit a PR and we are happy to review and accept it!

[walsvid]

The two point clouds used to calculate the Chamfer loss may have different numbers of points, so adding the bi-directional loss directly without reduction will cause the shape mismatch. If Add 'None' option for point_reduction in the chamfer_distance function, I think a better solution is to return the loss in the form of a tuple.

[yuvalH9]

@nematollahi, @gkioxari , @walsvid is there any new with this issue?

[haritha-j]

I've found that I need this as well, I'll submit a PR.