Feature: AI-assisted 3DGS Scene Segmentation

[johnbanq]

AI-assisted 3DGS scene segmentation has been a highly requested feature of this tool (since issue #6).

So let's make it happen! Here is the central issue for another attempt at this feature.

Goals

This attempt aims to achieve:

• make the user able to rapidly coarsely segment (prefers to oversegment than undersegment) an object from a loaded 3DGS scene using scribbles/point clicks at a few sparse view points
• The approach should be lightweight and still reasonably fast

Non-Goals

This attempt will not aim to achieve:

• precise, high-quality segmentation
• Further UI improvements for labelling and annotation

Related Issues & PRs

• issue Segment-Anything to remove noise #6 - an earlier request for this feature
• PR Adds Object Segmentation Tool #30 - an abandoned branch & PR for an preliemary attempt at this feature
• issue Gaussian Splat labelling workflow #216 - adapting supersplat for annotation and labelling

[johnbanq]

Existing Work

1. Gauzilla Pro

Link: https://www.gauzilla.xyz/

Single Object Segmentation

Approach (Guessed):

• takes an image at a viewpoint, use user clicks to segment the image using SAM, the mask can then be used to cull the viewing frustum to areas that might include the object.
• do this over two or more viewpoints, and then use the intersection of the viewing frustums to get a rough 3D segmentation.

Multi Object Segmentation

Approach (Guessed):

• takes an image at a viewpoint and a text prompt, use VLM to detect all objects fitting the prompt in the image, the boxes can then be used to cull the viewing frustum.
• do this over two or more viewpoints, and then use the intersection of the viewing frustums to get a rough 3D segmentation for every object.

Pro:

• Simple - this is probably the simplest way to do 3D segmentation from 2D images, simple makes it robust
• Efficient - only requires running 2D vision models over a few viewpoints, the models are popular so lots of lightweight explorations.

Con:

• Requires more work - annotation of multiple viewpoints is required
• Coarse - the segmentation is not very precise, see the butterfly scene, large amount of walls are included

2. Previous PR & Issue

Link: 8165ad2 of PR #30.

Approach:

• Adds a new tool that takes image at a viewpoint and user scribbles, use interactive segmenter from mediapipe to produce 2D mask, use 2D mask to to filter selected splats whose projected 2D center is in the mask.

Pro:

• Simple - simple and straight-forward

Con:

• Coarse - single view only, so a lot of extra splats are included

Link Issue #6

• (potentially relevant, using SAM etc. to remove noises / background from scene)
• 1. proposes to use DBSCAN clustering to remove floaters and noise from scene, promising, but slow. You can pile tricks to get it to work, but floaters are out of scope.
• 2. use SAMURAI to video segment the 3DGS training images before training. Requires re-training so out of scope.

Link: Issue #216

• (potentially relevant, utilizing supersplat for annotation)
• note: has DBSCAN clustering impl.

3. Academic Literature

3DGS segmentation is a pretty hot topic in academic research, I got these notable papers after a quick search in "lightweight rapid 3DGS segmentation":

• FlashSplat 2D to 3D Gaussian Splatting Segmentation Solved Optimally
• GaussianCut Interactive segmentation via graph cut for 3D Gaussian Splatting
• Gradient-Driven 3D Segmentation and Affordance Transfer in Gaussian Splatting Using 2D Masks
• Lifting by Gaussians A Simple Fast and Flexible Method for 3D Instance Segmentation

They can achieve high-quality segmentation with good speed, but they are quite complicated to implement.

I believe FlashSplat + SAGD could be an interesting starting point for a "refinement" stage after coarse segmentation. But that is not our immediate focus.

Proposed Plan

• we start with (guessed to be) Gauzilla Pro's multiview frustum-intersection approach, start with SAM masks, if undersegments, try object boxes. UI coded in Adds Object Segmentation Tool #30 could be reused.
• we then investigate FlashSplat + SAGD in the future.

[marklundin]

We had a very early PR for this #30

Realised you've already linked to this PR

[papinto]

Adding to the Existing Work, I was just researching this today and came across this paper: SAGA

[johnbanq]

A quick update on the prototype, I got the "simplest thing that can possibly work" working here: https://github.com/johnbanq/supersplat/tree/feat/segmentation-proto-v2

Demo: https://supersplat-git-feat-segmentation-proto-v2-johnbanqs-projects.vercel.app
Demo Data: https://drive.google.com/file/d/1JDFdk3Z5EIbx-K6pqPEiiLVJXGmrWSMB/view?usp=sharing
Video: https://www.youtube.com/watch?v=buDbqHVt4Og

With this prototype, I discovered some questions that might need discussion (@slimbuck a little help plz 😉?):

• Firstly, I am uncertain about the UI/UX of this feature:

  • Currently our version is based on Gauzilla Pro's design, the user use prompt points to 2D segment a single view, the 2D segmentation is then used to select all splats that fall in the mask, the user then move to another view, do another segmentation and so on, the result is then the AND of all the 3D masks among all viewpoints.
  • A more flexible version (iirc PR Adds Object Segmentation Tool #30) might be that we keep the 2D segment & 3D segment part, but instead of recording 3D mask at all viewpoints and AND it, we just let users choose to SET/AND/OR the current selection. This would make the feature simpler, more flexible and easier to implement things like undo/redo, but might require more documentation / guidance to the user to achieve the multiview segmentation presented in Gauzilla Pro.

• Secondly, I don't have any good idea about how to properly implement the segmentation overlay.

  • I am currently using a canvas and set the zIndex of everyone else to sandwich it between the main canvas and the UI floating on top (see johnbanq@2e0b01e#diff-3521eed3054b6487ff81166a3903ad15a8d4e13115e2322428effcf0eea196e2R73), this is ridiculously sketchy lol, any better ideas? I tried looking into how the existing tools (polygon/brush selection tool) works, but they seems to use svg to show the selection, while this feature produces pixel image masks.

Apart from the questions above, this prototype absolutely needs more polish, I will get back to it in a while!

[johnbanq]

Here is the alternative UX I mentioned above: https://github.com/johnbanq/supersplat/tree/feat/segmentation-v3-2.

Demo: https://precious-toffee-4fcb0c.netlify.app/ (segmentation model is a bit wonky, sorry)
Demo Data: https://drive.google.com/file/d/1JDFdk3Z5EIbx-K6pqPEiiLVJXGmrWSMB/view?usp=sharing
Video: https://youtu.be/0SV3eTQ_WR4

The difference is that this version is more flexible, the user just 2D segment from a single view and choose to set/and/or the splats falls in the 2D mask with the current selection.

The only concern is that the user might need some instruction to achieve that 3D segmentation experience in Gauzilla Pro, since we are just providing tool achieving that, not hard coding the workflow.

Should we go with this one or the previous one?