feat(notebooks): add proposal segmentation notebook and enhance verification

notebooks/verification.py

@@ -61,12 +61,13 @@ def setup_verification_context(
     np,
 ):
     scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
-    image_size = 256
+    image_size = 512
     num_rooms = 4
     views_per_room = 6
     meters_per_pixel = 0.05
 
     sam_max_masks = 5
+    sam_candidate_masks = 24
     sam_min_area = 32 * 32
     hash_bits = 512
     pipeline_batch_size = 64
@@ -102,6 +103,7 @@ def setup_verification_context(
         meters_per_pixel,
         pipeline_batch_size,
         room_nodes,
+        sam_candidate_masks,
         sam_max_masks,
         sam_min_area,
         sim,
@@ -142,6 +144,7 @@ def build_scene_graph_pipeline(
     np,
     pipeline_batch_size,
     room_nodes,
+    sam_candidate_masks,
     sam_max_masks,
     sam_min_area,
     sim,
@@ -149,6 +152,7 @@ def build_scene_graph_pipeline(
 ):
     from rich.progress import track
 
+    from compressors import MaskScoringConfig, rank_masks
     from compressors.pipeline import HashPipeline
     from compressors.proposal import extract_masked_region, generate_proposals_batch
     from simulator import save_object_image, save_room_view
@@ -176,7 +180,17 @@ def build_scene_graph_pipeline(
     verification_output_dir = cfg_manager.get().output.directory / "verification"
     verification_output_dir.mkdir(parents=True, exist_ok=True)
 
+    mask_scoring_config = MaskScoringConfig(
+        max_area_ratio=0.45,
+        reject_edge_touch_count=3,
+        reject_large_edge_touch_count=2,
+        reject_large_edge_area_ratio=0.12,
+        max_components=4,
+        min_largest_component_ratio=0.70,
+    )
+
     total_masks = 0
+    total_raw_masks = 0
     object_index = 0
 
     room_view_dataset = [
@@ -192,7 +206,7 @@ def build_scene_graph_pipeline(
         hash_pipeline.mask_generator,
         room_view_images,
         min_area=hash_pipeline.sam_min_mask_area,
-        max_masks=hash_pipeline.sam_max_masks,
+        max_masks=sam_candidate_masks,
         points_per_batch=hash_pipeline.sam_points_per_batch,
     )
     if len(masks_dataset) != len(room_view_dataset):
@@ -204,9 +218,17 @@ def build_scene_graph_pipeline(
         description="Building object dataset...",
     ):
         save_room_view(verification_output_dir, room_id, view_idx, image)
-        total_masks += len(masks)
+        total_raw_masks += len(masks)
 
-        for mask_idx, mask in enumerate(masks):
+        ranked_masks = rank_masks(
+            masks,
+            image_shape=(image.height, image.width),
+            config=mask_scoring_config,
+            max_masks=sam_max_masks,
+        )
+        total_masks += len(ranked_masks)
+
+        for mask_idx, mask in enumerate(ranked_masks):
             masked_image = extract_masked_region(image, mask["segment"])
             object_dataset.append(
                 (room_id, view_idx, mask_idx, mask["bbox"], masked_image)
@@ -268,7 +290,8 @@ def build_scene_graph_pipeline(
         )
 
     print(f"Total objects created: {len(scene_graph.objects)}")
-    print(f"Total processed masks: {total_masks}")
+    print(f"Total raw masks from SAM: {total_raw_masks}")
+    print(f"Total kept masks after ranking: {total_masks}")
     print(f"Saved object images to: {verification_output_dir}")
     return (scene_graph,)