CPU Instance iou

Author: nicolas-chaulet

cuda/include/metrics.h

@@ -0,0 +1,5 @@
+#pragma once
+#include <torch/extension.h>
+
+at::Tensor instance_iou_cuda(at::Tensor instance_idx, at::Tensor instance_offsets,
+                             at::Tensor instance_gt);

cuda/src/metrics.cpp

@@ -0,0 +1,27 @@
+#include "metrics.h"
+#include "compat.h"
+#include "utils.h"
+
+void instance_iou_kernel_wrapper(int b, int n, int m, const float* dataset, float* temp, int* idxs);
+
+at::Tensor instance_iou_cuda(at::Tensor instance_idx, at::Tensor instance_offsets,
+                             at::Tensor instance_gt)
+{
+    CHECK_CONTIGUOUS(instance_idx);
+    CHECK_CONTIGUOUS(instance_offsets);
+    CHECK_CONTIGUOUS(instance_gt);
+    CHECK_CUDA(instance_idx)
+    CHECK_CUDA(instance_offsets)
+    CHECK_CUDA(instance_gt)
+
+    auto num_gt_instances = instance_gt.max(0);
+    auto num_proposed_instances = instance_offsets.size(0);
+    at::Tensor output =
+        torch::zeros({num_proposed_instances, num_gt_instances},
+                     at::device(num_gt_instances.device()).dtype(at::ScalarType::Float));
+
+    instance_iou_kernel_wrapper(points.size(0), points.size(1), nsamples, points.DATA_PTR<float>(),
+                                tmp.DATA_PTR<float>(), output.DATA_PTR<float>());
+
+    return output;
+}

cuda/src/metrics.cu

@@ -0,0 +1,8 @@
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "cuda_utils.h"
+
+void instance_iou_kernel_wrapper(int b, int n, int m, const float* dataset, float* temp,
+                                            int* idxs);

test/test_cluster.py

@@ -34,18 +34,16 @@ def test_simple(self):
         self.assertEqual(clusters, [[0, 1, 2], [4, 5, 6]])
 
     def test_region_grow(self):
-        clusters = region_grow(
+        cluster_idx = region_grow(
             self.pos, self.labels, self.batch, radius=2, min_cluster_size=1
         )
-        self.assertEqual(len(clusters[0]), 2)
-        self.assertEqual(len(clusters[1]), 3)
-        self.assertEqual(len(clusters[10]), 1)
-        torch.testing.assert_allclose(clusters[0][0], torch.tensor([0, 1]))
-        torch.testing.assert_allclose(clusters[0][1], torch.tensor([4]))
-        torch.testing.assert_allclose(clusters[1][0], torch.tensor([2]))
-        torch.testing.assert_allclose(clusters[1][1], torch.tensor([3]))
-        torch.testing.assert_allclose(clusters[1][2], torch.tensor([5, 6]))
-        torch.testing.assert_allclose(clusters[10][0], torch.tensor([7]))
+        self.assertEqual(len(cluster_idx), 6)
+        torch.testing.assert_allclose(cluster_idx[0], torch.tensor([0, 1]))
+        torch.testing.assert_allclose(cluster_idx[1], torch.tensor([4]))
+        torch.testing.assert_allclose(cluster_idx[2], torch.tensor([2]))
+        torch.testing.assert_allclose(cluster_idx[3], torch.tensor([3]))
+        torch.testing.assert_allclose(cluster_idx[4], torch.tensor([5, 6]))
+        torch.testing.assert_allclose(cluster_idx[5], torch.tensor([7]))
 
 
 if __name__ == "__main__":

test/test_metrics.py

@@ -0,0 +1,29 @@
+import unittest
+import torch
+import os
+import sys
+import numpy as np
+
+ROOT = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..")
+sys.path.insert(0, ROOT)
+
+from torch_points_kernels.metrics import instance_iou
+
+
+class TestInstanceIou(unittest.TestCase):
+    def test_simple(self):
+        gt_instances = torch.tensor([1, 2, 1, 2, 2, 3, 0])
+        proposed_instances = [
+            torch.tensor([0, 2]),  # 100% instance 1
+            torch.tensor([1, 4]),  # 2/3 of instance 2
+            torch.tensor([3, 5]),  # 1/3 of instance 2 and 1/1 of instance 3
+        ]
+
+        ious = instance_iou(proposed_instances, gt_instances)
+        torch.testing.assert_allclose(
+            ious, torch.tensor([[1, 0, 0], [0, 2 / 3.0, 0], [0, 1.0 / 4.0, 1.0 / 2.0]])
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

torch_points_kernels/cluster.py

@@ -2,6 +2,7 @@
 from .torchpoints import ball_query_partial_dense
 import numpy as np
 import numba
+from typing import List
 
 
 @numba.jit(nopython=True)
@@ -50,7 +51,7 @@ def grow_proximity(pos, batch, nsample=16, radius=0.02, min_cluster_size=32):
 
 def region_grow(
     pos, labels, batch, ignore_labels=[], nsample=16, radius=0.02, min_cluster_size=32
-):
+) -> List[torch.Tensor]:
     """ Region growing clustering algorithm proposed in
     PointGroup: Dual-Set Point Grouping for 3D Instance Segmentation
     https://arxiv.org/pdf/2004.01658.pdf
@@ -76,7 +77,7 @@ def region_grow(
     assert pos.shape[0] == labels.shape[0]
 
     unique_labels = torch.unique(labels)
-    clusters = {}
+    clusters = []
     ind = torch.arange(0, pos.shape[0])
     for l in unique_labels:
         if l in ignore_labels:
@@ -95,10 +96,8 @@ def region_grow(
 
         # Remap indices to original coordinates
         if len(label_clusters):
-            remaped_clusters = []
             for cluster in label_clusters:
                 cluster = torch.tensor(cluster).to(pos.device)
-                remaped_clusters.append(local_ind[cluster])
-            clusters[l.item()] = remaped_clusters
+                clusters.append(local_ind[cluster])
 
     return clusters

torch_points_kernels/metrics.py

@@ -0,0 +1,60 @@
+import torch
+from typing import List
+import numpy as np
+import numba
+
+
+@numba.jit(nopython=True, parallel=True)
+def _instance_iou_cpu(instance_idx, instance_offsets, instance_gt):
+    num_instances = np.max(instance_gt)
+    gt_instance_sizes = []
+    for instance_id in numba.prange(1, num_instances + 1):
+        gt_instance_sizes.append(np.sum(instance_gt == instance_id))
+    iou = np.zeros((len(instance_offsets), num_instances))
+    old_offset = 0
+    for proposed_instance, offset in enumerate(instance_offsets):
+        instance = instance_idx[old_offset:offset]
+        for instance_id in numba.prange(1, num_instances + 1):
+            intersection = 0
+            for idx in instance:
+                if instance_gt[idx] == instance_id:
+                    intersection += 1
+            iou[proposed_instance, instance_id - 1] = intersection / float(
+                len(instance) + gt_instance_sizes[instance_id - 1] - intersection
+            )
+        old_offset = offset
+    return iou
+
+
+def instance_iou(instance_idx: List[torch.Tensor], gt_instances: torch.Tensor):
+    """ Computes the IoU between each proposed instance in instance_idx and ground truth instances. Returns a 
+    tensor of shape [instance_idx.shape[0], num_instances] that contains the iou between the proposed instances and all gt instances
+    Instance label 0 is reserved for non instance points
+
+    Parameters
+    ----------
+    instance_idx : List[torch.Tensor]
+        List of instances. Each tensor in this list is a proposed and contains the index of the points 
+        that belong to that particular instance
+    gt_instances : torch.Tensor
+        Ground truth instances, contains the index of the instance for each point
+    num_instances : int
+        Number of instances in ground truth
+
+    Returns
+    -------
+    ious: 
+    """
+    instance_offsets = []
+    cum_offset = 0
+    for instance in instance_idx:
+        cum_offset += instance.shape[0]
+        instance_offsets.append(cum_offset)
+    instance_idx = torch.cat(instance_idx)
+    if gt_instances.is_cuda:
+        pass
+    else:
+        res = _instance_iou_cpu(
+            instance_idx.numpy(), np.asarray(instance_offsets), gt_instances.numpy()
+        )
+        return torch.tensor(res).float()