Skip to content

feat(tflite): decode segmentation masks in TFLite inference - #1053

Merged
Borda merged 11 commits into
roboflow:developfrom
omkar-334:tflite-segmentation-support
May 19, 2026
Merged

Borda merged 11 commits into
roboflow:developfrom
omkar-334:tflite-segmentation-support

Conversation

@omkar-334

Copy link
Copy Markdown
Contributor

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 masks output alongside dets and labels. After #1041, the TFLite graph for a segmentation model converts and includes that masks tensor, but _run_inference decoded 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):

Variant Conversion FP32 FP16 dynamic INT8
RFDETRSegNano OK 0.983 / 0.979 0.981 / 0.978 0.967 / 0.974
RFDETRSegSmall OK 0.961 / 0.947 0.984 / 0.966 0.971 / 0.966
RFDETRSegMedium OK 0.972 / 0.969 0.983 / 0.963 0.940 / 0.944
RFDETRSegLarge OK 0.983 / 0.976 0.961 / 0.952 0.956 / 0.916
RFDETRSegPreview OK 0.973 / 0.972 0.972 / 0.973 0.968 / 0.970

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

  • Bug fix (non-breaking change that fixes an issue)

Testing

  • I have tested this change locally
  • I have added/updated tests for this change

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have commented my code where necessary, particularly in hard-to-understand areas
  • My changes generate no new warnings or errors
  • I have updated the documentation accordingly (if applicable)

@codecov

codecov Bot commented May 19, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 96.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 80%. Comparing base (6769f30) to head (0f80bfa).

❌ 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:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@Borda Borda added the enhancement New feature or request label May 19, 2026
@Borda
Borda requested a review from Copilot May 19, 2026 09:34

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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-4 masks output and populate Detections.mask for 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.

Comment thread src/rfdetr/export/_tflite/inference.py
omkar-334 and others added 2 commits May 19, 2026 15:20
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
Borda previously approved these changes May 19, 2026
Borda and others added 4 commits May 19, 2026 12:41
… 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>

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated no new comments.

@Borda
Borda merged commit 8ae6767 into roboflow:develop May 19, 2026
30 checks passed
@Borda Borda mentioned this pull request May 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Question: tflite support

3 participants