reduce mask details

SVivdich02 · SVivdich02 · commit 8ae5cdf91882 · 2024-03-04T23:41:06.000+03:00
diff --git a/example.ipynb b/example.ipynb
@@ -26,7 +26,7 @@
     "    **{\n",
     "        \"dataset\": kitti,\n",
     "        \"start_index\": 20,\n",
-    "        \"end_index\": 23,\n",
+    "        \"end_index\": 24,\n",
     "        \"start_image_index_offset\": 3,\n",
     "        \"cam_name\": \"cam2\",\n",
     "        \"R\": 16,\n",
@@ -50,6 +50,30 @@
     "points2instances = InitInstancesMatrixProcessor().process(config, init_pcd)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "points2instances.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy\n",
+    "\n",
+    "from src.utils.pcd_utils import color_pcd_by_labels\n",
+    "from src.utils.pcd_utils import visualize_pcd\n",
+    "\n",
+    "colored_pcd = color_pcd_by_labels(copy.deepcopy(init_pcd), points2instances[:, 0])\n",
+    "visualize_pcd(colored_pcd)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -79,7 +103,8 @@
     "for processor in processors:\n",
     "    pcd, points2instances = processor.process(config, pcd, points2instances)\n",
     "\n",
-    "pcd_for_clustering = copy.deepcopy(pcd)"
+    "pcd_for_clustering = copy.deepcopy(pcd)\n",
+    "points2instances_pcd_for_clustering = copy.deepcopy(points2instances)"
    ]
   },
   {
@@ -107,7 +132,7 @@
     "points = np.asarray(pcd.points)\n",
     "spatial_distance = cdist(points, points)\n",
     "\n",
-    "dist, masks = sam_label_distance(points2instances, spatial_distance, 5, 10, 2)"
+    "dist, masks = sam_label_distance(points2instances, spatial_distance, 3, 3, 5)"
    ]
   },
   {
@@ -156,7 +181,7 @@
    "source": [
     "from src.services.normalized_cut_service import normalized_cut\n",
     "\n",
-    "T = 0.07\n",
+    "T = 0.02\n",
     "eigenval = 2\n",
     "\n",
     "clusters = normalized_cut(\n",
@@ -198,6 +223,18 @@
     "visualize_pcd(pcd_colored)"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from src.utils.pcd_utils import color_pcd_by_labels\n",
+    "\n",
+    "pcd_src_colored = color_pcd_by_labels(pcd_for_clustering, points2instances_pcd_for_clustering[:, 2])\n",
+    "visualize_pcd(pcd_src_colored)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/labelling_results/date0403-dist335-T0d02-ei2.png b/labelling_results/date0403-dist335-T0d02-ei2.png
diff --git a/src/services/preprocessing/init/instances_matrix.py b/src/services/preprocessing/init/instances_matrix.py
@@ -17,8 +17,11 @@
 import zope.interface
 
 from src.services.preprocessing.common.interface import IProcessor
+from src.utils.geometry_utils import calculate_area
 from src.utils.pcd_utils import get_subpcd
 from src.utils.pcd_utils import get_visible_points
+from src.utils.sam_mask_utils import find_intersection_mask
+from src.utils.sam_mask_utils import find_union_mask
 
 
 @zope.interface.implementer(IProcessor)
@@ -56,7 +59,8 @@ def build_points2instances_matrix(
         points2instances = np.zeros((N, end_image_index - start_image_index), dtype=int)
 
         for view_id, view in enumerate(range(start_image_index, end_image_index)):
-            masks = dataset.get_image_instances(cam_name, view)
+            full_masks = dataset.get_image_instances(cam_name, view)
+            masks = self.reduce_detail(full_masks)
             image_labels = self.masks_to_image(masks)
 
             T = dataset.get_lidar_pose(view)
@@ -111,3 +115,52 @@ def masks_to_image(self, masks):
         for i, mask in enumerate(masks):
             image_labels[mask["segmentation"]] = i + 1
         return image_labels
+
+    def reduce_detail(self, masks, intersection_to_union_ratio_threshold=0.35):
+        merged_mask = []
+        merged_indices = []
+
+        for i in range(len(masks)):
+            if i in merged_indices:
+                continue
+
+            area_bbox_i = calculate_area(masks[i]['bbox'])
+
+            indices_merged_with_i = []
+            for j in range(i + 1, len(masks)):
+                if j in merged_indices:
+                    continue
+
+                area_bbox_j = calculate_area(masks[j]['bbox'])
+
+                intersection_mask = find_intersection_mask(masks[i], masks[j])
+                if intersection_mask == None:
+                    continue
+                area_intersection = intersection_mask['area']
+
+                area_bbox_intersection = calculate_area(intersection_mask['bbox'])
+                area_bbox_union = area_bbox_i + area_bbox_j - area_bbox_intersection
+                IU_ratio = area_bbox_intersection / area_bbox_union
+
+                if (IU_ratio >= intersection_to_union_ratio_threshold
+                    or area_intersection / masks[i]['area'] >= 0.6
+                    or area_intersection / masks[j]['area'] >= 0.6):
+                    masks[i] = find_union_mask(masks[i], masks[j])
+                    indices_merged_with_i.append(j)
+
+            if indices_merged_with_i:
+                merged_mask.append(masks[i])
+
+                merged_indices.append(i)
+                for ind in indices_merged_with_i:
+                    merged_indices.append(ind)
+
+        masks_result = []
+        for ind, mask in enumerate(masks):
+            if ind not in merged_indices:
+                masks_result.append(mask)
+
+        for mask in merged_mask:
+            masks_result.append(mask)
+
+        return masks_result
diff --git a/src/utils/geometry_utils.py b/src/utils/geometry_utils.py
@@ -0,0 +1,56 @@
+# Copyright (c) 2023, Sofia Vivdich and Anastasiia Kornilova
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def find_intersection(bbox1, bbox2):
+    x1_bbox1, y1_bbox1, w_bbox1, h_bbox1 = bbox1
+    x1_bbox2, y1_bbox2, w_bbox2, h_bbox2 = bbox2
+
+    x2_bbox1 = x1_bbox1 + w_bbox1
+    y2_bbox1 = y1_bbox1 + h_bbox1
+    x2_bbox2 = x1_bbox2 + w_bbox2
+    y2_bbox2 = y1_bbox2 + h_bbox2
+
+    if (x1_bbox1 > x2_bbox2 or x2_bbox1 < x1_bbox2 or y1_bbox1 > y2_bbox2 or y2_bbox1 < y1_bbox2):
+        return None
+    
+    x_left = max(x1_bbox1, x1_bbox2)
+    x_right = min(x2_bbox1, x2_bbox2)
+    y_top = max(y1_bbox1, y1_bbox2)
+    y_bottom = min(y2_bbox1, y2_bbox2)
+
+    return x_left, y_top, x_right - x_left, y_bottom - y_top
+
+
+def find_union(bbox1, bbox2):
+    x1_bbox1, y1_bbox1, w_bbox1, h_bbox1 = bbox1
+    x1_bbox2, y1_bbox2, w_bbox2, h_bbox2 = bbox2
+
+    x2_bbox1 = x1_bbox1 + w_bbox1
+    y2_bbox1 = y1_bbox1 + h_bbox1
+    x2_bbox2 = x1_bbox2 + w_bbox2
+    y2_bbox2 = y1_bbox2 + h_bbox2
+
+    x_left = min(x1_bbox1, x1_bbox2)
+    y_top = min(y1_bbox1, y1_bbox2)
+    x_right = max(x2_bbox1, x2_bbox2)
+    y_bottom = max(y2_bbox1, y2_bbox2)
+
+    return x_left, y_top, x_right - x_left, y_bottom - y_top
+
+
+def calculate_area(bbox):
+    w = bbox[2]
+    h = bbox[3]
+    return w * h
diff --git a/src/utils/sam_mask_utils.py b/src/utils/sam_mask_utils.py
@@ -0,0 +1,44 @@
+# Copyright (c) 2023, Sofia Vivdich and Anastasiia Kornilova
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+
+from src.utils.geometry_utils import find_intersection
+from src.utils.geometry_utils import find_union
+
+
+def find_intersection_mask(mask1, mask2):
+    bbox = find_intersection(mask1['bbox'], mask2['bbox'])
+    if bbox == None:
+        return None
+    segmentation = mask1['segmentation'] * mask2['segmentation']
+    area = segmentation.sum()
+
+    intersection_mask = copy.deepcopy(mask1)
+    intersection_mask['segmentation'] = segmentation
+    intersection_mask['bbox'] = bbox
+    intersection_mask['area'] = area
+    return intersection_mask
+
+
+def find_union_mask(mask1, mask2):
+    segmentation = mask1['segmentation'] + mask2['segmentation']
+    bbox = find_union(mask1['bbox'], mask2['bbox'])
+    area = segmentation.sum()
+
+    union_mask = copy.deepcopy(mask1)
+    union_mask['segmentation'] = segmentation
+    union_mask['bbox'] = bbox
+    union_mask['area'] = area
+    return union_mask
