feat(tflite): decode segmentation masks in TFLite inference - #1053
Merged
Merged
Conversation
omkar-334
requested review from
Borda,
SkalskiP,
isaacrob and
probicheaux
as code owners
May 19, 2026 02:30
Codecov Report❌ Patch coverage is ❌ Your project check has failed because the head coverage (80%) is below the target coverage (95%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #1053 +/- ##
=======================================
Coverage 80% 80%
=======================================
Files 101 101
Lines 8695 8719 +24
=======================================
+ Hits 6987 7010 +23
- Misses 1708 1709 +1 🚀 New features to boost your workflow:
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds end-to-end instance segmentation support to the TFLite inference helper by decoding the exported masks output into supervision.Detections.mask, aligning TFLite inference behavior with the PyTorch post-processing path.
Changes:
- Add
_decode_masks()to upsample mask logits to image size and threshold them into boolean masks. - Extend
_run_inference()to detect an optional rank-4masksoutput and populateDetections.maskfor segmentation exports. - Add targeted unit tests covering mask decoding and the segmentation vs detection output behavior; update TFLite converter docs to reflect validated segmentation support.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/rfdetr/export/_tflite/inference.py |
Implements mask decoding and wires it into _run_inference() when a masks output is present. |
tests/export/test_tflite_inference.py |
Adds unit tests for _decode_masks() and verifies masks are set only for 3-output segmentation exports. |
src/rfdetr/export/_tflite/converter.py |
Updates module/API docstrings to reflect validated segmentation outputs in TFLite exports. |
PILImage.Resampling was introduced in Pillow 9.1; older environments raise AttributeError. Replace with a module-level fallback that prefers Resampling.BILINEAR when available and falls back to the legacy PILImage.BILINEAR constant. [resolve roboflow#1] @Copilot / @Borda (gh): roboflow#1053 (comment) --- Co-authored-by: Claude Code <noreply@anthropic.com>
Borda
previously approved these changes
May 19, 2026
… type - Add `if mask_logits.ndim != 3: raise ValueError(...)` guard at top of `_decode_masks` so mismatched tensors from the rank-4 heuristic fail loudly rather than producing silent shape corruption downstream - Tighten return type annotation from `NDArray[Any]` to `NDArray[np.bool_]` and `dtype=bool` → `dtype=np.bool_` for consistency - Soften docstring claim from "Mirrors PostProcess.forward" to "Approximates" and document that Pillow bilinear and PyTorch F.interpolate may differ at border pixels due to distinct half-pixel conventions Resolve /review finding roboflow#2 (foundry:sw-engineer HIGH) + roboflow#7 (foundry:doc-scribe MEDIUM) from .reports/review/2026-05-19T10-06-08Z/review-report.md --- Co-authored-by: Claude Code <noreply@anthropic.com>
…didates Previously the shape-based fallback called `next(...)` over all rank-4 outputs, silently picking the first one. If a future model variant or onnx2tf version emits an auxiliary rank-4 tensor, the wrong tensor would be decoded as masks without any diagnostic. Now the fallback collects all rank-4 candidates: - exactly 1 → accepted as the mask output (same as before) - ≥ 2 → log a warning and skip mask decode (fail safe, not silent corruption) The name-based path (`"masks" in output name`) is unchanged and takes priority. Resolve /review finding roboflow#1 (foundry:sw-engineer HIGH) from .reports/review/2026-05-19T10-06-08Z/review-report.md --- Co-authored-by: Claude Code <noreply@anthropic.com>
When keep.any() is False (all scores below threshold), mask_idx may still be set for a segmentation model, causing _decode_masks to receive a (0, Hm, Wm) input and returning a (0, H, W) bool array stored as Detections.mask. Callers checking `dets.mask is None` would behave differently from `len(dets) == 0`. Now mask decode is gated on `keep.any()`, so segmentation models return `Detections.mask = None` when zero detections survive — consistent with the detection-only model behaviour. Resolve /review finding roboflow#6 (foundry:sw-engineer MEDIUM) from .reports/review/2026-05-19T10-06-08Z/review-report.md --- Co-authored-by: Claude Code <noreply@anthropic.com>
… edge cases Resolve /review findings roboflow#4, roboflow#5, roboflow#8, roboflow#9, roboflow#10 from .reports/review/2026-05-19T10-06-08Z/review-report.md - test_run_inference_name_based_mask_detection: exercises primary name-based path (`"masks" in output name`) — previously dead code under test; every existing seg test used `Identity_N` which bypassed it - test_run_inference_seg_model_no_detections_returns_none_mask: seg model + all logits below threshold → `Detections.mask is None` (post-keep.any() gate) - test_decode_masks_raises_on_wrong_rank: _decode_masks rejects rank-4 input with a clear ValueError (exercises new guard from prior commit) - test_decode_masks_exact_zero_logit_decodes_to_false: boundary at strict >0; uniform-zero logit map → all-False (not all-True) - test_decode_masks_non_square_logit_input: (3,7,14) logits → (3,28,56) output, guards PIL bilinear resize for non-square feature maps - test_decode_masks_parity_positive_negative_regions: high-magnitude ±10 logits verify correct True/False regions after bilinear upsample; parity regression guard for PostProcess.forward equivalent contract --- Co-authored-by: Claude Code <noreply@anthropic.com>
Borda
approved these changes
May 19, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds instance mask decoding to the TFLite inference helper so segmentation models work end to end after TFLite export. Builds on the TFLite export support added in #1041.
Fixes #1035
RF-DETR segmentation models export a third
masksoutput alongsidedetsandlabels. After #1041, the TFLite graph for a segmentation model converts and includes thatmaskstensor, but_run_inferencedecoded only boxes and labels and ignored any extra output. A user exporting a segmentation model to TFLite could get boxes back but not masks.Verification
Aggregate parity vs the PyTorch baseline (box IoU mean / mask IoU mean):
FP32 and FP16 track the baseline closely for both boxes and masks.
Dynamic-range INT8 is close, with marginally lower mask fidelity.
Type of Change
Testing
Checklist