Adding confidence to SAM predictions from box prompts

[prernadh]

What changes are proposed in this pull request?

Added confidence scores to SAM predictions that are generate from bounding box prompts

How is this patch tested? If it is not, please explain why.

Tested manually by applying model on a dataset and confirming that confidence scores are saved on the predictions.

Release Notes

Is this a user-facing change that should be mentioned in the release notes?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release
  notes for FiftyOne users.

(Details in 1-2 sentences. You can just refer to another PR with a description
if this PR is part of a larger change.)

What areas of FiftyOne does this PR affect?

• App: FiftyOne application changes
• Build: Build and test infrastructure changes
• Core: Core fiftyone Python library changes
• Documentation: FiftyOne documentation changes
• Other: SAM utils changes

Summary by CodeRabbit

• New Features

  • Enhanced prediction outputs for the segmentation models to include confidence scores alongside existing data (boxes, labels, masks).
  • Improved consistency in output formats across different model methods.

• Bug Fixes

  • Updated method signatures to ensure accurate return values for both segmentation models.

[coderabbitai]

Walkthrough

The pull request introduces changes to the SegmentAnythingModel and SegmentAnything2ImageModel classes in the fiftyone/utils/sam.py and fiftyone/utils/sam2.py files, respectively. Both classes' _forward_pass_boxes methods have been updated to include a new output key, "scores", which captures confidence scores from their respective prediction methods. The _forward_pass_points method in SegmentAnythingModel has also been modified to include similar score outputs, ensuring consistency across the methods.

Changes

File Path	Change Summary
fiftyone/utils/sam.py	Updated _forward_pass_boxes and _forward_pass_points methods in SegmentAnythingModel to include a new "scores" key in the output dictionary.
fiftyone/utils/sam2.py	Updated _forward_pass_boxes method in SegmentAnything2ImageModel to include a new "scores" key in the output dictionary.

Possibly related PRs

• Adding MedSAM2 to the model zoo #4828: The changes in fiftyone/utils/sam2.py related to the _forward_pass_boxes method in the SegmentAnything2ImageModel class are directly relevant, as both the main PR and this retrieved PR modify the output structure to include a "scores" key in the returned dictionary.

Suggested reviewers

• brimoor

Poem

In the garden where models play,
New scores bloom in bright array.
Boxes and points, with confidence shown,
A clearer path for seeds we've sown.
Hopping forth with joy so bright,
Our predictions take their flight! 🐇✨

---

📜 Recent review details

Configuration used: .coderabbit.yaml
Review profile: CHILL

📥 Commits

Files that changed from the base of the PR and between 00c8ffb and c6a2e33.

📒 Files selected for processing (2)

• fiftyone/utils/sam.py (1 hunks)
• fiftyone/utils/sam2.py (2 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• fiftyone/utils/sam.py

🧰 Additional context used

🔇 Additional comments (2)

fiftyone/utils/sam2.py (2)

Line range hint 1-577: Overall, the changes look good and achieve the PR objectives.

The modifications to include confidence scores in the SAM2 predictions are well-implemented and localized to the _forward_pass_boxes method. This approach minimizes the risk of unintended side effects on other parts of the codebase.

A few points to consider for future improvements:

1. Adding documentation for the new scores output in the method's docstring.
2. Implementing a confidence threshold as suggested earlier.
3. Updating any relevant tests to cover the new scores output.

---

Line range hint 204-219: LGTM! Consider adding a confidence threshold.

The addition of confidence scores to the SAM predictions is well-implemented and aligns with the PR objectives. Good job on maintaining backward compatibility while enhancing the model's output.

Consider adding an optional confidence threshold parameter to filter low-confidence predictions. This could be implemented as follows:

def _forward_pass_boxes(self, imgs, confidence_threshold=0.0):
    # ... existing code ...
    
    masks, scores, _ = sam2_predictor.predict(
        point_coords=None,
        point_labels=None,
        box=sam_boxes[None, :],
        multimask_output=False,
    )
    if masks.ndim == 3:
        masks = np.expand_dims(masks, axis=0)
    
    # Apply confidence threshold
    mask = scores >= confidence_threshold
    outputs.append(
        {
            "boxes": input_boxes[mask],
            "labels": labels[mask],
            "masks": torch.tensor(masks, device=sam2_predictor.device)[mask],
            "scores": torch.tensor(scores, device=sam2_predictor.device)[mask],
        }
    )

This change would allow users to filter out low-confidence predictions if desired.

Let's verify the impact of these changes on the rest of the codebase:

This will help ensure that the changes don't break existing functionality and are consistent with other parts of the codebase.

✅ Verification successful

Verified: Changes are correctly isolated to video processing. No modifications needed for image processing.

The addition of the confidence_threshold parameter to the _forward_pass_boxes method in the video processing context does not impact other parts of the codebase. All existing functionalities remain intact, and the changes align with the intended enhancements.

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Search for uses of _forward_pass_boxes method
rg "def _forward_pass_boxes" fiftyone/utils/sam2.py -A 10
# Search for calls to _forward_pass_boxes method
rg "\._forward_pass_boxes\(" fiftyone/utils/sam2.py -A 5

Length of output: 1363

---

Script:

#!/bin/bash
# Search for calls to _forward_pass_boxes method with image parameters
rg "\._forward_pass_boxes\(" fiftyone/utils/sam2.py -A 5
# Search for calls to _forward_pass_boxes method with video parameters
rg "\._forward_pass_boxes\(" fiftyone/utils/sam2.py -A 5

Length of output: 685

---

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 0

🧹 Outside diff range and nitpick comments (2)

fiftyone/utils/sam.py (1)

274-286: LGTM! Consider adding a comment for clarity.

The changes successfully implement the addition of confidence scores to the SAM predictions, aligning with the PR objectives. The scores variable is correctly captured from the predict_torch method and added to the output dictionary.

For improved clarity, consider adding a brief comment explaining the significance of the scores variable:

 masks, scores, _ = sam_predictor.predict_torch(
     point_coords=None,
     point_labels=None,
     boxes=transformed_boxes,
     multimask_output=False,
 )
+# scores represent the confidence of each prediction
 outputs.append(
     {
         "boxes": input_boxes,
         "labels": labels,
         "masks": masks,
         "scores": scores,
     }
 )

fiftyone/utils/sam2.py (1)

Line range hint 204-219: LGTM! Consider a minor readability improvement.

The changes successfully implement the addition of confidence scores to the SAM2 predictions, aligning with the PR objective. The scores are correctly handled and added to the output dictionary.

To improve code readability, consider extracting the device assignment to a separate line:

+            device = sam2_predictor.device
             masks, scores, _ = sam2_predictor.predict(
                 point_coords=None,
                 point_labels=None,
                 box=sam_boxes[None, :],
                 multimask_output=False,
             )
             if masks.ndim == 3:
                 masks = np.expand_dims(masks, axis=0)
             outputs.append(
                 {
                     "boxes": input_boxes,
                     "labels": labels,
-                    "masks": torch.tensor(masks, device=sam2_predictor.device),
-                    "scores": torch.tensor(
-                        scores, device=sam2_predictor.device
-                    ),
+                    "masks": torch.tensor(masks, device=device),
+                    "scores": torch.tensor(scores, device=device),
                 }
             )

This change reduces repetition and makes the code slightly more maintainable.

📜 Review details

Configuration used: .coderabbit.yaml
Review profile: CHILL

📥 Commits

Files that changed from the base of the PR and between 022a7e8 and 00c8ffb.

📒 Files selected for processing (2)

• fiftyone/utils/sam.py (1 hunks)
• fiftyone/utils/sam2.py (2 hunks)

🧰 Additional context used

🔇 Additional comments (2)

fiftyone/utils/sam.py (2)

Line range hint 1-486: Overall assessment: Changes successfully implement confidence scores.

The modifications to fiftyone/utils/sam.py effectively add confidence scores to SAM predictions from box prompts, meeting the PR objectives. The implementation is consistent across different input types (_forward_pass_boxes and _forward_pass_points), and the existing _forward_pass_auto method already includes a similar score. These changes enhance the output of the model's prediction methods by providing valuable confidence information for each prediction.

---

Line range hint 1-486: Verify consistency of score implementation across methods.

The addition of scores to the output dictionary in _forward_pass_boxes is consistent with the existing implementation in _forward_pass_points. The _forward_pass_auto method already includes a score (predicted_iou) in its output, maintaining uniformity across different input types.

To ensure complete consistency, run the following script to check if all relevant methods include a score or confidence value in their outputs:

This script will help confirm that all forward pass methods consistently include a score or confidence value in their outputs.

[mythrandire]

Not a review comment, just a tangential observation - @prernadh any reason why we shouldn't also include the final return value (low_res_masks_np) from the predict method? I don't see a reason to not expose it, seeing as the change mirrors the addition of scores.

[mythrandire]

No comments, fairly straightforward addition to expose prediction IOU scores.

[prernadh]

Not a review comment, just a tangential observation - @prernadh any reason why we shouldn't also include the final return value (low_res_masks_np) from the predict method? I don't see a reason to not expose it, seeing as the change mirrors the addition of scores.

I would like to avoid saving too many additional fields in the sample unless necessary. My understanding of the low_res_masks_np is that its useful if you are using this predict method as an intermediate step and need to pass these low resolution masks further downstream to a model etc. I don't think that is required currently but if a user asks for it, it should be a small change