Add SAM2 models and AI Batch Mode #1493
Conversation
|
sam-cpp-macos is the Segment Anything Model 2 CPP Wrapper for macOS and Ubuntu CPU/GPU. On macOS CPU use-case.
sam2_polygon.mp4 |
WalkthroughThis update expands segmentation capabilities by introducing new model classes and a complete segmentation model implementation using the SAM2 architecture. Two new classes, one optimized for accuracy and one for speed, extend the existing model. A new file implements segmentation functionality with image encoding, decoding, and embedding caching. Additionally, a new "ai_batch" mode is integrated into the GUI, with corresponding UI actions, configurations, and canvas event handling. Minor formatting cleanups in utility functions are also included. Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant MW as MainWindow
participant C as Canvas
participant AM as AI Model
U->>MW: Select "AI Batch Mode"
MW->>C: Activate AI Batch Mode
U->>C: Draw shape (mouse events)
C->>AM: Finalise drawing and send points
AM-->>C: Return predicted polygon/mask
C->>MW: Update display with new shape
MW-->>U: Render updated canvas
sequenceDiagram
participant C as Canvas
participant SAM as SegmentAnything2Model
participant EN as SAM2ImageEncoder
participant DE as SAM2ImageDecoder
C->>SAM: set_image(image)
SAM->>EN: Compute image embedding
EN-->>SAM: Return embedding
SAM->>DE: Predict mask/polygon from points
DE-->>SAM: Return segmentation result
SAM-->>C: Provide final segmentation output
Poem
Tip CodeRabbit's docstrings feature is now available as part of our Pro Plan! Simply use the command ✨ Finishing Touches
🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (7)
labelme/ai/__init__.py (2)
90-102: Definition ofSAM2HieraL.
Good introduction of a new class extendingSegmentAnything2Modelwith specialized encoder/decoder paths for the large model. Consider verifying file checksums or versioning if you anticipate frequent updates to these ONNX files.
103-114: Definition ofSAM2HieraT.
Similar toSAM2HieraL, but optimized for speed using tiny model variants. Overall structure is consistent. Also consider logging or verifying model downloads in case of updates or potential network issues.labelme/widgets/canvas.py (1)
785-814: Painting logic forai_batch.
Uses skimage’sfind_contoursandapproximate_polygonto outline multiple polygons, painting each contour. This is an efficient approach for batch previews, though be mindful of performance with large images.labelme/ai/segment_anything2_model.py (4)
17-22: Consider thread-safety in multi-threaded usage.
This class holds a_threadreference and uses a lock, but as soon as a newset_imagecall happens, the_threadreference might be overwritten. If multiple threads callset_imagein parallel, the logic may become interleaved. Clarify whether this class is intended for single-threaded or truly concurrent use, and add additional synchronization if necessary.
36-44: Use LRU logic or a hashed key for memory efficiency.
Storing entire raw images as the cache key (self._image.tobytes()) can significantly impact memory usage for large images. Also, usingpopitem(last=False)inOrderedDictremoves items in insertion order, not necessarily least recently used. Consider using an LRU approach or a hashed key to optimize memory usage and cache management.
52-62: Simplify unnecessary array reshaping.
best_maskis wrapped innp.array([[best_mask]])only to be indexed back at[0,0]. To improve clarity and reduce extra operations, remove these redundant steps and work directly with the mask array.- best_mask = np.array([[best_mask]]) - best_mask = best_mask[0,0]
151-156: Consider adding error handling for inference.
Currently, ifself.session.run()fails for any reason, no exception handling occurs. This may leave the session in an inconsistent state or raise unhandled errors. Add a try-except block or logging to handle inference failures gracefully.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
labelme/ai/__init__.py(2 hunks)labelme/ai/_utils.py(1 hunks)labelme/ai/segment_anything2_model.py(1 hunks)labelme/app.py(8 hunks)labelme/config/default_config.yaml(1 hunks)labelme/widgets/canvas.py(12 hunks)
✅ Files skipped from review due to trivial changes (1)
- labelme/ai/_utils.py
🔇 Additional comments (15)
labelme/config/default_config.yaml (1)
84-84: Newai_batchconfiguration option introduced.
This new boolean toggle integrates seamlessly with the surrounding AI-related settings, and no immediate issues are observed.labelme/ai/__init__.py (2)
4-4: Added import forSegmentAnything2Model.
This ensures the base class is accessible for the newly introduced SAM2 classes. No concerns here.
122-123: Added new classes to theMODELSlist.
Makes the new SAM2 variants discoverable by the rest of the system. No concerns.labelme/app.py (5)
398-413: New "Create AI-Batch" mode action.
This action properly parallels the existing AI creation modes. The callback to initialize the AI model is consistent with other AI-based modes.
667-667: Inclusion ofcreateAiBatchModein the actions structure.
This ensures the new mode's action is exposed to the application’s toolset.
706-706: Extended context menu withcreateAiBatchMode.
Properly follows the pattern used by other creation actions.
727-727: EnabledcreateAiBatchModeon shape load.
Consistent with other AI-based modes, ensuring it is available when images are loaded.
998-998: Enabling thecreateAiBatchModeon file load.
This completes the wire-up by allowing the mode to be active once the user opens a file.labelme/widgets/canvas.py (5)
66-66:"ai_batch": falsecrosshair setting.
Adds crosshair handling for the new mode. No issues; consistent with other mode flags.
261-261: Added "ai_batch" to the creation mode check.
Allows the code path for AI-based shape logic to properly trigger.
289-289: Extends conditional for AI-based drawings to includeai_batch.
Ensures the line logic includes the new mode. No immediate issues.
447-447: Extended multi-mode check forai_batchin mousePressEvent.
Preserves the existing code style and ensures the new mode’s logic remains consistent with other AI modes.
863-905: Finalizing shapes forai_batch.
- Iterates through each contour to create multiple polygons in batch.
- Re-initializes
self.currentto a new shape after appending each one, ensuring each contour is stored.Implementation looks correct for multi-polygon batch mode. Keep an eye on large or noisy masks, as excessive contours could degrade performance.
labelme/ai/segment_anything2_model.py (2)
23-35: Verify old threads won't be prematurely overshadowed.
Whenset_imageis called repeatedly in quick succession, the old worker thread may still be running unless a prior_get_image_embeddinghas been invoked. This could lead to race conditions or overshadowing. Confirm that the design prevents simultaneous computation for multiple images to avoid concurrency issues.
64-65: Confirm that remove_small_objects accepts float values formin_size.
skimage.morphology.remove_small_objectstypically expects an integermin_size. Though it may accept float inputs, confirm correctness to avoid future compatibility problems.
| import numpy as np | ||
| import skimage | ||
|
|
There was a problem hiding this comment.
💡 Verification agent
❓ Verification inconclusive
New imports for numpy, skimage, and related modules.
These are essential for mask handling, contour detection, and polygon approximations. Ensure the project dependencies reflect these libraries.
Would you like to run a quick script to confirm that these dependencies are listed in your environment or requirements?
🏁 Script executed:
#!/bin/bash
# Check if numpy and scikit-image are installed in the environment
pip show numpy
pip show scikit-imageLength of output: 121
Dependency Update Required: Add numpy and scikit-image to the Project Dependencies
We verified via the pip commands that the required packages (numpy and scikit-image) are not currently installed in the environment. Since these libraries are essential for mask handling, contour detection, and polygon approximations—and they are imported in labelme/widgets/canvas.py (lines 13-15)—please make sure that the project’s dependency configuration (for example, in your requirements.txt or setup.py) includes these packages.
- Location:
labelme/widgets/canvas.py(lines 13-15) - Action Required:
- Add
numpyandscikit-image(the latter is imported asskimage) to your dependency file. - Verify that any environments used (development, CI, production) install these dependencies correctly.
- Add
| def encode(self, cv_image: np.ndarray) -> list[np.ndarray]: | ||
| original_size = cv_image.shape[:2] | ||
| high_res_feats_0, high_res_feats_1, image_embed = self.encoder(cv_image) | ||
| return { |
There was a problem hiding this comment.
Fix return type annotation.
The method signature says encode(self, cv_image: np.ndarray) -> list[np.ndarray], but it returns a dictionary. Update either the signature or the return structure to maintain consistency.
- def encode(self, cv_image: np.ndarray) -> list[np.ndarray]:
+ def encode(self, cv_image: np.ndarray) -> dict[str, np.ndarray | tuple[int, int]]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def encode(self, cv_image: np.ndarray) -> list[np.ndarray]: | |
| original_size = cv_image.shape[:2] | |
| high_res_feats_0, high_res_feats_1, image_embed = self.encoder(cv_image) | |
| return { | |
| def encode(self, cv_image: np.ndarray) -> dict[str, np.ndarray | tuple[int, int]]: | |
| original_size = cv_image.shape[:2] | |
| high_res_feats_0, high_res_feats_1, image_embed = self.encoder(cv_image) | |
| return { |
| def process_output( | ||
| self, outputs: list[np.ndarray] | ||
| ) -> tuple[list[ndarray | Any], ndarray[Any, Any]]: | ||
|
|
||
| scores = outputs[1].squeeze() | ||
| masks = outputs[0][0] | ||
|
|
||
| return (masks, | ||
| scores, | ||
| self.orig_im_size | ||
| ) | ||
|
|
There was a problem hiding this comment.
Correct mismatch in tuple size vs. type annotation.
process_output indicates -> tuple[list[ndarray | Any], ndarray[Any, Any]], yet returns three elements (masks, scores, self.orig_im_size). Update the annotation or split the functionality to be consistent with the declared return type.
- ) -> tuple[list[ndarray | Any], ndarray[Any, Any]]:
+ ) -> tuple[list[np.ndarray], np.ndarray, tuple[int, int]]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def process_output( | |
| self, outputs: list[np.ndarray] | |
| ) -> tuple[list[ndarray | Any], ndarray[Any, Any]]: | |
| scores = outputs[1].squeeze() | |
| masks = outputs[0][0] | |
| return (masks, | |
| scores, | |
| self.orig_im_size | |
| ) | |
| def process_output( | |
| self, outputs: list[np.ndarray] | |
| ) -> tuple[list[np.ndarray], np.ndarray, tuple[int, int]]: | |
| scores = outputs[1].squeeze() | |
| masks = outputs[0][0] | |
| return (masks, | |
| scores, | |
| self.orig_im_size | |
| ) |
| def predict_masks(self, embedding, points, labels) -> list[np.ndarray]: | ||
| points, labels = np.array(points), np.array(labels) | ||
|
|
||
| image_embedding = embedding["image_embedding"] | ||
| high_res_feats_0 = embedding["high_res_feats_0"] | ||
| high_res_feats_1 = embedding["high_res_feats_1"] | ||
| original_size = embedding["original_size"] | ||
| self.decoder.set_image_size(original_size) | ||
| masks, scores, orig_im_size = self.decoder( | ||
| image_embedding, | ||
| high_res_feats_0, | ||
| high_res_feats_1, | ||
| points, | ||
| labels, | ||
| ) | ||
|
|
||
| return masks, scores, orig_im_size | ||
|
|
There was a problem hiding this comment.
Align declared return type with actual triple.
predict_masks declares -> list[np.ndarray], but it actually returns (masks, scores, orig_im_size). Adjust the type hint to reflect the three-element tuple being returned.
- def predict_masks(self, embedding, points, labels) -> list[np.ndarray]:
+ def predict_masks(self, embedding, points, labels) -> tuple[list[np.ndarray], np.ndarray, tuple[int, int]]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def predict_masks(self, embedding, points, labels) -> list[np.ndarray]: | |
| points, labels = np.array(points), np.array(labels) | |
| image_embedding = embedding["image_embedding"] | |
| high_res_feats_0 = embedding["high_res_feats_0"] | |
| high_res_feats_1 = embedding["high_res_feats_1"] | |
| original_size = embedding["original_size"] | |
| self.decoder.set_image_size(original_size) | |
| masks, scores, orig_im_size = self.decoder( | |
| image_embedding, | |
| high_res_feats_0, | |
| high_res_feats_1, | |
| points, | |
| labels, | |
| ) | |
| return masks, scores, orig_im_size | |
| def predict_masks(self, embedding, points, labels) -> tuple[list[np.ndarray], np.ndarray, tuple[int, int]]: | |
| points, labels = np.array(points), np.array(labels) | |
| image_embedding = embedding["image_embedding"] | |
| high_res_feats_0 = embedding["high_res_feats_0"] | |
| high_res_feats_1 = embedding["high_res_feats_1"] | |
| original_size = embedding["original_size"] | |
| self.decoder.set_image_size(original_size) | |
| masks, scores, orig_im_size = self.decoder( | |
| image_embedding, | |
| high_res_feats_0, | |
| high_res_feats_1, | |
| points, | |
| labels, | |
| ) | |
| return masks, scores, orig_im_size |
SAM2 Large Models
https://github.com/jakep72/labelme/releases/download/SAM2/sam2_large.encoder.onnx
https://github.com/jakep72/labelme/releases/download/SAM2/sam2_large.decoder.onnx
SAM2 Tiny Models
https://github.com/jakep72/labelme/releases/download/SAM2/sam2_hiera_tiny.encoder.onnx
https://github.com/jakep72/labelme/releases/download/SAM2/sam2_hiera_tiny.decoder.onnx
Summary by CodeRabbit
New Features
Chores