Add YOLO26 semantic segmentation export support#217
Conversation
📝 WalkthroughWalkthroughAdds end-to-end support for exporting YOLO26 semantic segmentation models. A new ChangesYOLOv26 Semantic Segmentation Export Pipeline
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
☂️ Python Coverage
Overall Coverage
New FilesNo new covered files... Modified Files
|
Test Results 6 files 6 suites 23m 35s ⏱️ Results for commit 501846c. ♻️ This comment has been updated with latest results. |
klemen1999
left a comment
klemen1999
left a comment
There was a problem hiding this comment.
Overall LGTM, just one nitpick and general question:
Do we potentially need any ultralytic submodule version bump? Just checking...
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tools/modules/exporter.py (1)
247-253: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd a class-count consistency check at the semantic archive boundary.
class_listandn_classesare accepted independently, but there is no guard here. If they diverge, metadata becomes inconsistent and can break label mapping downstream. Please validate this inmake_semantic_seg_nn_archive(Line 247 onward), not only in callers.Suggested patch
def make_semantic_seg_nn_archive( self, class_list: List[str], n_classes: int, is_softmax: bool = False, encoding: Encoding = Encoding.RGB, ): @@ + if len(class_list) != n_classes: + raise ValueError( + f"Number of class names ({len(class_list)}) does not match n_classes ({n_classes})." + ) + self.f_nn_archive = (self.output_folder / f"{self.model_name}.tar.xz").resolve() output_specs = self.get_output_specs()Also applies to: 300-303
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tools/modules/exporter.py` around lines 247 - 253, The `make_semantic_seg_nn_archive` method accepts `class_list` and `n_classes` independently without validating consistency between them. Add a guard check at the beginning of the `make_semantic_seg_nn_archive` method to ensure that the length of `class_list` matches the value of `n_classes`, and raise a clear error if they diverge. This prevents metadata inconsistencies from propagating downstream into label mapping operations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@tools/modules/exporter.py`:
- Around line 247-253: The `make_semantic_seg_nn_archive` method accepts
`class_list` and `n_classes` independently without validating consistency
between them. Add a guard check at the beginning of the
`make_semantic_seg_nn_archive` method to ensure that the length of `class_list`
matches the value of `n_classes`, and raise a clear error if they diverge. This
prevents metadata inconsistencies from propagating downstream into label mapping
operations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0810cdaa-b2df-40ec-8a59-effebfdcf6f6
📒 Files selected for processing (3)
tests/test_end2end.pytools/modules/exporter.pytools/yolo/yolo26_exporter.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tools/yolo/yolo26_exporter.py
- tests/test_end2end.py
I don't think we need a bumping right now. The current submodule pointer already includes SemanticSegment, and the targeted yolo26n-sem NNArchive test passes with the existing submodule version |
2 participants
Purpose
Add support for exporting YOLO26 semantic segmentation models from
tools.This enables
yolo26n-sem.ptto be exported as an ONNX NNArchive with a semantic segmentation head configuration, so the resulting archive can be converted through HubAI and used directly withParsingNeuralNetwork/SegmentationParser.Specification
This PR adds:
yolov26-semversion detection.yolov26-semin the main conversion flow.Yolo26Exporter.make_segmentation_nn_archive()helper inExporterto generate NNArchives withSegmentationParser+HeadSegmentationMetadata.The YOLO26 semantic segmentation checkpoint is detected through the
SemanticSegmenthead in the model checkpoint.The generated archive keeps the existing YOLO26 output naming convention and uses
output_yolo26as the semantic segmentation output.The generated ONNX output is:
The generated NNArchive head config is:
No changes were needed in
depthai-nodes, since the existingSegmentationParseralready supports a single semantic segmentation logits output and produces aSegmentationMask.I also verified that the existing YOLO26 detect / instance segmentation / pose output-name behavior stays unchanged.
Dependencies & Potential Impact
No new dependencies.
Impact should be limited to YOLO26 semantic segmentation export. Existing YOLO26 export paths are not intended to change.
Deployment Plan
No special deployment required.
Rollback is reverting this PR.
Testing & Validation
Tests:
Manual export validation:
yolo26n-sem.ptthroughtools.output_yolo26 [1, 19, 416, 416].SegmentationParser.["output_yolo26"].is_softmax=false.Hub conversion validation:
The converted archive configs were inspected and confirmed to preserve:
Runtime validation:
The converted archives were tested with a DepthAI pipeline using:
The parser output was:
and the visualizer rendered semantic segmentation masks / overlays correctly.
I also validated the RVC2 archive locally with
oak-examples/neural-networks/generic-exampleon OAK-1 using-overlay.AI Usage
Assisted-by: ChatGPT
Submitted code was reviewed by a human: YES
The author is taking responsibility for the contribution: YES
Summary by CodeRabbit
Release Notes
New Features
Tests