From 2262c4c29cfaf6e1c44443d059dc617672f159ab Mon Sep 17 00:00:00 2001
From: asbe <asbjorn.berge@sintef.no>
Date: Wed, 6 Nov 2024 14:15:48 +0100
Subject: [PATCH] Update quantization.cpp

Fixes a bug where conflicting labels are handled incorrectly.

Minimal script for reproducing bug

import MinkowskiEngine as ME
import numpy as np
#-----Dummy inputs
#(Other examples without conflicting labels, and with one point at a time behaves as expected)
#Case 0: 2 pairs of redundant points, 1) with same labels 2) with different labels.
#Expected behaviour: coords_aug_new = [1, 0, 0], labels_new = [11, 255], feats = [0,0,101], return_index = [0, 2], return_mapping = [0, 0, ]
#Result: one ignore label, but in the wrong position.
coords_aug = np.float32(np.array([
[1, 0, 0],[1, 0, 0],
[2, 0, 0], [2, 0, 0]
]))
labels = np.array([
11,11,
20,21,
])
feats = np.array([
[0, 0, 101],[0, 0, 101],
[0, 0, 101], [0, 0, 101]
])
return_args_tmp = ME.utils.sparse_quantize(
coords_aug, feats, labels=labels.astype(np.int32), ignore_label=255, return_index = True, return_inverse = True), #quantization_size=1)
(coords_aug_new,feats_new,labels_new,voxel_return_index, voxel_return_mapping)=return_args_tmp[0]
print("New labels: ", labels_new)
print("New coords: ", coords_aug_new)
print("Return index", voxel_return_index)
print("Return mapping", voxel_return_mapping)
---
 src/quantization.cpp | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/src/quantization.cpp b/src/quantization.cpp
index 0311bde2..322e451b 100644
--- a/src/quantization.cpp
+++ b/src/quantization.cpp
@@ -175,23 +175,21 @@ std::vector<std::vector<int>> quantize_label(int const *const p_coords,
   for (int row = 0; row < nrows; ++row) {
     auto key = coordinate<coordinate_type>(p_coords + ncols * row);
     const auto it = map.find(key);
-    auto iter_success =
-        map.insert(value_type(key, mapped_type(n_unique, p_labels[row])));
-    if (iter_success.second) {
-      // success
+    if (it == map.end()) {
+      // New coordinate
+      map.insert(value_type(key, mapped_type(n_unique, p_labels[row])));
       mapping.push_back(row);
       colabels.push_back(p_labels[row]);
       inverse_mapping[row] = n_unique++;
     } else {
-      auto &keyval = *(iter_success.first);
-      auto &val = keyval.second;
-      // Set the label
-      if (val.second != p_labels[row] && val.second != invalid_label) {
-        // When the labels differ
+      // Existing coordinate
+      auto &val = it->second;
+      if (val.second != p_labels[row]) {
+        // Conflicting labels, set to invalid_label
         val.second = invalid_label;
-        colabels[inverse_mapping[val.first]] = invalid_label;
+        colabels[val.first] = invalid_label;
       }
-      inverse_mapping[row] = val.first; // row
+      inverse_mapping[row] = val.first;
     }
   }